import { trait } from "@thilawyn/traitify-ts"
import { AbstractClass } from "type-fest"
import { z } from "zod"
import { ZodSchemaAbstractClass } from "../shapes/ZodSchemaClass"
import { Extend, StaticMembers } from "../util"


export const ExtendableZodSchemaObject = trait
    .implement(Super => class ExtendableZodSchemaObject extends Super {
        static extend<
            Super                  extends ZodSchemaAbstractClass<SuperInstance, SuperSchemaT, SuperSchemaUnknownKeys, SuperSchemaCatchall, SuperValues, SuperDefaultValues>,
            SuperInstance          extends SuperValues,

            SuperSchemaT           extends z.ZodRawShape,
            SuperSchemaUnknownKeys extends z.UnknownKeysParam,
            SuperSchemaCatchall    extends z.ZodTypeAny,
            SuperValues            extends {},
            SuperDefaultValues     extends Partial<SuperValues>,

            SchemaT                extends z.ZodRawShape,
            SchemaUnknownKeys      extends z.UnknownKeysParam,
            SchemaCatchall         extends z.ZodTypeAny,
            Values                 extends SuperValues,
            DefaultValues          extends Partial<Values>,
        >(
            this: Super | ZodSchemaAbstractClass<SuperInstance, SuperSchemaT, SuperSchemaUnknownKeys, SuperSchemaCatchall, SuperValues, SuperDefaultValues>,

            props: {
                schema: (props: {
                    schema:  Super["schema"]
                    shape:   Super["schema"]["shape"]
                }) => z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>

                defaultValues: (defaultValues: SuperDefaultValues) => DefaultValues
            },
        ): (
            AbstractClass<
                Extend<[SuperInstance, Values]>,
                [values: Values]
            > &
            Extend<[
                StaticMembers<Super>,
                {
                    readonly schema:         z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>,
                    readonly defaultValues:  DefaultValues,
                },
            ]>
        ) {
            const schema        = props.schema({ schema: this.schema, shape: this.schema.shape })
            const defaultValues = props.defaultValues(this.defaultValues)

            return class extends this {
                static readonly schema        = schema
                static readonly defaultValues = defaultValues
            } as any
        }
    })
    .build()