import type { z, core as zodCore } from "zod";

// Utility type to "prettify" the resulting type to make it better serializable
// and avoid issues like: https://stackoverflow.com/a/68475392/5506547
// > The inferred type of this node exceeds the maximum length the compiler will serialize. An explicit type annotation is needed.
type Prettify<T> = { [K in keyof T]: T[K] } & {};

type OptionalKeys<T> = {
    [K in keyof T]-?: object extends Pick<T, K> ? K : never;
}[keyof T];

type MergeOptionalMarkers<Source, Target> = Prettify<
    Omit<Target, Extract<OptionalKeys<Source>, keyof Target>> &
        Partial<Pick<Target, Extract<OptionalKeys<Source>, keyof Target>>>
>;

/**
 * A refined Zod object is a Zod object that has been refined through the usage of a
 * `createRefinableSchema`. In general, all methods which mutate the object and discard refinements
 * are disabled.
 *
 * This is very important in terms of refinements and validations as using `.pick()` outside of the
 * `createRefinableSchema` function will discard all refinements.
 *
 * Example:
 *
 * ```ts
 * const createExampleSchema = createRefinableSchema({
 *     // [... schema definition]
 *     refine: (schema, context) => {
 *         // [... refinement logic]
 *     },
 * });
 *
 * const incorrectSchema = createExampleSchema(context).pick({ startDate: true });
 *                                               ^ Wrong: All refinements will be discarded and no validation will be performed!
 *
 * const correctSchema = createExampleSchema((schema) => schema.pick({ startDate: true }), context);
 *                                               ^ Correct: The refinement will be placed on the picked schema.
 * ```
 *
 * It additionally forbids the usage of the `shape` property on the refined Zod object as otherwise this could be "misused"
 * within a spread syntax as described here: https://zod.dev/api#extend
 *
 * Example:
 *
 * ```ts
 * const DogWithBreed = z.object({
 *   ...createAnimalSchema(context).shape,
 *   breed: z.string(),
 * });
 *
 * // Instead, do something like this:
 * const DogWithBreed = createAnimalSchema((schema) => schema.extend({ breed: z.string() }), context);
 *
 * // Or, if you want not to use `.extend()` and a more performant solution with object destructuring:
 * const DogWithBreed = createAnimalSchema((schema) => z.object({
 *   ...schema.shape, // <-- within the middleware you can safely use .shape to access the schema's properties
 *   breed: z.string(),
 * }), context);
 * ```
 *
 * @see https://github.com/colinhacks/zod/issues/5192
 * @see https://github.com/colinhacks/zod/discussions/4706
 * @see https://github.com/colinhacks/zod/issues/5425
 */
type RefinedZodObject<ZO extends z.ZodObject> = Pick<
    ZO,
    | "_zod"
    // | "and"
    // | "array"                  // We need to override the return type to reflect `this` to be also RefinedZodObject
    | "brand"
    | "catch"
    | "catchall"
    // | "check"
    // | "clone"
    | "decode"
    | "decodeAsync"
    | "def"
    // | "default"
    | "describe"
    | "description"
    | "encode"
    | "encodeAsync"
    // | "extend"
    | "keyof"
    // | "loose"
    // | "meta"                    // We need to override the return type to reflect `this` to be also RefinedZodObject
    // | "nonoptional"
    // | "nullable"
    // | "nullish"
    // | "omit"
    // | "openapi"
    | "optional"
    // | "or"
    // | "overwrite"
    | "parse"
    | "parseAsync"
    // | "partial"
    // | "pick"
    // | "pipe"
    // | "prefault"
    | "readonly"
    // | "refine"
    // | "register"
    // | "required"
    | "safeDecode"
    | "safeDecodeAsync"
    | "safeEncode"
    | "safeEncodeAsync"
    | "safeExtend"
    | "safeParse"
    | "safeParseAsync"
    // | "shape"
    // | "spa"
    // | "strict"
    // | "strip"
    // | "superRefine"
    // | "transform"
    | "type"
    | "~standard"
    | "_def"
    | "_input"
    | "_output"
    | "isNullable"
    | "isOptional"
    // | "merge"
    // | "passthrough"
> & {
    array: () => z.ZodArray<RefinedZodObject<ZO>>;
    meta(): z.core.$replace<z.core.GlobalMeta, ZO> | undefined;
    meta(...args: Parameters<ZO["meta"]>): RefinedZodObject<ZO>;
};

// Add explicit types to avoid hitting the maximum type length serialization
type CreateRefinableSchemaType<
    Args extends any[],
    SchemaInputWithRecursion extends z.ZodObject,
    SchemaOutputWithRecursion extends z.ZodObject,
> = ((...args: undefined[]) => RefinedZodObject<SchemaOutputWithRecursion>) &
    ((...args: Args) => RefinedZodObject<SchemaOutputWithRecursion>) &
    (<Result extends z.ZodObject>(
        middleware: (data: SchemaInputWithRecursion) => Result,
        ...args: Args
    ) => RefinedZodObject<Result>) & {
        identifier: symbol;
    };

