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


type CreateArgs<Input extends object> = (
    HasRequiredKeys<Input> extends true
        ?      [values: Input, params?: Partial<z.ParseParams>]
        : [] | [values: Input, params?: Partial<z.ParseParams>]
)


export const ZodSchemaObject = <
    T             extends z.ZodRawShape,
    Catchall      extends z.ZodTypeAny,
    Values        extends object,
    PartialValues extends Partial<Values>,
>(
    schemaWithDefaults: z.ZodObject<T, "strip", Catchall, Values, PartialValues>,
) => trait
    .implement(Super => {
        class ZodSchemaObjectImpl extends Super {
            static readonly schema             = stripZodObjectDefaults(schemaWithDefaults)
            static readonly schemaWithDefaults = schemaWithDefaults


            static pipeSchemaIntoInstance<
                Instance          extends Values,

                SchemaT           extends z.ZodRawShape,
                SchemaUnknownKeys extends z.UnknownKeysParam,
                SchemaCatchall    extends z.ZodTypeAny,
                SchemaOutput      extends Values,
                SchemaInput,
            >(
                this:    Class<Instance, [values: Values]>,
                schema:  z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, SchemaOutput, SchemaInput>,
            ) {
                return schema.transform(values => new this(values))
            }

            static pipeInstanceIntoSchema<
                Self              extends Class<Values, [values: Values]>,

                SchemaT           extends z.ZodRawShape,
                SchemaUnknownKeys extends z.UnknownKeysParam,
                SchemaCatchall    extends z.ZodTypeAny,
                SchemaOutput,
                SchemaInput       extends Values,
            >(
                this:    Self,
                schema:  z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, SchemaOutput, SchemaInput>,
            ) {
                return z.instanceof(this).pipe(schema)
            }


            static create<
                Instance extends Values
            >(
                this: Class<Instance, [values: Values]>,
                ...[values, params]: CreateArgs<PartialValues>
            ) {
                const t = implStaticInstantiableThis(ZodSchemaObjectImpl, this)

                return t
                    .pipeSchemaIntoInstance(t.schemaWithDefaults)
                    .parse(values, params)
            }

            static createPromise<
                Instance extends Values
            >(
                this: Class<Instance, [values: Values]>,
                ...[values, params]: CreateArgs<PartialValues>
            ) {
                const t = implStaticInstantiableThis(ZodSchemaObjectImpl, this)

                return t
                    .pipeSchemaIntoInstance(t.schemaWithDefaults)
                    .parseAsync(values, params)
            }

            static createEffect<
                Instance extends Values
            >(
                this: Class<Instance, [values: Values]>,
                ...[values, params]: CreateArgs<PartialValues>
            ) {
                const t = implStaticInstantiableThis(ZodSchemaObjectImpl, this)

                return parseZodSchemaEffect(
                    t.pipeSchemaIntoInstance(t.schemaWithDefaults),
                    values,
                    params,
                )
            }
        }

        return ZodSchemaObjectImpl
    })
    .build()


export type ZodSchemaObjectTrait<
    T             extends z.ZodRawShape,
    Catchall      extends z.ZodTypeAny,
    Values        extends object,
    PartialValues extends Partial<Values>,
> = (
    ReturnType<
        typeof ZodSchemaObject<T, Catchall, Values, PartialValues>
    >
)