import * as ast from "../ast.ts";
import { FileReporter } from "../diagnostics.ts";

export class Syms {
    constructor(
        private symMap: Map<number, Sym>,
    ) {}

    get(node: ast.Node): Sym {
        if (!this.symMap.has(node.id)) {
            throw new Error(`'${node.kind.tag}' not resolved`);
        }
        return this.symMap.get(node.id)!;
    }
}

export type Sym =
    | { tag: "Error" }
    | { tag: "Bool"; value: boolean }
    | { tag: "Builtin"; id: string }
    | { tag: "Fn"; stmt: ast.NodeWithKind<"FnStmt"> }
    | {
        tag: "FnParam";
        stmt: ast.NodeWithKind<"FnStmt">;
        param: ast.NodeWithKind<"Param">;
        idx: number;
    }
    | {
        tag: "Let";
        stmt: ast.NodeWithKind<"LetStmt">;
        param: ast.NodeWithKind<"Param">;
    }
    | { tag: "Loop"; stmt: ast.Node }
    | { tag: "BuiltinTy"; ident: string };

export function resolve(
    file: ast.Node,
    reporter: FileReporter,
): Syms {
    let syms = ResolverSyms.root();
    const resols = new Map<number, Sym>();

    const loopStack: ast.Node[] = [];

    file.visit({
        visit(node) {
            const k = node.kind;

            if (k.tag === "File" || k.tag === "Block") {
                syms = ResolverSyms.forkFrom(syms);
                for (const stmt of k.stmts) {
                    if (stmt.is("FnStmt")) {
                        syms.define(stmt.kind.ident, { tag: "Fn", stmt });
                    }
                }
                node.visitBelow(this);
                syms = syms.parent!;
                return "break";
            }

            if (k.tag === "FnStmt") {
                ast.assertNodeWithKind(node, "FnStmt");
                syms = ResolverSyms.forkFrom(syms);
                for (const [idx, param] of k.params.entries()) {
                    ast.assertNodeWithKind(param, "Param");
                    const sym: Sym = { tag: "FnParam", stmt: node, param, idx };
                    syms.define(param.kind.ident, sym);
                    resols.set(param.id, sym);
                }
                node.visitBelow(this);
                syms = syms.parent!;
                return "break";
            }

            if (k.tag === "LetStmt") {
                const stmt = node as ast.NodeWithKind<"LetStmt">;
                const param = k.param as ast.NodeWithKind<"Param">;
                const sym: Sym = { tag: "Let", stmt, param };
                syms.define(param.kind.ident, sym);
                resols.set(param.id, sym);
            }

            if (k.tag === "IdentExpr") {
                const sym = syms.resolveExpr(k.ident);
                if (sym === null) {
                    reporter.error(node.loc, `undefined symbol '${k.ident}'`);
                    reporter.abort();
                }
                resols.set(node.id, sym);
            }

            if (k.tag === "WhileStmt") {
                loopStack.push(node);
                node.visitBelow(this);
                loopStack.pop();
                return "break";
            }

            if (k.tag === "BreakStmt") {
                const loopNode = loopStack.at(-1);
                if (!loopNode) {
                    reporter.error(node.loc, `break outside loop`);
                    reporter.abort();
                }
                resols.set(node.id, { tag: "Loop", stmt: loopNode });
            }

            if (k.tag === "IdentTy") {
                const sym = syms.resolveTy(k.ident);
                if (sym === null) {
                    reporter.error(node.loc, `undefined symbol '${k.ident}'`);
                    reporter.abort();
                }
                resols.set(node.id, sym);
            }
        },
    });

    return new Syms(resols);
}

class ResolverSyms {
    static root(): ResolverSyms {
        return new ResolverSyms(
            new Map([
                ["len", { tag: "Builtin", id: "len" }],
                ["print", { tag: "Builtin", id: "print" }],
            ]),
            null,
        );
    }
    static forkFrom(parent: ResolverSyms): ResolverSyms {
        return new ResolverSyms(
            new Map(),
            parent,
        );
    }

    private constructor(
        private syms = new Map<string, Sym>(),
        public parent: ResolverSyms | null,
    ) {}

    define(ident: string, sym: Sym) {
        this.syms.set(ident, sym);
    }

    resolveExpr(ident: string): Sym | null {
        if (ident === "false" || ident === "true") {
            return { tag: "Bool", value: ident === "true" };
        }

        if (this.syms.has(ident)) {
            return this.syms.get(ident)!;
        }
        if (this.parent) {
            return this.parent.resolveExpr(ident);
        }
        return null;
    }

    resolveTy(ident: string): Sym | null {
        const builtins: string[] = [
            "void",
            "bool",
            "str",
            "i8",
            "i16",
            "i32",
            "i64",
            "isize",
            "u8",
            "u16",
            "u32",
            "u64",
            "usize",
        ];
        if (builtins.includes(ident)) {
            return { tag: "BuiltinTy", ident } as Sym;
        }

        if (this.syms.has(ident)) {
            return this.syms.get(ident)!;
        }
        if (this.parent) {
            return this.parent.resolveExpr(ident);
        }
        return null;
    }
}