/**
 * Since zod v4 the `refine()` and `check()` methods are not passthrough the e.g. `pick()` and `omit()` methods.
 * This function allows to create a schema that is refineable and can be used with the `pick()` and `omit()` methods.
 *
 * **Attention**: In your passed `refine` method make sure to check for the object's `shape`. Example:
 *
 * ```ts
 * const createSupportResponseMessageSchema = createRefinableSchema(
 *     () =>
 *         z.object({
 *             id: z.string(),
 *             message: z.string(),
 *             // [...]
 *         }),
 *     (schema) =>
 *         schema.refine(
 *             ({ startDate, endDate }) => {
 *                 // Check if the object has the `startDate` and `endDate` properties
 *                 if (schema.shape.startDate && schema.shape.endDate) {
 *                     return startDate <= endDate;
 *                 } else {
 *                     return true;
 *                 }
 *             },
 *             {
 *                 message: "Start date must be before end date",
 *                 path: ["startDate"],
 *             },
 *         ),
 * );
 * ```
 *
 * If you want to require additional context to your refinable schema (e.g. the `i18next` `t` function),
 * you can pass it as an argument to the `createRefinableSchema` function. Example:
 *
 * ```ts
 * const createSupportResponseMessageSchema = createRefinableSchema(
 *     ({ t }: CreateContractContext) =>
 *         z.object({
 *             id: z.string(),
 *             message: z.string(),
 *             // [...]
 *         }),
 * ```
 *
 * If you want to use this schema in a contract as an OpenAPI component, you should not use the `meta()` method in your
 * schema definition. Instead you should use the `meta()` method when you use the schema in a contract. Otherwise you will be
 * running into errors that the ID is already registered.
 *
 * ### Recursion
 *
 * If you want to use a schema in a recursive way, you can use the `recursionSchemas` option together with the `memoize` function (which
 * uses the mechanics of recursive-objects we learned in https://zod.dev/api?id=recursive-objects). As this gets too complex
 * to explain here, please refer to the [`README.md`](../../README.md) for an example.
 *
 * @see https://github.com/colinhacks/zod/discussions/4706
 * @see https://github.com/colinhacks/zod/issues/4427
 */
function createRefinableSchema<
    SchemaInput extends z.ZodObject,
    Args extends any[],
    SchemaOutput extends z.ZodObject = SchemaInput,
    RecursionMap = {
        [K in keyof SchemaInput["shape"] as SchemaInput["shape"][K] extends z.ZodCustom ? K : never]: z.ZodObject;
    } & {
        [K in keyof SchemaInput["shape"] as SchemaInput["shape"][K] extends z.ZodOptional<z.ZodCustom>
            ? K
            : never]?: z.ZodObject;
    },
    SchemaInputWithRecursion extends z.ZodObject = z.ZodObject<
        MergeOptionalMarkers<RecursionMap, zodCore.util.Extend<SchemaInput["shape"], RecursionMap>>
    >,
    SchemaOutputWithRecursion extends z.ZodObject = z.ZodObject<
        MergeOptionalMarkers<RecursionMap, zodCore.util.Extend<SchemaOutput["shape"], RecursionMap>>
    >,
>({
    schema,
    refine,
    recursionSchemas,
}: {
    schema: (...args: Args) => SchemaInput;
    refine?: (schema: SchemaInput, ...args: Args) => SchemaOutput;
    recursionSchemas?: (memoize: <T extends (...args: any[]) => any>(fn: T) => T, ...args: Args) => RecursionMap;
}): CreateRefinableSchemaType<Args, SchemaInputWithRecursion, SchemaOutputWithRecursion> {
    const identifier = Symbol("recursionMap");

    const fn = (...args: any[]) => {
        let middleware: (data: SchemaInput) => z.ZodObject;
        if (typeof args[0] === "function") {
            middleware = args.shift();
        }

        let s = schema(...(args as Args)) as z.ZodObject;

        // CU-8699tgz5g Runtime check: Ensure the schema hasn't been refined
        // The schema parameter should only return a base ZodObject, not a refined one
        // Refinements should be done in the separate `refine` parameter
        if (s && typeof s === "object" && "_def" in s && (s as any)._def?.checks) {
            throw new Error(
                "createRefinableSchema: Do not call .refine() within the 'schema' parameter. Move refinements to the separate 'refine' parameter instead.",
            );
        }

        let memoizeRecursionMap = new Map<symbol, any>();

        // Check if there is already a memoization map for recursion
        for (const arg of args) {
            if (arg instanceof Map) {
                const { value } = arg.keys().next();
                if (typeof value === "symbol" && value.description === "recursionMap") {
                    args.splice(args.indexOf(arg), 1);
                    memoizeRecursionMap = arg;
                    break;
                }
            }
        }

        if (recursionSchemas) {
            const memoize = (originalFn: (...args: any[]) => any) => {
                const { identifier } = originalFn as any;

                return (...args: any[]) => {
                    const cache = memoizeRecursionMap.get(identifier);
                    if (cache) {
                        return cache;
                    }

                    const resultOfOriginalFn = originalFn(...args, memoizeRecursionMap);
                    return resultOfOriginalFn;
                };
            };

            // With the recursion schema function we get the Zod schema which we need to replace
            // with the predefined `z.custom` schemas.
            const recursionSchemasMap = recursionSchemas(memoize as any, ...(args as Args));
            s = s
                .omit(
                    Object.keys(recursionSchemasMap).reduce(
                        (acc, key) => {
                            acc[key] = true;
                            return acc;
                        },
                        {} as Record<string, true>,
                    ),
                )
                .extend(recursionSchemasMap);

            memoizeRecursionMap.set(identifier, s);
        }

        const m = middleware ? middleware(s as any) : s;
        const r = (refine ? refine(m as any, ...(args as Args)) : m) as SchemaOutputWithRecursion;

        return r;
    };

    fn.identifier = identifier;

    return fn as CreateRefinableSchemaType<Args, SchemaInputWithRecursion, SchemaOutputWithRecursion>;
}

export type { CreateRefinableSchemaType, OptionalKeys, MergeOptionalMarkers, RefinedZodObject };
export { createRefinableSchema };
