import {CamelCase} from "type-fest";
import {ZodRawShape} from "zod/lib/types";
import {z, ZodDiscriminatedUnion, ZodLiteral, ZodObject, ZodSchema} from "zod";
import {
    EntityCore,
    EntityCoreShape,
    EntityCoreShapeType,
    EntitySymbol,
    Extender,
    UnsafeEntitySymbol,
    UnsafeZentity,
    Zentity,
    ZodShapeType
} from "./zentity";
import {toCamelCase} 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;


type ShapeRecord<S extends ZodRawShape> = Record<string, S>

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

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


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

type Clean<T> = T extends undefined ? {} : T

type OShape<B extends ZodRawShape, O extends ShapeRecord<ZodRawShape>, K extends keyof O> = O[K] & Clean<B>


type OType<B extends ZodRawShape, O extends ShapeRecord<ZodRawShape>, K extends keyof O> = ZodShapeType<OShape<B, O, K>>

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


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


type ZentityNewFn<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>, TD extends TransitionDef<KS>> = {
    [key in (keyof KS & string) as CC<`new`, key>]: (params: OType<BS, KS, key> & {
        id: string
    }) => Full<BS, KS, D, KSE, TD, key>
}
type ZentityCreateFn<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>, TD extends TransitionDef<KS>> = { [key in (keyof KS & string) as CC<`create`, key>]: (params: OType<BS, KS, key> & EntityCore) => Full<BS, KS, D, KSE, TD, key> }
type ZentityUnsafeFn<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>, TD extends TransitionDef<KS>> = { [key in (keyof KS & string) as CC<`unsafe`, key>]: (params: OType<BS, KS, key> & EntityCore) => UnsafeFull<BS, KS, D, KSE, TD, key> }
type ZentityGuardFn<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>, TD extends TransitionDef<KS>> = { [key in (keyof KS & string) as CC<'is', key>]: (data: any) => data is Full<BS, KS, D, KSE, TD, key> }

type ZentityShape<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string> = { [key in (keyof KS & string) as CC<key, 'shape'>]: OShape<BS, KS, key> & { [K in D]: ZodLiteral<key> } & typeof EntityCoreShape }
type ZentitySchema<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string> = { [key in (keyof KS & string) as CC<key, 'schema'>]: ZodObject<OShape<BS, KS, key> & { [K in D]: ZodLiteral<key> } & typeof EntityCoreShape> }


type Full<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KE extends KeyedExtender<BS, KS, any>, TD extends TransitionDef<KS>, key extends keyof KS> =
    Zentity<{ [Q in D]: key } & OType<BS, KS, key>, ReturnType<KE[key]> & Transition<BS, KS, KE, TD, key>>

type UnsafeFull<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KE extends KeyedExtender<BS, KS, any>, TD extends TransitionDef<KS>, key extends keyof KS> =
    UnsafeZentity<{ [Q in D]: key } & OType<BS, KS, key>, ReturnType<KE[key]> & Transition<BS, KS, KE, TD, key>>


type DiscriminatedZentityClass<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KE extends KeyedExtender<BS, KS, any>, TD extends TransitionDef<KS>> =
    {
        create: <T extends Flatten<KS, D>>(params: T & EntityCore & ZodShapeType<BS>) => Zentity<ZodShapeType<BS> & ExtractType<T, D>, ReturnType<KE[T[D]]> & Transition<BS, KS, KE, TD, T[D]>>,
        new: <T extends Flatten<KS, D>>(params: T & ZodShapeType<BS> & {
            id: string
        }) => Zentity<ExtractType<T, D>, ReturnType<KE[T[D]]> & Transition<BS, KS, KE, TD, T[D]>>,
        unsafe: <T extends Flatten<KS, D>>(params: T & ZodShapeType<BS>) => UnsafeZentity<ExtractType<T, D>, ReturnType<KE[T[D]]> & Transition<BS, KS, KE, TD, T[D]>>,
        baseShape: BS & typeof EntityCoreShape,
        schema: ZodDiscriminatedUnion<D, [ToZodD<BS, KS, D>, ...ToZodD<BS, KS, D>[]]>,
    }
    & ZentityNewFn<BS, KS, D, KE, TD>
    & ZentityCreateFn<BS, KS, D, KE, TD>
    & ZentityUnsafeFn<BS, KS, D, KE, TD>
    & ZentityGuardFn<BS, KS, D, KE, TD>
    & ZentitySchema<BS, KS, D>
    & ZentityShape<BS, KS, D>


