import { ImplStatic, TraitInstance, TraitStaticMembers, trait } from "@thilawyn/traitify-ts"
import { Class, Jsonifiable } from "type-fest"
import { z } from "zod"
import { parseZodSchemaEffect } from "../util"
import { ZodSchemaObjectTrait } from "./ZodSchemaObject"


export type OfClass<
    Instance      extends OfClassInstance<T, Catchall, Values, PartialValues>,

    T             extends z.ZodRawShape,
    Catchall      extends z.ZodTypeAny,
    Values        extends object,
    PartialValues extends Partial<Values>,
> = (
    Class<Instance, [values: Values]> &
    TraitStaticMembers<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>>
)

export type OfClassInstance<
    T                 extends z.ZodRawShape,
    Catchall          extends z.ZodTypeAny,
    Values            extends object,
    PartialValues     extends Partial<Values>,
> = (
    Values & TraitInstance<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>>
)


export const JsonifiedZodSchemaObject = <
    Of                extends OfClass<OfInstance, T, Catchall, Values, PartialValues>,
    OfInstance        extends OfClassInstance<T, Catchall, Values, PartialValues>,
    T                 extends z.ZodRawShape,
    Catchall          extends z.ZodTypeAny,
    Values            extends object,
    PartialValues     extends Partial<Values>,

    JsonifyT          extends z.ZodRawShape,
    JsonifyCatchall   extends z.ZodTypeAny,
    DejsonifyT        extends z.ZodRawShape,
    DejsonifyCatchall extends z.ZodTypeAny,
    JsonifiedValues   extends object & Jsonifiable,
>(
    of: Of | OfClass<OfInstance, T, Catchall, Values, PartialValues>,

    props: {
        jsonifySchema: (
            schema: typeof of.schema
        ) => z.ZodObject<JsonifyT, "strip", JsonifyCatchall, JsonifiedValues, Values>,

        dejsonifySchema: (
            schema: typeof of.schema
        ) => z.ZodObject<DejsonifyT, "strip", DejsonifyCatchall, Values, JsonifiedValues>,
    },
) => trait
    .implement(Super => class JsonifiedZodSchemaObjectImpl extends Super {
        declare ["constructor"]: typeof JsonifiedZodSchemaObjectImpl

        static readonly of              = of as Of
        static readonly jsonifySchema   = props.jsonifySchema(of.schema)
        static readonly dejsonifySchema = props.dejsonifySchema(of.schema)


        static pipeSchemaIntoInstance<
            Instance          extends JsonifiedValues,

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


        static jsonify<
            Instance extends JsonifiedValues
        >(
            this: (
                Class<Instance, [values: JsonifiedValues]> &
                ImplStatic<typeof JsonifiedZodSchemaObjectImpl>
            ),
            values: Values,
            params?: Partial<z.ParseParams>,
        ) {
            return this
                .pipeSchemaIntoInstance(this.jsonifySchema)
                .parse(values, params)
        }

        static jsonifyPromise<
            Instance extends JsonifiedValues
        >(
            this: (
                Class<Instance, [values: JsonifiedValues]> &
                ImplStatic<typeof JsonifiedZodSchemaObjectImpl>
            ),
            values: Values,
            params?: Partial<z.ParseParams>,
        ) {
            return this
                .pipeSchemaIntoInstance(this.jsonifySchema)
                .parseAsync(values, params)
        }

        static jsonifyEffect<
            Instance extends JsonifiedValues
        >(
            this: (
                Class<Instance, [values: JsonifiedValues]> &
                ImplStatic<typeof JsonifiedZodSchemaObjectImpl>
            ),
            values: Values,
            params?: Partial<z.ParseParams>,
        ) {
            return parseZodSchemaEffect(
                this.pipeSchemaIntoInstance(this.jsonifySchema),
                values,
                params,
            )
        }


        dejsonify(params?: Partial<z.ParseParams>) {
            return this.constructor.of.pipeSchemaIntoInstance(
                this.constructor.dejsonifySchema
            ).parse(this, params)
        }

        dejsonifyPromise(params?: Partial<z.ParseParams>) {
            return this.constructor.of.pipeSchemaIntoInstance(
                this.constructor.dejsonifySchema
            ).parseAsync(this, params)
        }

        dejsonifyEffect(params?: Partial<z.ParseParams>) {
            return parseZodSchemaEffect(this.constructor.dejsonifySchema, this, params)
        }
    })
    .build()