import {CamelCase} from "type-fest";
import {ZodRawShape} from "zod/lib/types";
import {z, ZodDiscriminatedUnion, ZodLiteral, ZodNumber, ZodObject, ZodSchema, ZodString} from "zod";
import {
    EntityCore,
    EntityCoreShape,
    EntityCoreShapeType,
    EntitySymbol,
    Extender,
    UnsafeEntitySymbol,
    UnsafeZentity,
    Zentity,
    ZodShapeType
} from "./zentity";
import {Clean, toCamelCase, typedObjectEntries, typedObjectKeys, ValueUnion} from "./utils";

type ExtractType<T, D extends string> = T extends { [K in D]: infer V }
    ? { [K in keyof T]: K extends D ? V : T[K] }
    : never;


export type VariantsShape<S extends ZodRawShape> = Record<string, S>

export type KeyedExtender<B extends ZodRawShape, KS extends VariantsShape<any>, R extends object> = {
    [K in keyof KS]?: Extender<KS[K] & B, R>
}

type Flatten<O extends VariantsShape<ZodRawShape>, D extends string> = ValueUnion<{
    [K in keyof O]: ZodShapeType<O[K]> & { [Q in D]: K }
}>


type CamelCaseOf<P extends string, K extends string> = CamelCase<`${P}-${K}`>


type VariantShape<BASE_SHAPE extends ZodRawShape, VARIANTS_SHAPE extends VariantsShape<ZodRawShape>, VARIANT_KEY extends keyof VARIANTS_SHAPE> =
    VARIANTS_SHAPE[VARIANT_KEY]
    & Clean<BASE_SHAPE>


type VariantType<BASE_SHAPE extends ZodRawShape, VARIANTS_SHAPE extends VariantsShape<ZodRawShape>, VARIANT_KEY extends keyof VARIANTS_SHAPE> = ZodShapeType<VariantShape<BASE_SHAPE, VARIANTS_SHAPE, VARIANT_KEY>>

type ReturnType<E> = E extends (arg: any) => infer R ? R : {}


type ToZodD<BS extends ZodRawShape, KS extends VariantsShape<any>, D extends string> = ZodObject<ValueUnion<{
    [key in keyof KS]: KS[key] & { [K in D]: ZodLiteral<any> }
}> & BS & EntityCoreShapeType>


type PolymorphicNewFn<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>> = {
    [key in (keyof KS & string) as CamelCaseOf<`new`, key>]: (params: VariantType<BS, KS, key> & {
        id: string
    }) => PolymorphicEntityType<BS, KS, D, KSE, key>
}

type PolymorphicCreateFn<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>> = { [key in (keyof KS & string) as CamelCaseOf<`create`, key>]: (params: VariantType<BS, KS, key> & EntityCore) => PolymorphicEntityType<BS, KS, D, KSE, key> }
type PolymorphicUnsafeFn<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>> = { [key in (keyof KS & string) as CamelCaseOf<`unsafe`, key>]: (params: VariantType<BS, KS, key> & EntityCore) => UnsafePolymorphicEntityType<BS, KS, D, KSE, key> }
type PolymorphicGuardFn<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>> = { [key in (keyof KS & string) as CamelCaseOf<'is', key>]: (data: any) => data is PolymorphicEntityType<BS, KS, D, KSE, key> }

type PolymorphicShape<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string> = { [key in (keyof KS & string) as CamelCaseOf<key, 'shape'>]: VariantShape<BS, KS, key> & { [K in D]: ZodLiteral<key> } & typeof EntityCoreShape }
type PolymorphicSchema<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string> = { [key in (keyof KS & string) as CamelCaseOf<key, 'schema'>]: ZodObject<VariantShape<BS, KS, key> & { [K in D]: ZodLiteral<key> } & typeof EntityCoreShape> }


type PolymorphicEntityType<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KE extends KeyedExtender<BS, KS, any>, key extends keyof KS> =
    Zentity<{ [Q in D]: key } & VariantType<BS, KS, key>, ReturnType<KE[key]>>

type UnsafePolymorphicEntityType<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KE extends KeyedExtender<BS, KS, any>, key extends keyof KS> =
    UnsafeZentity<{ [Q in D]: key } & VariantType<BS, KS, key>, ReturnType<KE[key]>>