type ValueUnion<T> = Required<T[keyof T]>;

type TransitionDef<KS extends ShapeRecord<ZodRawShape>> = {
    [key in keyof KS]: (keyof Omit<KS, key>) []
}


type ExclusiveToB<B, A> = {
    [K in keyof B as K extends keyof A ? never : K]: B[K];
};

type Transition<B extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, KE extends KeyedExtender<B, KS, any>, TD extends TransitionDef<KS>, K extends keyof KS> = {
    [key in (TD[K][number] & string) as CC<`to`, key>]: (param: ExclusiveToB<ZodShapeType<KS[key]>, ZodShapeType<KS[K]>>) => Zentity<OType<B, KS, key>, ReturnType<KE[key]> & Transition<B, KS, KE, TD, key>>
}


function typedObjectKeys<O extends object>(object: O): (keyof O)[] {
    return Object.keys(object) as (keyof O)[]
}

function typedObjectEntries<O extends object>(object: O): [keyof O, O[keyof O]][] {
    return Object.entries(object) as [keyof O, O[keyof O]][]
}


function toObject<T, K extends string, Y>(list: T[], fn: (item: T) => [K, Y]): Record<K, Y> {
    return list.reduce((acc, item) => {
        const [key, value] = fn(item)
        return {
            ...acc,
            [String(key)]: value
        }
    }, {} as Record<K, Y>)
}


export function statefulEntity<BS extends ZodRawShape, KS extends ShapeRecord<ZodRawShape>, D extends string, KSE extends KeyedExtender<BS, KS, any>, T extends TransitionDef<KS>>(options: {
    discriminator: D,
    base: BS,
    states: KS,
    extend: KSE,
    transitions: T,
}): DiscriminatedZentityClass<BS, KS, D, KSE, T> {
    const baseShape = {
        ...EntityCoreShape,
        ...options.base,
    }
    const unionSchema = z.discriminatedUnion<D, any>(
        options.discriminator,
        typedObjectEntries(options.states).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),
            ...transitionFnsForState(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 transitionFnsForState = (key: keyof KS, params: any) => {
        const toKeys = options.transitions[key] ?? []

        return toKeys
            .reduce((acc, key) => {
                return {
                    ...acc,
                    [toCamelCase(`to-${String(key)}`)]: transitionFn(key, params)
                }
            }, {})
    }

    const transitionFn = (to: keyof KS, params: any) => {
        return (extraParams: any) => {
            const destSchema = z.object({
                ...baseShape,
                [options.discriminator]: z.literal(to),
                ...options.states[to]
            })
            return createFn(destSchema, {
                ...params,
                ...extraParams,
                [options.discriminator]: to
            })
        }
    }


    const stateKeys = typedObjectKeys(options.states)


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


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

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


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

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


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

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


    return {
        ...stateCreateFn,
        ...stateNewFn,
        ...stateUnsafeFn,
        ...stateShapes,
        ...stateSchemas,
        ...stateIsFn,
        baseShape,
        schema: unionSchema,
        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 ShapeRecord<ZodRawShape>, D extends string> = ValueUnion<{
    [key in keyof KS]: ZodShapeType<KS[key]> & { [Q in D]: key }
}>

export type inferStatefulZentityTypeFromClass<T> = T extends DiscriminatedZentityClass<infer BS, infer KS, infer D, infer KE, infer TD> ?
    Zentity<UnionOfStateShapeTypesWithDiscriminator<KS, D> & Clean<BS>, ReturnType<KE[keyof KS]> & Transition<BS, KS, KE, TD, keyof KS>> : never;