import { ImplStatic, trait } from "@thilawyn/traitify-ts"
import { Class, HasRequiredKeys } from "type-fest"
import { z } from "zod"
import { parseZodSchemaEffect } 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 = <
    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:  z.ZodObject<SchemaWithDefaultValuesT, SchemaWithDefaultValuesUnknownKeys, SchemaWithDefaultValuesCatchall, Values, PartialValues>,
) => trait
    .implement(Super => class ZodSchemaObject extends Super {
        static readonly schema                  = schema
        static readonly schemaWithDefaultValues = schemaWithDefaultValues

        static transform<
            Instance    extends Values,

            T           extends z.ZodRawShape,
            UnknownKeys extends z.UnknownKeysParam,
            Catchall    extends z.ZodTypeAny,
            Output      extends Values,
            Input,
        >(
            this:    Class<Instance, [values: Values]>,
            schema:  z.ZodObject<T, UnknownKeys, Catchall, Output, Input>,
        ) {
            return schema.transform(values => new this(values))
        }


        static create<
            Instance extends Values
        >(
            this: (
                Class<Instance, [values: Values]> &
                ImplStatic<typeof ZodSchemaObject>
            ),
            ...[values, params]: CreateArgs<PartialValues>
        ) {
            return this
                .transform(this.schemaWithDefaultValues)
                .parse(values, params)
        }

        static createPromise<
            Instance extends Values
        >(
            this: (
                Class<Instance, [values: Values]> &
                ImplStatic<typeof ZodSchemaObject>
            ),
            ...[values, params]: CreateArgs<PartialValues>
        ) {
            return this
                .transform(this.schemaWithDefaultValues)
                .parseAsync(values, params)
        }

        static createEffect<
            Instance extends Values
        >(
            this: (
                Class<Instance, [values: Values]> &
                ImplStatic<typeof ZodSchemaObject>
            ),
            ...[values, params]: CreateArgs<PartialValues>
        ) {
            return parseZodSchemaEffect(
                this.transform(this.schemaWithDefaultValues),
                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
        >
    >
)