import { EType, Expr, Stmt } from "./ast.ts";
import {
    AstVisitor,
    visitEType,
    visitExpr,
    visitParam,
    VisitRes,
    visitStmt,
    visitStmts,
} from "./ast_visitor.ts";
import { printStackTrace, Reporter } from "./info.ts";
import {
    EntryModSyms,
    FnSyms,
    LeafSyms,
    ModSyms,
    Syms,
} from "./resolver_syms.ts";
import { Pos } from "./token.ts";

export class Resolver implements AstVisitor<[Syms]> {
    public constructor(private reporter: Reporter) {
    }

    public resolve(stmts: Stmt[]): VisitRes {
        const syms = new EntryModSyms();
        this.scout(stmts, syms);
        visitStmts(stmts, this, syms);
        return "stop";
    }

    private scout(stmts: Stmt[], syms: Syms) {
        for (const stmt of stmts) {
            if (stmt.kind.type === "fn") {
                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,
                });
            } else if (stmt.kind.type === "type_alias") {
                const ident = stmt.kind.param.ident;
                if (syms.definedLocally(ident)) {
                    this.reportAlreadyDefined(ident, stmt.pos, syms);
                    return;
                }
                syms.define(ident, {
                    ident,
                    type: "type_alias",
                    pos: stmt.kind.param.pos,
                    stmt,
                    param: stmt.kind.param,
                });
            }
        }
    }

    visitModStmt(stmt: Stmt, syms: Syms): VisitRes {
        if (stmt.kind.type !== "mod") {
            throw new Error("expected let statement");
        }
        const modSyms = new ModSyms(syms);
        const { mod, ident } = stmt.kind;
        this.scout(mod.ast, modSyms);
        visitStmts(mod.ast, this, modSyms);

        if (syms.definedLocally(ident)) {
            this.reportAlreadyDefined(ident, stmt.pos, syms);
            return;
        }
        syms.define(ident, {
            type: "mod",
            ident,
            pos: stmt.pos,
            syms: modSyms,
        });

        return "stop";
    }

    visitLetStmt(stmt: Stmt, syms: Syms): VisitRes {
        if (stmt.kind.type !== "let") {
            throw new Error("expected let statement");
        }
        visitExpr(stmt.kind.value, this, 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,
        });
        return "stop";
    }

    visitTypeAliasStmt(stmt: Stmt, syms: Syms): VisitRes {
        if (stmt.kind.type !== "type_alias") {
            throw new Error("expected type_alias statement");
        }
        // nothing to do here
    }

    visitFnStmt(stmt: Stmt, syms: Syms): VisitRes {
        if (stmt.kind.type !== "fn") {
            throw new Error("expected fn statement");
        }
        const fnScopeSyms = new FnSyms(syms);
        for (const param of stmt.kind.genericParams ?? []) {
            if (fnScopeSyms.definedLocally(param.ident)) {
                this.reportAlreadyDefined(param.ident, param.pos, syms);
                continue;
            }
            fnScopeSyms.define(param.ident, {
                ident: param.ident,
                type: "generic",
                pos: param.pos,
                stmt,
                genericParam: param,
            });
        }
        for (const param of stmt.kind.params) {
            if (fnScopeSyms.definedLocally(param.ident)) {
                this.reportAlreadyDefined(param.ident, param.pos, syms);
                continue;
            }
            visitParam(param, this, fnScopeSyms);
            fnScopeSyms.define(param.ident, {
                ident: param.ident,
                type: "fn_param",
                pos: param.pos,
                param,
            });
        }
        if (stmt.kind.returnType) {
            visitEType(stmt.kind.returnType, this, fnScopeSyms);
        }
        visitExpr(stmt.kind.body, this, fnScopeSyms);
        return "stop";
    }

    visitIdentExpr(expr: Expr, syms: Syms): VisitRes {
        if (expr.kind.type !== "ident") {
            throw new Error("expected ident");
        }
        const ident = expr.kind.ident;
        const symResult = syms.get(ident);
        if (!symResult.ok) {
            this.reportUseOfUndefined(ident, expr.pos, syms);
            return;
        }
        const sym = symResult.sym;
        expr.kind = { type: "sym", ident, sym };
        return "stop";
    }

    visitPathExpr(expr: Expr, syms: Syms): VisitRes {
        if (expr.kind.type !== "path") {
            throw new Error("expected ident");
        }
        visitExpr(expr.kind.subject, this, syms);
        if (expr.kind.subject.kind.type !== "sym") {
            throw new Error("this error is not handled properly");
        }
        const subjectSym = expr.kind.subject.kind.sym;
        if (subjectSym.type !== "mod") {
            this.reporter.reportError({
                reporter: "Resolver",
                msg: `path expression are not implemented for '${subjectSym.type}' symbols`,
                pos: expr.pos,
            });
            printStackTrace();
            return "stop";
        }
        const getRes = subjectSym.syms.get(expr.kind.ident);
        if (!getRes.ok) {
            this.reportUseOfUndefined(
                expr.kind.ident,
                expr.pos,
                subjectSym.syms,
            );
            return "stop";
        }

        expr.kind = {
            type: "sym",
            ident: expr.kind.ident,
            sym: getRes.sym,
        };

        return "stop";
    }

    visitBlockExpr(expr: Expr, syms: Syms): VisitRes {
        if (expr.kind.type !== "block") {
            throw new Error();
        }
        const childSyms = new LeafSyms(syms);
        this.scout(expr.kind.stmts, childSyms);
        visitStmts(expr.kind.stmts, this, childSyms);
        if (expr.kind.expr) {
            visitExpr(expr.kind.expr, this, childSyms);
        }
        return "stop";
    }

    visitForExpr(expr: Expr, syms: Syms): VisitRes {
        if (expr.kind.type !== "for") {
            throw new Error();
        }
        const childSyms = new LeafSyms(syms);
        if (expr.kind.decl) visitStmt(expr.kind.decl, this, syms);
        if (expr.kind.cond) visitExpr(expr.kind.cond, this, syms);
        if (expr.kind.incr) visitStmt(expr.kind.incr, this, syms);
        visitExpr(expr.kind.body, this, childSyms);

        return "stop";
    }

    visitIdentEType(etype: EType, syms: Syms): VisitRes {
        if (etype.kind.type !== "ident") {
            throw new Error();
        }
        const ident = etype.kind.ident;
        const symResult = syms.get(ident);
        if (!symResult.ok) {
            this.reportUseOfUndefined(ident, etype.pos, syms);
            return;
        }
        const sym = symResult.sym;
        etype.kind = { type: "sym", ident, sym };
        return "stop";
    }

    private reportUseOfUndefined(ident: string, pos: Pos, _syms: Syms) {
        this.reporter.reportError({
            reporter: "Resolver",
            msg: `use of undefined symbol '${ident}'`,
            pos,
        });
        printStackTrace();
    }

    private reportAlreadyDefined(ident: string, pos: Pos, syms: Syms) {
        this.reporter.reportError({
            reporter: "Resolver",
            msg: `symbol already defined '${ident}'`,
            pos,
        });
        const prev = syms.get(ident);
        if (!prev.ok) {
            throw new Error("expected to be defined");
        }
        if (!prev.sym.pos) {
            return;
        }
        this.reporter.addNote({
            reporter: "Resolver",
            msg: `previous definition of '${ident}'`,
            pos: prev.sym.pos,
        });
        printStackTrace();
    }
}