import { Expr, Stmt, Sym } from "./ast.ts";
import { Pos } from "./Token.ts";

type SymMap = { [ident: string]: Sym };

class Syms {
    private syms: SymMap = {};

    public constructor(private parent?: Syms) {}

    public define(ident: string, sym: Sym) {
        this.syms[ident] = sym;
    }

    public definedLocally(ident: string): boolean {
        return ident in this.syms;
    }

    public get(ident: string): { ok: true; sym: Sym } | { ok: false } {
        if (ident in this.syms) {
            return { ok: true, sym: this.syms[ident] };
        }
        if (this.parent) {
            return this.parent.get(ident);
        }
        return { ok: false };
    }
}

export class Resolver {
    private root = new Syms();

    public resolve(stmts: Stmt[]) {
        const scopeSyms = new Syms(this.root);
        for (const stmt of stmts) {
            this.resolveStmt(stmt, scopeSyms);
        }
    }

    private resolveExpr(expr: Expr, syms: Syms) {
        if (expr.kind.type === "error") {
            return;
        }
        if (expr.kind.type === "ident") {
            this.resolveIdentExpr(expr, syms);
            return;
        }
        if (expr.kind.type === "binary") {
            this.resolveExpr(expr.kind.left, syms);
            this.resolveExpr(expr.kind.right, syms);
            return;
        }
        if (expr.kind.type === "block") {
            const childSyms = new Syms(syms);
            for (const stmt of expr.kind.stmts) {
                this.resolveStmt(stmt, childSyms);
            }
            if (expr.kind.expr) {
                this.resolveExpr(expr.kind.expr, childSyms);
            }
            return;
        }
        if (expr.kind.type === "group") {
            this.resolveExpr(expr.kind.expr, syms);
            return;
        }
        if (expr.kind.type === "field") {
            this.resolveExpr(expr.kind.subject, syms);
            return;
        }
        if (expr.kind.type === "index") {
            this.resolveExpr(expr.kind.subject, syms);
            this.resolveExpr(expr.kind.value, syms);
            return;
        }
        if (expr.kind.type === "call") {
            this.resolveExpr(expr.kind.subject, syms);
            for (const e of expr.kind.args) {
                this.resolveExpr(e, syms);
            }
            return;
        }
        if (expr.kind.type === "unary") {
            this.resolveExpr(expr.kind.subject, syms);
            return;
        }
        if (expr.kind.type === "if") {
            this.resolveExpr(expr.kind.cond, syms);
            this.resolveExpr(expr.kind.truthy, syms);
            if (expr.kind.falsy !== undefined) {
                this.resolveExpr(expr.kind.falsy, syms);
            }
            return;
        }
        if (expr.kind.type === "loop") {
            this.resolveExpr(expr.kind.body, syms);
            return;
        }
        if (
            expr.kind.type === "int" || expr.kind.type === "bool" ||
            expr.kind.type === "null" || expr.kind.type === "string" ||
            expr.kind.type === "sym"
        ) {
            return;
        }
    }

    private resolveIdentExpr(expr: Expr, syms: Syms) {
        if (expr.kind.type !== "ident") {
            throw new Error("expected ident");
        }
        const ident = expr.kind;
        const symResult = syms.get(ident.value);
        if (!symResult.ok) {
            this.reportUseOfUndefined(ident.value, expr.pos, syms);
            return;
        }
        const sym = symResult.sym;
        expr.kind = {
            type: "sym",
            ident: ident.value,
            defType: sym.type,
        };
        if (sym.stmt) {
            expr.kind.stmt = sym.stmt;
        }
        if (sym.param) {
            expr.kind.param = sym.param;
        }
    }

    private resolveStmt(stmt: Stmt, syms: Syms) {
        if (stmt.kind.type === "error") {
            return;
        }
        if (stmt.kind.type === "let") {
            this.resolveLetStmt(stmt, syms);
            return;
        }
        if (stmt.kind.type === "fn") {
            this.resolveFnStmt(stmt, syms);
            return;
        }
        if (stmt.kind.type === "return") {
            if (stmt.kind.expr) {
                this.resolveExpr(stmt.kind.expr, syms);
            }
            return;
        }
        if (stmt.kind.type === "break") {
            if (stmt.kind.expr !== undefined) {
                this.resolveExpr(stmt.kind.expr, syms);
            }
            return;
        }
        if (stmt.kind.type === "assign") {
            this.resolveExpr(stmt.kind.subject, syms);
            this.resolveExpr(stmt.kind.value, syms);
            return;
        }
        if (stmt.kind.type === "expr") {
            this.resolveExpr(stmt.kind.expr, syms);
            return;
        }
    }

    private resolveLetStmt(stmt: Stmt, syms: Syms) {
        if (stmt.kind.type !== "let") {
            throw new Error("expected let statement");
        }
        this.resolveExpr(stmt.kind.value, syms);
        const ident = stmt.kind.param.ident;
        if (syms.definedLocally(ident)) {
            this.reportAlreadyDefined(ident, stmt.pos, syms);
            return;
        }
        syms.define(ident, {
            ident,
            type: "let",
            pos: stmt.kind.param.pos,
            stmt,
            param: stmt.kind.param,
        });
    }

    private resolveFnStmt(stmt: Stmt, syms: Syms) {
        if (stmt.kind.type !== "fn") {
            throw new Error("expected fn statement");
        }
        if (syms.definedLocally(stmt.kind.ident)) {
            this.reportAlreadyDefined(stmt.kind.ident, stmt.pos, syms);
            return;
        }
        const ident = stmt.kind.ident;
        syms.define(ident, {
            ident: stmt.kind.ident,
            type: "fn",
            pos: stmt.pos,
            stmt,
        });
        const fnScopeSyms = new Syms(syms);
        for (const param of stmt.kind.params) {
            if (fnScopeSyms.definedLocally(param.ident)) {
                this.reportAlreadyDefined(param.ident, param.pos, syms);
                continue;
            }
            fnScopeSyms.define(param.ident, {
                ident: param.ident,
                type: "fn_param",
                pos: param.pos,
                param,
            });
        }
        this.resolveExpr(stmt.kind.body, fnScopeSyms);
    }

    private reportUseOfUndefined(ident: string, pos: Pos, syms: Syms) {
        console.error(
            `use of undefined symbol '${ident}' at ${pos.line}${pos.col}`,
        );
    }

    private reportAlreadyDefined(ident: string, pos: Pos, syms: Syms) {
        console.error(
            `symbol already defined '${ident}', at ${pos.line}${pos.col}`,
        );
        const prev = syms.get(ident);
        if (!prev.ok) {
            throw new Error("expected to be defined");
        }
        if (!prev.sym.pos) {
            return;
        }
        const { line: prevLine, col: prevCol } = prev.sym.pos;
        console.error(
            `previous definition of '${ident}' at ${prevLine}:${prevCol}`,
        );
    }
}