zod-schema-class/src/builders/ZodSchemaClassExtender.ts

import { expression } from "@thilawyn/traitify-ts"
import { AbstractClass, Simplify } from "type-fest"
import { JsonifiableObject } from "type-fest/source/jsonifiable"
import { z } from "zod"
import { JsonifiableZodSchemaAbstractClass } from "../shapes/JsonifiableZodSchemaClass"
import { ZodSchemaAbstractClass } from "../shapes/ZodSchemaClass"
import { Extend, Override, StaticMembers } from "../util"


export class ZodSchemaClassExtender<
    Superclass extends AbstractClass<object>,
    Subclass   extends AbstractClass<object>,
> {
    declare ["constructor"]: typeof ZodSchemaClassExtender

    constructor(
        readonly superclass:  Superclass,
        readonly subclass:    Subclass,
    ) {}


    schema<
        Super                  extends ZodSchemaAbstractClass<any, SuperSchemaT, SuperSchemaUnknownKeys, SuperSchemaCatchall, SuperValues, SuperDefaultValues>,

        SuperSchemaT           extends z.ZodRawShape,
        SuperSchemaUnknownKeys extends z.UnknownKeysParam,
        SuperSchemaCatchall    extends z.ZodTypeAny,

        SuperValues            extends object,
        SuperDefaultValues     extends Partial<SuperValues>,

        SchemaT                extends z.ZodRawShape,
        SchemaUnknownKeys      extends z.UnknownKeysParam,
        SchemaCatchall         extends z.ZodTypeAny,

        Values                 extends SuperValues,
        DefaultValues          extends Partial<Values>,
    >(
        this: ZodSchemaClassExtender<
            Super | ZodSchemaAbstractClass<any, SuperSchemaT, SuperSchemaUnknownKeys, SuperSchemaCatchall, SuperValues, SuperDefaultValues>,
            any
        >,

        props: {
            schema:         (schema: Super["schema"]) => z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>
            defaultValues:  (defaultValues: SuperDefaultValues) => DefaultValues
        },
    ) {
        const schema        = props.schema(this.superclass.schema)
        const defaultValues = props.defaultValues(this.superclass.defaultValues)

        class Schemas extends (this.superclass as AbstractClass<object, any[]>) {
            static readonly schema        = schema
            static readonly defaultValues = defaultValues
        }

        return new this.constructor(
            this.superclass as Super,

            Schemas as unknown as AbstractClass<
                Simplify<
                    Extend<[InstanceType<Super>, Values]>
                >,

                [values: Values]
            > &
            Simplify<
                Override<[
                    StaticMembers<Super>,
                    StaticMembers<typeof Schemas>,
                ]>
            >
        )
    }

    jsonifiable<
        /** Superclass jsonifiable schemas */
        Super                           extends JsonifiableZodSchemaAbstractClass<
            any,

            SuperJsonifySchemaT,
            SuperJsonifySchemaUnknownKeys,
            SuperJsonifySchemaCatchall,

            SuperDejsonifySchemaT,
            SuperDejsonifySchemaUnknownKeys,
            SuperDejsonifySchemaCatchall,

            SuperJsonifiedValues,
            SuperValues
        >,

        SuperJsonifySchemaT             extends z.ZodRawShape,
        SuperJsonifySchemaUnknownKeys   extends z.UnknownKeysParam,
        SuperJsonifySchemaCatchall      extends z.ZodTypeAny,

        SuperDejsonifySchemaT           extends z.ZodRawShape,
        SuperDejsonifySchemaUnknownKeys extends z.UnknownKeysParam,
        SuperDejsonifySchemaCatchall    extends z.ZodTypeAny,

        SuperJsonifiedValues            extends JsonifiableObject,
        SuperValues                     extends object,

        /** New schemas */
        Self                           extends ZodSchemaAbstractClass<any, SelfSchemaT, SelfSchemaUnknownKeys, SelfSchemaCatchall, SelfValues, SelfDefaultValues>,

        SelfSchemaT                    extends z.ZodRawShape,
        SelfSchemaUnknownKeys          extends z.UnknownKeysParam,
        SelfSchemaCatchall             extends z.ZodTypeAny,

        SelfValues                     extends object,
        SelfDefaultValues              extends Partial<SelfValues>,

        /* New jsonifiable schemas */
        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 SuperJsonifiedValues,
        Values                          extends SelfValues,
    >(
        this: ZodSchemaClassExtender<
            Super | JsonifiableZodSchemaAbstractClass<
                any,

                SuperJsonifySchemaT,
                SuperJsonifySchemaUnknownKeys,
                SuperJsonifySchemaCatchall,

                SuperDejsonifySchemaT,
                SuperDejsonifySchemaUnknownKeys,
                SuperDejsonifySchemaCatchall,

                SuperJsonifiedValues,
                SuperValues
            >,

            Self | ZodSchemaAbstractClass<any, SelfSchemaT, SelfSchemaUnknownKeys, SelfSchemaCatchall, SelfValues, SelfDefaultValues>
        >,

        props: {
            jsonifySchema: (
                schema:         Self["schema"],
                jsonifySchema:  Super["jsonifySchema"],
            ) => z.ZodObject<JsonifySchemaT, JsonifySchemaUnknownKeys, JsonifySchemaCatchall, JsonifiedValues, Values>

            dejsonifySchema: (
                schema:           Self["schema"],
                dejsonifySchema:  Super["dejsonifySchema"],
            ) => z.ZodObject<DejsonifySchemaT, DejsonifySchemaUnknownKeys, DejsonifySchemaCatchall, Values, JsonifiedValues>
        },
    ) {
        const jsonifySchema   = props.jsonifySchema(this.subclass.schema, this.superclass.jsonifySchema)
        const dejsonifySchema = props.dejsonifySchema(this.subclass.schema, this.superclass.dejsonifySchema)

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

        return new this.constructor(
            this.superclass as Super,

            JsonifiableSchemas as unknown as AbstractClass<
                Simplify<
                    Override<[InstanceType<Self>, JsonifiableSchemas]>
                >,

                ConstructorParameters<
                    ZodSchemaAbstractClass<any, SelfSchemaT, SelfSchemaUnknownKeys, SelfSchemaCatchall, SelfValues, SelfDefaultValues>
                >
            > &
            Simplify<
                Override<[
                    StaticMembers<Self>,
                    StaticMembers<typeof JsonifiableSchemas>,
                ]>
            >,
        )
    }


    toClass() {
        return this.subclass
    }

    toExpressionBuilder() {
        return expression.extends(this.subclass)
    }
}