import { ImplStaticThis, trait } from "@thilawyn/traitify-ts"
import { Class, HasRequiredKeys } from "type-fest"
import { z } from "zod"
import { StaticMembers, parseZodSchemaEffect } from "../util"


type ParseParamsArgs = [] | [params: Partial<z.ParseParams>]

type NewZodSchemaInstanceArgs<Input extends object> = (
    HasRequiredKeys<Input> extends true
        ? [values: Input, ...args: ParseParamsArgs]
        : [] | [values: Input, ...args: ParseParamsArgs]
)


export const ZodSchemaObject = <
    SchemaT                            extends z.ZodRawShape,
    SchemaUnknownKeys                  extends z.UnknownKeysParam,
    SchemaCatchall                     extends z.ZodTypeAny,

    SchemaWithDefaultValuesT           extends z.ZodRawShape,
    SchemaWithDefaultValuesUnknownKeys extends z.UnknownKeysParam,
    SchemaWithDefaultValuesCatchall    extends z.ZodTypeAny,

    Values                             extends object,
    PartialValues                      extends Partial<Values>,
>(
    schema: z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>,

    schemaWithDefaultValues: (
        schema: z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>
    ) => z.ZodObject<SchemaWithDefaultValuesT, SchemaWithDefaultValuesUnknownKeys, SchemaWithDefaultValuesCatchall, Values, PartialValues>,
) => trait
    .implement(Super => class ZodSchemaObject extends Super {
        static readonly schema                   = schema
        static readonly schemaWithDefaultsValues = schemaWithDefaultValues(schema)

        static instantiationSchema<
            Self     extends StaticMembers<ImplStaticThis<typeof ZodSchemaObject>>,
            Instance extends Values,
        >(
            this: Class<Instance, [values: Values]> & Self
        ) {
            return this.schema.transform(values => new this(values))
        }

        static instantiationSchemaWithDefaultValues<
            Self     extends StaticMembers<ImplStaticThis<typeof ZodSchemaObject>>,
            Instance extends Values,
        >(
            this: Class<Instance, [values: Values]> & Self
        ) {
            return this.schemaWithDefaultsValues.transform(values => new this(values))
        }


        static create<
            Self     extends StaticMembers<ImplStaticThis<typeof ZodSchemaObject>>,
            Instance extends Values,
        >(
            this: Class<Instance, [values: Values]> & Self,
            ...[values, params]: NewZodSchemaInstanceArgs<PartialValues>
        ) {
            return this.instantiationSchemaWithDefaultValues().parse(values, params)
        }

        static createPromise<
            Self     extends StaticMembers<ImplStaticThis<typeof ZodSchemaObject>>,
            Instance extends Values,
        >(
            this: Class<Instance, [values: Values]> & Self,
            ...[values, params]: NewZodSchemaInstanceArgs<PartialValues>
        ) {
            return this.instantiationSchemaWithDefaultValues().parseAsync(values, params)
        }

        static createEffect<
            Self     extends StaticMembers<ImplStaticThis<typeof ZodSchemaObject>>,
            Instance extends Values,
        >(
            this: Class<Instance, [values: Values]> & Self,
            ...[values, params]: NewZodSchemaInstanceArgs<PartialValues>
        ) {
            return parseZodSchemaEffect(this.instantiationSchemaWithDefaultValues(), values, params)
        }
    })
    .build()


export type ZodSchemaObjectTrait<
    SchemaT                            extends z.ZodRawShape,
    SchemaUnknownKeys                  extends z.UnknownKeysParam,
    SchemaCatchall                     extends z.ZodTypeAny,

    SchemaWithDefaultValuesT           extends z.ZodRawShape,
    SchemaWithDefaultValuesUnknownKeys extends z.UnknownKeysParam,
    SchemaWithDefaultValuesCatchall    extends z.ZodTypeAny,

    Values                             extends object,
    PartialValues                      extends Partial<Values>,
> = (
    ReturnType<
        typeof ZodSchemaObject<
            SchemaT,
            SchemaUnknownKeys,
            SchemaCatchall,

            SchemaWithDefaultValuesT,
            SchemaWithDefaultValuesUnknownKeys,
            SchemaWithDefaultValuesCatchall,

            Values,
            PartialValues
        >
    >
)