import { ImplStatic, expression } from "@thilawyn/traitify-ts"
import { AbstractClass } from "type-fest"
import { z } from "zod"
import { ZodSchemaObject } from "./ZodSchemaObject"


export const ExtendableZodSchemaObject = <
    T             extends z.ZodRawShape,
    Catchall      extends z.ZodTypeAny,
    Values        extends object,
    PartialValues extends Partial<Values>,
>(
    schemaWithDefaults: z.ZodObject<T, "strip", Catchall, Values, PartialValues>,
) => expression
    .expresses(ZodSchemaObject(schemaWithDefaults))
    .build()
    .subtrait()
        .implement(Super => class ExtendableZodSchemaObjectImpl extends Super {
            static extend<
                Self                  extends AbstractClass<ExtendableZodSchemaObjectImpl> & ImplStatic<typeof ExtendableZodSchemaObjectImpl>,

                ExtendedT             extends z.ZodRawShape,
                ExtendedCatchall      extends z.ZodTypeAny,
                ExtendedValues        extends Values,
                ExtendedPartialValues extends Partial<ExtendedValues>,
            >(
                this: Self,

                schemaWithDefaults: (
                    schemaWithDefaults: typeof this.schemaWithDefaults
                ) => z.ZodObject<ExtendedT, "strip", ExtendedCatchall, ExtendedValues, ExtendedPartialValues>,
            ) {
                return expression
                    .extends(this)
                    .expresses(ExtendableZodSchemaObject(schemaWithDefaults(this.schemaWithDefaults)))
            }
        })
        .build()