open_in_new content_copy 
/* eslint-disable @typescript-eslint/no-explicit-any */
import { Request, Response } from 'express'
import createHttpError from 'http-errors'
import { z, ZodError, ZodIssue } from 'zod'

export interface CreateHandler<
  Params extends z.ZodTypeAny = any,
  Body extends z.ZodTypeAny = any,
  Query extends z.ZodTypeAny = any,
  ResJson extends z.ZodTypeAny = any,
  OmittedFields extends 'params' | 'body' | 'query' | '' = '',
> {
  (): {
    resolve(
      f: (
        req: Request<z.input<Params>, any, z.input<Body>>,
        res: Response<ResJson>,
      ) => any,
    ): (req: Request, res: Response) => void
    inputParams<T extends z.ZodTypeAny>(
      paramsSchema: T,
    ): Omit<
      ReturnType<
        CreateHandler<T, Body, Query, ResJson, OmittedFields | 'params'>
      >,
      OmittedFields | 'params'
    >
    inputBody<T extends z.ZodTypeAny>(
      bodySchema: T,
    ): Omit<
      ReturnType<
        CreateHandler<Params, T, Query, ResJson, OmittedFields | 'body'>
      >,
      OmittedFields | 'body'
    >
    inputQuery<T extends z.ZodTypeAny>(
      querySchema: T,
    ): Omit<
      ReturnType<
        CreateHandler<Params, Body, T, ResJson, OmittedFields | 'body'>
      >,
      OmittedFields | 'query'
    >
    responseJson<T extends z.ZodTypeAny>(
      resJsonSchema: T,
    ): Omit<
      ReturnType<
        CreateHandler<Params, Body, Query, z.input<T>, OmittedFields | 'body'>
      >,
      OmittedFields | 'query'
    >
  }
}

type HandlerThis = {
  paramsSchema: z.ZodAny
  bodySchema: z.ZodAny
  querySchema: z.ZodAny
  resJsonSchema: z.ZodAny
}

const formatZodError = {
  all: (prefix: string, zodError: ZodError) =>
    Object.values(
      zodError.flatten(
        (issue: ZodIssue) => `[${prefix}${issue.path}]: ${issue.message}`,
      ).fieldErrors,
    ).join(', '),
  first: (prefix: string, zodError: ZodError) =>
    Object.values(
      zodError.flatten(
        (issue: ZodIssue) => `[${prefix}${issue.path}]: ${issue.message}`,
      ).fieldErrors,
    )[0],
}

export const createHandler: CreateHandler = function () {
  return {
    inputParams(paramsSchema) {
      ;(<any>this).paramsSchema = paramsSchema
      return this
    },
    inputBody(bodySchema) {
      ;(<any>this).bodySchema = bodySchema
      return this
    },
    inputQuery(querySchema) {
      ;(<any>this).querySchema = querySchema
      return this
    },
    responseJson(resJsonSchema) {
      ;(<any>this).resJsonSchema = resJsonSchema
      return this
    },
    resolve(this: HandlerThis, f) {
      return async (req, res) => {
        if (this.paramsSchema) {
          const paramsValidationResult = this.paramsSchema.safeParse(req.params)
          if (!paramsValidationResult.success) {
            return res
              .status(400)
              .json(formatZodError.all('param.', paramsValidationResult.error))
          }

          req.params = paramsValidationResult.data
        }

        if (this.bodySchema) {
          const bodyValidationResult = this.bodySchema.safeParse(req.body)
          if (!bodyValidationResult.success) {
            return res
              .status(400)
              .json(formatZodError.all('body.', bodyValidationResult.error))
          }

          req.body = bodyValidationResult.data
        }

        if (this.querySchema) {
          const queryValidationResult = this.querySchema.safeParse(req.query)
          if (!queryValidationResult.success) {
            return res
              .status(400)
              .json(formatZodError.all('query.', queryValidationResult.error))
          }

          req.query = queryValidationResult.data
        }

        // Because express is very old and can't handle async exceptions, let's do it from this helper
        const resJson: any = res.json
        try {
          if (this.resJsonSchema) {
            const newRes: any = Object.assign(res, { json: undefined })
            newRes.json = (json: any) => {
              const resJsonValidationResult = this.resJsonSchema.safeParse(json)
              if (!resJsonValidationResult.success) {
                res.status(400)
                resJson.call(
                  res,
                  formatZodError.all('resJson.', resJsonValidationResult.error),
                )
                return
              }
              resJson.call(res, json)
            }
            await f(req, newRes as any)
            return
          }

          await f(req, res)
          return
        } catch (error) {
          if (createHttpError.isHttpError(error)) {
            res.status(error.status)
            resJson.call(res, { message: error.message })
          }
          return
        }
      }
    },
  }
}