// based on https://github.com/honojs/middleware/blob/main/packages/zod-validator/src/index.ts
import type { ZodSafeParseResult as ZodSafeParseResult$1 } from 'zod/v4'
import type { Context, Env, Input, MiddlewareHandler, TypedResponse, ValidationTargets } from 'hono'
import { zValidator as zValidatorBase } from '@hono/zod-validator'
import * as v3 from 'zod/v3'
import * as v4 from 'zod/v4/core'

type Awaitable<T> = T | Promise<T>
type HasUndefined<T> = undefined extends T ? true : false

type ZodSchema = v3.ZodType | v4.$ZodType
type ZodError<T extends ZodSchema> = T extends v4.$ZodType ? v4.$ZodError<v4.output<T>> : v3.ZodError
type ZodSafeParseResult<T, T2, T3 extends ZodSchema> = T3 extends v4.$ZodType ? ZodSafeParseResult$1<T> : v3.SafeParseReturnType<T, T2>
type zInput<T> = T extends v3.ZodType ? v3.input<T> : T extends v4.$ZodType ? v4.input<T> : never
type zOutput<T> = T extends v3.ZodType ? v3.output<T> : T extends v4.$ZodType ? v4.output<T> : never
type zInfer<T> = T extends v3.ZodType ? v3.infer<T> : T extends v4.$ZodType ? v4.infer<T> : never
type Hook<T, E extends Env, P extends string, Target extends keyof ValidationTargetsWithResponse = keyof ValidationTargetsWithResponse, O = {}, Schema extends ZodSchema = any> = (result: ({
  success: true
  data: T
} | {
  success: false
  error: ZodError<Schema>
  data: T
}) & {
  target: Target
}, c: Context<E, P>) => Awaitable<Response | void | TypedResponse<O>>

type ValidationTargetsWithResponse = ValidationTargets & { response: any }

/**
 * zValidator wraps `zValidator` from `@hono/zod-validator` and extends it to support response validation. It forwards
 * query parameter, path parameter, and request body validation to `@hono/zod-validator`.
 */
export const zValidator =
  <
    T extends ZodSchema,
    Target extends keyof ValidationTargetsWithResponse,
    E extends Env,
    P extends string,
    In = zInput<T>,
    Out = zOutput<T>,
    I extends Input = {
      in: HasUndefined<In> extends true
        ? {
          [K in Target]?: In extends ValidationTargetsWithResponse[K]
            ? In
            : {
              [K2 in keyof In]?: In[K2] extends ValidationTargetsWithResponse[K][K2]
                ? In[K2]
                : ValidationTargetsWithResponse[K][K2]
            }
        }
        : {
          [K in Target]: In extends ValidationTargetsWithResponse[K]
            ? In
            : {
              [K2 in keyof In]: In[K2] extends ValidationTargetsWithResponse[K][K2]
                ? In[K2]
                : ValidationTargetsWithResponse[K][K2]
            }
        }
      out: { [K in Target]: Out }
    },
    V extends I = I,
    InferredValue = zInfer<T>
  > (
    target: Target,
    schema: T,
    hook?: Hook<InferredValue, E, P, Target, {}, T>
  ): MiddlewareHandler<E, P, V> =>
    async (c, next) => {
      if (target !== 'response'){
        return zValidatorBase<
          T,
          keyof ValidationTargets,
          E,
          P,
          In,
          Out,
          I,
          V,
          InferredValue
        >(
          target,
          schema,
          hook as Hook<InferredValue, E, P, keyof ValidationTargets, {}, T>
        )(c, next)
      }

      await next()

      if (
        c.res.status !== 200 ||
        !c.res.headers.get('Content-Type')?.includes('application/json')
      ) {
        return
      }

      let value: unknown
      try {
        value = await c.res.json()
      } catch {
        const message = 'Malformed JSON in response'
        c.res = new Response(message, { status: 400 })

        return
      }

      const { success, data, error } = await (schema as v3.ZodType).safeParseAsync(value) as ZodSafeParseResult<InferredValue, Out, T>

      if (hook) {
        const hookResult = await hook({
          target, success,
          data: data as InferredValue,
          error: error as ZodError<T>
        }, c)

        if (typeof hookResult === 'object' && hookResult != null) {
          if (hookResult instanceof Response) {
            c.res = new Response(hookResult.body, hookResult)
          } else if ('response' in hookResult && hookResult.response instanceof Response) {
            c.res = new Response(hookResult.response.body, hookResult.response)
          }
        }
      }

      if (!success) {
        c.res = new Response(JSON.stringify({ success, data, error }), {
          status: 400,
          headers: {
            'Content-Type': 'application/json'
          }
        })
      } else {
        c.res = new Response(JSON.stringify(data), c.res)
      }
    }