import { Trait, TraitExpression, TraitExpressionBuilder } from "@thilawyn/traitify-ts"
import { AbstractClass } from "type-fest"
import { JsonifiableObject } from "type-fest/source/jsonifiable"
import { z } from "zod"
import { ZodSchemaAbstractClass } from "../shapes/ZodSchemaClass"
import { DejsonifiableZodSchemaObject } from "../traits/DejsonifiableZodSchemaObject"
import { JsonifiableZodSchemaObject } from "../traits/JsonifiableZodSchemaObject"
import { ZodSchemaObject } from "../traits/ZodSchemaObject"
import { StaticMembers } from "../util"


export class ZodSchemaClassBuilder<
    Superclass   extends AbstractClass<object>,
    const Traits extends readonly Trait<any, any, any, any>[],
>
extends TraitExpressionBuilder<Superclass, Traits> {
    declare ["constructor"]: typeof ZodSchemaClassBuilder

    schema<
        Super                              extends AbstractClass<object, []> & { schema?: never, schemaWithDefaultValues?: never },

        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>,
    >(
        this: ZodSchemaClassBuilder<Super, Traits>,

        { schema, schemaWithDefaultValues }: {
            schema: z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>

            schemaWithDefaultValues: (
                schema: z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>
            ) => z.ZodObject<SchemaWithDefaultValuesT, SchemaWithDefaultValuesUnknownKeys, SchemaWithDefaultValuesCatchall, Values, PartialValues>
        },
    ) {
        class ZodSchemaObjectConstructor extends (this.expressionSuperclass as AbstractClass<object, []>) {
            constructor(values: Values) {
                super()
                Object.assign(this, values)
            }
        }

        return new this.constructor(
            ZodSchemaObjectConstructor as unknown as (
                AbstractClass<
                    InstanceType<Super> & Values,
                    [values: Values]
                > &
                StaticMembers<Super>
            ),

            [...this.expressionTraits, ZodSchemaObject(schema, schemaWithDefaultValues(schema))],
        )
    }

    jsonifiable<
        Super                              extends ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, SchemaWithDefaultValuesT, SchemaWithDefaultValuesUnknownKeys, SchemaWithDefaultValuesCatchall, Values, PartialValues>
                                         & { jsonifySchema?: never, dejsonifySchema?: never },
        Instance                           extends Values,

        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>,

        JsonifySchemaT                     extends z.ZodRawShape,
        JsonifySchemaUnknownKeys           extends z.UnknownKeysParam,
        JsonifySchemaCatchall              extends z.ZodTypeAny,

        DejsonifySchemaT                   extends z.ZodRawShape,
        DejsonifySchemaUnknownKeys         extends z.UnknownKeysParam,
        DejsonifySchemaCatchall            extends z.ZodTypeAny,

        JsonifiedValues                    extends JsonifiableObject,
    >(
        this: ZodSchemaClassBuilder<
            Super | ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, SchemaWithDefaultValuesT, SchemaWithDefaultValuesUnknownKeys, SchemaWithDefaultValuesCatchall, Values, PartialValues>,
            Traits
        >,

        props: {
            jsonifySchema: (
                schema: z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>
            ) => z.ZodObject<JsonifySchemaT, JsonifySchemaUnknownKeys, JsonifySchemaCatchall, JsonifiedValues, Values>

            dejsonifySchema: (
                schema: z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>
            ) => z.ZodObject<DejsonifySchemaT, DejsonifySchemaUnknownKeys, DejsonifySchemaCatchall, Values, JsonifiedValues>
        },
    ) {
        const jsonifySchema   = props.jsonifySchema(this.expressionSuperclass.schema)
        const dejsonifySchema = props.dejsonifySchema(this.expressionSuperclass.schema)

        class JsonifiableSchemas extends (this.expressionSuperclass as AbstractClass<object>) {
            static readonly jsonifySchema   = jsonifySchema
                   readonly jsonifySchema   = jsonifySchema
            static readonly dejsonifySchema = dejsonifySchema
                   readonly dejsonifySchema = dejsonifySchema
        }

        return new this.constructor(
            JsonifiableSchemas as unknown as (
                AbstractClass<
                    Instance & JsonifiableSchemas,

                    // TODO: for some reason, ConstructorParameters<Super> does not work here. Maybe try to find a fix?
                    ConstructorParameters<ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, SchemaWithDefaultValuesT, SchemaWithDefaultValuesUnknownKeys, SchemaWithDefaultValuesCatchall, Values, PartialValues>>
                > &
                StaticMembers<Super> &
                StaticMembers<typeof JsonifiableSchemas>
            ),

            [
                ...this.expressionTraits,
                JsonifiableZodSchemaObject,
                DejsonifiableZodSchemaObject,
            ],
        )
    }
}

export interface ZodSchemaClassBuilder<
    Superclass extends AbstractClass<object>,
    Traits     extends readonly Trait<any, any, any, any>[],
> {
    extends<
        Super extends AbstractClass<object>
    >(
        //                                        \/ Ensures `extends` can only be called once at the beginning
        this:        ZodSchemaClassBuilder<typeof TraitExpression.NullSuperclass, Traits>,
        superclass:  Super,
    ): ZodSchemaClassBuilder<Super, Traits>

    expresses<
        const T extends readonly Trait<
            TraitExpression<
                typeof TraitExpression.NullSuperclass,
                readonly Trait<any, any, any, any>[]
            >,
            any,
            any,
            any
        >[]
    >(
        ...traits: T
    ): ZodSchemaClassBuilder<
        Superclass,
        TraitExpressionBuilder.ExpressesReturnTypeTraits<Traits, T>
    >

    expressesFirst<
        const T extends readonly Trait<
            TraitExpression<
                typeof TraitExpression.NullSuperclass,
                readonly Trait<any, any, any, any>[]
            >,
            any,
            any,
            any
        >[]
    >(
        ...traits: T
    ): ZodSchemaClassBuilder<
        Superclass,
        TraitExpressionBuilder.ExpressesFirstReturnTypeTraits<Traits, T>
    >
}


export const zodSchemaClass = new ZodSchemaClassBuilder(TraitExpression.NullSuperclass, [])