export type PolymorphicEntityClass<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KE extends KeyedExtender<BS, KS, any>> =
    {
        name: string,
        create: <T extends Flatten<KS, D>>(params: T & EntityCore & ZodShapeType<BS>) => Zentity<ZodShapeType<BS> & ExtractType<T, D>, ReturnType<KE[T[D]]>>,
        new: <T extends Flatten<KS, D>>(params: T & ZodShapeType<BS> & {
            id: string
        }) => Zentity<ExtractType<T, D>, ReturnType<KE[T[D]]>>,
        unsafe: <T extends Flatten<KS, D>>(params: T & ZodShapeType<BS>) => UnsafeZentity<ExtractType<T, D>, ReturnType<KE[T[D]]>>,
        baseShape: BS & typeof EntityCoreShape,
        schema: ZodDiscriminatedUnion<D, [ToZodD<BS, KS, D>, ...ToZodD<BS, KS, D>[]]>,
        variants: (keyof KS)[],
    }
    & PolymorphicNewFn<BS, KS, D, KE>
    & PolymorphicCreateFn<BS, KS, D, KE>
    & PolymorphicUnsafeFn<BS, KS, D, KE>
    & PolymorphicGuardFn<BS, KS, D, KE>
    & PolymorphicSchema<BS, KS, D>
    & PolymorphicShape<BS, KS, D>


export function polymorphicEntity<BS extends ZodRawShape, KS extends VariantsShape<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>>(options: {
    name: string,
    discriminator: D,
    base: BS,
    variants: KS,
    extend: KSE,
}): PolymorphicEntityClass<BS, KS, D, KSE> {
    const baseShape = {
        ...EntityCoreShape,
        ...options.base,
    }
    const unionSchema = z.discriminatedUnion<D, any>(
        options.discriminator,
        typedObjectEntries(options.variants).map(([key, shape]) => {
            return z.object({
                ...shape,
                ...baseShape,
                [options.discriminator]: z.literal(key)
            })
        })
    )

    const genEntityCore = () => ({
        createdAt: new Date(),
        updatedAt: new Date(),
    })


    const createFn = (schema: ZodSchema, params: any, unsafe: boolean = false) => {
        const data = unsafe ? params : schema.parse(params)

        return {
            _brand: UnsafeEntitySymbol,
            _safe: EntitySymbol,
            ...data,
            copy: copyFn(schema, data),
            validate: validateFn(schema, data),
            toPlain: () => data,
            ...options.extend[params[options.discriminator]]?.(data),
        }
    }
    const copyFn = (schema: ZodSchema, params: any) => {
        return (copyParams: any) => {
            return createFn(schema, {
                ...params,
                ...copyParams,
            })
        }
    }


    const validateFn = (schema: ZodSchema, params: any) => () => {
        return createFn(schema, params)
    }


    const stateKeys = typedObjectKeys(options.variants)


    const stateShapes: PolymorphicShape<BS, KS, D> = stateKeys
        .reduce((acc, key) => {
            return {
                ...acc,
                [toCamelCase(`${String(key)}-shape`)]: {
                    ...baseShape,
                    ...options.variants[key],
                    [options.discriminator]: z.literal(key),
                }
            }
        }, {}) as any;


    const stateSchemas: PolymorphicSchema<BS, KS, D> = stateKeys
        .reduce((acc, key) => {
            return {
                ...acc,
                [toCamelCase(`${String(key)}-schema`)]: z.object({
                    ...baseShape,
                    ...options.variants[key],
                    [options.discriminator]: z.literal(key),
                })
            }
        }, {}) as any;

    const stateCreateFn: PolymorphicCreateFn<BS, KS, D, KSE> = stateKeys
        .reduce((acc, key) => {
            return {
                ...acc,
                [toCamelCase(`create-${String(key)}`)]: (params: VariantType<BS, KS, typeof key> & EntityCore) => {
                    return createFn(z.object({
                        ...baseShape,
                        ...options.variants[key],
                        [options.discriminator]: z.literal(key),
                    }), {...params, [options.discriminator]: key})
                }
            }
        }, {}) as any;


    const stateNewFn: PolymorphicNewFn<BS, KS, D, KSE> = stateKeys
        .reduce((acc, key) => {
            return {
                ...acc,
                [toCamelCase(`new-${String(key)}`)]: (params: VariantType<BS, KS, typeof key>) => {
                    return createFn(z.object({
                            ...baseShape,
                            ...options.variants[key],
                            [options.discriminator]: z.literal(key),
                        }), {
                            ...params,
                            [options.discriminator]: key,
                            ...genEntityCore()
                        },
                    )
                }
            }
        }, {}) as any;

    const stateUnsafeFn: PolymorphicUnsafeFn<BS, KS, D, KSE> = stateKeys
        .reduce((acc, key) => {
            return {
                ...acc,
                [toCamelCase(`unsafe-${String(key)}`)]: (params: VariantType<BS, KS, typeof key> & EntityCore) => {
                    return createFn(z.object({
                            ...baseShape,
                            ...options.variants[key],
                            [options.discriminator]: z.literal(key),
                        }), {
                            ...params,
                            [options.discriminator]: key,
                            ...genEntityCore()
                        },
                        true,
                    )
                }
            }
        }, {}) as any;


    const stateIsFn: PolymorphicGuardFn<BS, KS, D, KSE> = stateKeys
        .reduce((acc, key) => {
            return {
                ...acc,
                [toCamelCase(`is-${String(key)}`)]: (params: any): params is PolymorphicEntityType<BS, KS, D, KSE, typeof key> => {
                    const schema = z.object({
                        ...baseShape,
                        ...options.variants[key],
                        [options.discriminator]: z.literal(key),
                    })

                    return schema.safeParse(params).success
                }
            }
        }, {}) as any;


    return {
        name: options.name,
        ...stateCreateFn,
        ...stateNewFn,
        ...stateUnsafeFn,
        ...stateShapes,
        ...stateSchemas,
        ...stateIsFn,
        baseShape,
        schema: unionSchema,
        variants: Object.keys(options.variants) as (keyof KS)[],
        create: <T extends Flatten<KS, D>>(params: T & EntityCore & ZodShapeType<BS>) => {
            return createFn(unionSchema, params)
        },
        new: <T extends Flatten<KS, D>>(params: T & ZodShapeType<BS> & { id: string }) => {
            return createFn(unionSchema, {
                ...params,
                ...genEntityCore(),
            })
        },
        unsafe: <T extends Flatten<KS, D>>(params: T & ZodShapeType<BS>) => {
            return createFn(unionSchema, params, true)
        },
    }
}


