interface MemoOptions<F extends Fn, S extends unknown[] = unknown[]> {
    serialize?: (...args: Parameters<F>) => S;
}
interface MemoAsyncOptions<F extends Fn> extends MemoOptions<F> {
    external?: {
        get: (args: Parameters<F>) => Promise<Awaited<ReturnType<F>> | undefined | null>;

        set: (args: Parameters<F>, value: Awaited<ReturnType<F>>) => Promise<void>;

        remove: (args: Parameters<F>) => Promise<void>;

        clear: () => Promise<void>;
    };
}

type Fn = (...params: any[]) => any;

type AsyncFn = (...params: any[]) => Promise<any>;

interface MemoFunc<F extends Fn> {
    (...args: Parameters<F>): ReturnType<F>;

    get(...args: Parameters<F>): ReturnType<F>;

    raw(...args: Parameters<F>): ReturnType<F>;

    clear(...args: Parameters<F> | []): void | Promise<void>;
}

export const enum State {
    Empty,
    Ok,
    Waiting,
    Error
}

export interface Node<T extends Fn> {
    state: State;
    value: ReturnType<T> | undefined;
    error: unknown;
    primitive: Map<any, Node<T>>;
    reference: WeakMap<any, Node<T>>;
    callbacks?: Set<{ res: (value: ReturnType<T>) => void; rej: (error: unknown) => void }>;
}

function makeNode<T extends Fn>(): Node<T> {
    return {
        state: State.Empty,
        value: undefined,
        error: undefined,
        primitive: new Map(),
        reference: new WeakMap()
    };
}

function clearNode<T extends Fn>(node: Node<T> | undefined) {
    if (node) {
        node.state = State.Empty;
        node.value = undefined;
        node.error = undefined;
        node.primitive = new Map();
        node.reference = new WeakMap();
    }
}
function isPrimitiveType(value: unknown) {
    return (typeof value !== 'object' && typeof value !== 'function') || value === null;
}
function walkBase<T extends Fn, P extends any[] = Parameters<T>>(
    node: Node<T>,
    args: P,
    hooks: { makeNode: () => Node<T> | undefined }
): Node<T> | undefined {
    let cur = node;
    for (const arg of args) {
        if (isPrimitiveType(arg)) {
            if (cur.primitive.has(arg)) {
                cur = cur.primitive.get(arg)!;
            } else {
                const newNode = hooks.makeNode();
                if (newNode) {
                    cur.primitive.set(arg, newNode);
                    cur = newNode;
                } else {
                    return undefined;
                }
            }
        } else {
            if (cur.reference.has(arg)) {
                cur = cur.reference.get(arg)!;
            } else {
                const newNode = hooks.makeNode();
                if (newNode) {
                    cur.reference.set(arg, newNode);
                    cur = newNode;
                } else {
                    return undefined;
                }
            }
        }
    }
    return cur;
}

function walkAndCreate<T extends Fn, P extends any[] = Parameters<T>>(
    node: Node<T>,
    args: P
) {
    return walkBase(node, args, { makeNode })!;
}

function walkOrBreak<T extends Fn, P extends any[] = Parameters<T>>(node: Node<T>, args: P) {
    return walkBase(node, args, { makeNode: () => undefined });
}
export function memoAsync<F extends AsyncFn>(
    fn: F,
    options: MemoAsyncOptions<F> = {}
): MemoFunc<F> {
    const root = makeNode<F>();

    const memoFunc = async function (...args: Parameters<F>) {
        const path = options.serialize ? options.serialize(...args) : args;
        const cur = walkAndCreate<F, any[]>(root, path);

        if (cur.state === State.Ok) {
            return cur.value;
        } else if (cur.state === State.Error) {
            throw cur.error;
        } else if (cur.state === State.Waiting) {
            return new Promise((res, rej) => {
                if (!cur.callbacks) {
                    cur.callbacks = new Set();
                }
                cur.callbacks!.add({ res, rej });
            });
        } else {
            try {
                cur.state = State.Waiting;

                const external = options.external ? await options.external.get(args) : undefined;
                const value = external !== undefined && external !== null ? external : await fn(...args);

                cur.state = State.Ok;
                cur.value = value;

                if (options.external) {
                    await options.external.set(args, value);
                }

                for (const callback of cur.callbacks ?? []) {
                    callback.res(value);
                }

                return value;
            } catch (error) {
                cur.state = State.Error;
                cur.error = error;

                for (const callback of cur.callbacks ?? []) {
                    callback.rej(error);
                }

                throw error;
            }
        }
    } as MemoFunc<F>;

    memoFunc.get = (...args) => {
        return memoFunc(...args);
    };

    memoFunc.raw = (...args) => {
        return fn(...args) as ReturnType<F>;
    };

    memoFunc.clear = async (...args) => {
        if (args.length === 0) {
            clearNode(root);
            if (options.external) {
                await options.external.clear();
            }
        } else {
            const cur = walkOrBreak<F>(root, args as Parameters<F>);
            clearNode(cur);
            if (options.external) {
                await options.external.remove(args as Parameters<F>);
            }
        }
    };

    return memoFunc;
}