type UnionOfStateShapeTypesWithDiscriminator<KS extends VariantsShape<ZodRawShape>, D extends string> = ValueUnion<{
    [key in keyof KS]: ZodShapeType<KS[key]> & { [Q in D]: key }
}>

export type inferPolymorphicEntityTypeFromClass<T> = T extends PolymorphicEntityClass<infer BS, infer KS, infer D, infer KE> ?
    Zentity<UnionOfStateShapeTypesWithDiscriminator<KS, D> & ZodShapeType<Clean<BS>>, ReturnType<KE[keyof KS]>> : never;


export type inferSpecificPolymorphicEntityTypeFromClass<T extends {
    variants: any
}, K extends T['variants'][number]> = T extends PolymorphicEntityClass<infer BS, infer KS, infer D, infer KE> ?
    K extends keyof KS ?
        Zentity<ZodShapeType<KS[K]> & ZodShapeType<BS>, ReturnType<KE[K]>> : never : never;


const TestClass = polymorphicEntity({
    name: 'Test',
    discriminator: 'status',
    base: {
        gameNo: z.number()
    },
    extend: {},
    variants: {
        initial: {
            indexer: z.string()
        },
        finished: {
            multiplier: z.number().min(1),
        }
    },
})

type Test = inferPolymorphicEntityTypeFromClass<typeof TestClass>

type TestInitial = inferSpecificPolymorphicEntityTypeFromClass<typeof TestClass, "finished">

function testFn(test: Test, initial: TestInitial) {
    if (test.status == 'finished') {
        return test.gameNo + test.multiplier
    } else if (2 > 3) {
        return test.indexer.split("")
    }

    return initial.multiplier * 5

}