import { Ctx, DefId, File, IdentId, idKey, Mod } from "../ctx.ts";
import * as ast from "../ast/ast.ts";
import {
    Block,
    Expr,
    Ident,
    Item,
    Pat,
    Path,
    PathSegment,
    QPath,
    Stmt,
    Ty,
} from "./hir.ts";
import { exhausted, Res as Result, todo } from "../util.ts";
import { Rib, RibKind } from "./rib.ts";
import { Res } from "./res.ts";

export class AstLowerer {
    private tyRibs: Rib[] = [];
    private valRibs: Rib[] = [];
    private currentFile!: File;
    private currentMod!: Mod;

    public constructor(
        private ctx: Ctx,
    ) {}

    public lower() {
        const file = this.ctx.entryFile();
        this.currentFile = file;
        this.currentMod = this.ctx.internMod({
            defKind: { tag: "mod" },
            ident: this.ctx.internIdent("root"),
            items: new Set(),
            defs: new Map(),
        });
        const ast = this.ctx.fileInfo(file).ast!;
        this.lowerFile(ast);
    }

    private lowerFile(file: ast.File) {
        this.lowerStmts(file.stmts);
    }

    private lowerStmts(stmts: ast.Stmt[]): Stmt[] {
        return stmts.map((stmt) => this.lowerStmt(stmt));
    }

    private lowerStmt(stmt: ast.Stmt): Stmt {
        const span = stmt.span;
        const kind = stmt.kind;
        switch (kind.tag) {
            case "error":
                return this.ctx.internStmt(kind, span);
            case "item":
                return this.ctx.internStmt({
                    tag: "item",
                    item: this.lowerItem(kind.item),
                }, span);
            case "let":
                return this.lowerLetStmt(stmt, kind);
            case "return":
                return this.ctx.internStmt({
                    tag: "return",
                    expr: kind.expr && this.lowerExpr(kind.expr),
                }, span);
            case "break":
                return this.ctx.internStmt({
                    tag: "break",
                    expr: kind.expr && this.lowerExpr(kind.expr),
                }, span);
            case "continue":
                return this.ctx.internStmt({
                    tag: "continue",
                }, span);
            case "assign":
                return this.ctx.internStmt({
                    tag: "assign",
                    assignType: kind.assignType,
                    subject: this.lowerExpr(kind.subject),
                    value: this.lowerExpr(kind.value),
                }, span);
            case "expr":
                return this.ctx.internStmt({
                    tag: "expr",
                    expr: this.lowerExpr(kind.expr),
                }, span);
        }
        exhausted(kind);
    }

    private lowerLetStmt(stmt: ast.Stmt, kind: ast.LetStmt): Stmt {
        const expr = kind.expr && this.lowerExpr(kind.expr);
        const ty = kind.ty && this.lowerTy(kind.ty);
        this.pushRib({ tag: "normal" });
        const pat = this.lowerPat(kind.pat);
        return this.ctx.internStmt({ tag: "let", pat, ty, expr }, stmt.span);
    }

    private lowerItem(item: ast.Item): Item {
        const { ident, kind, pub, span } = item;
        switch (kind.tag) {
            case "error":
                return this.ctx.internItem({ ident, kind, pub, span });
            case "mod_block": {
                const parent = this.currentMod;
                const mod = this.currentMod = this.ctx.internMod({
                    parent,
                    ident: ident.id,
                    defKind: { tag: "mod" },
                    items: new Set(),
                    defs: new Map(),
                });
                const point = this.ribPoint();
                this.pushRib({ tag: "mod", mod });
                const _block = this.lowerBlock(kind.block);
                this.returnToRibPoint(point);
                this.currentMod = parent;
                return this.ctx.internItem({
                    ident,
                    kind: { tag: "mod", mod },
                    pub,
                    span,
                });
            }
            case "mod_file": {
                const parent = this.currentMod;
                const mod = this.currentMod = this.ctx.internMod({
                    parent,
                    ident: ident.id,
                    defKind: { tag: "mod" },
                    items: new Set(),
                    defs: new Map(),
                });
                const point = this.ribPoint();
                this.pushRib({ tag: "mod", mod });
                const parentFile = this.currentFile;
                const fileInfo = this.ctx.fileInfo(kind.file!);
                this.lowerFile(fileInfo.ast!);
                this.returnToRibPoint(point);
                this.currentFile = parentFile;
                this.currentMod = parent;
                if (this.tyRib().bindings.has(idKey(ident.id))) {
                    throw new Error();
                }
                this.tyRib().bindings.set(idKey(ident.id), {
                    tag: "def",
                    id: mod.id,
                    kind: { tag: "mod" },
                });
                return this.ctx.internItem({
                    ident,
                    kind: { tag: "mod", mod },
                    pub,
                    span,
                });
            }
            case "enum":
                return todo();
            case "struct":
                return todo();
            case "fn": {
                return this.ctx.internItem({
                    ident,
                    kind: { ...todo() },
                    pub,
                    span,
                });
            }
            case "use":
                return todo();
            case "type_alias":
                return todo();
        }
        exhausted(kind);
    }

    private lowerExpr(expr: ast.Expr): Expr {
        const span = expr.span;
        const kind = expr.kind;
        switch (kind.tag) {
            case "error":
                return this.ctx.internExpr(kind, span);
            case "path":
                return this.ctx.internExpr({
                    tag: "path",
                    path: this.lowerPath(kind.path, kind.qty),
                }, span);
            case "null":
                return this.ctx.internExpr(kind, span);
            case "int":
                return this.ctx.internExpr(kind, span);
            case "bool":
                return this.ctx.internExpr(kind, span);
            case "str":
                return this.ctx.internExpr(kind, span);
            case "group":
                return this.ctx.internExpr({
                    tag: "group",
                    expr: this.lowerExpr(kind.expr),
                }, span);
            case "array":
                return this.ctx.internExpr({
                    tag: "array",
                    exprs: kind.exprs.map((expr) => this.lowerExpr(expr)),
                }, span);
            case "repeat":
                return this.ctx.internExpr({
                    tag: "repeat",
                    expr: this.lowerExpr(kind.expr),
                    length: this.lowerExpr(kind.length),
                }, span);
            case "struct":
                return this.ctx.internExpr({
                    tag: "struct",
                    path: kind.path && this.lowerPath(kind.path),
                    fields: kind.fields.map(({ ident, expr, span }) => ({
                        ident,
                        expr: this.lowerExpr(expr),
                        span,
                    })),
                }, span);
            case "ref":
                return this.ctx.internExpr({
                    tag: "ref",
                    expr: this.lowerExpr(kind.expr),
                    refType: kind.refType,
                    mut: kind.mut,
                }, span);
            case "deref":
                return this.ctx.internExpr({
                    tag: "deref",
                    expr: this.lowerExpr(kind.expr),
                }, span);
            case "elem":
                return this.ctx.internExpr({
                    tag: "elem",
                    expr: this.lowerExpr(kind.expr),
                    elem: kind.elem,
                }, span);
            case "field":
                return this.ctx.internExpr({
                    tag: "field",
                    expr: this.lowerExpr(kind.expr),
                    ident: kind.ident,
                }, span);
            case "index":
                return this.ctx.internExpr({
                    tag: "index",
                    expr: this.lowerExpr(kind.expr),
                    index: this.lowerExpr(kind.index),
                }, span);
            case "call":
                return this.ctx.internExpr({
                    tag: "call",
                    expr: this.lowerExpr(kind.expr),
                    args: kind.args.map((arg) => this.lowerExpr(arg)),
                }, span);
            case "unary":
                return this.ctx.internExpr({
                    tag: "unary",
                    expr: this.lowerExpr(kind.expr),
                    unaryType: kind.unaryType,
                }, span);
            case "binary":
                return this.ctx.internExpr({
                    tag: "binary",
                    left: this.lowerExpr(kind.left),
                    right: this.lowerExpr(kind.right),
                    binaryType: kind.binaryType,
                }, span);
            case "block":
                return this.ctx.internExpr({
                    tag: "block",
                    block: this.lowerBlock(kind.block),
                }, span);
            case "if":
                return this.ctx.internExpr({
                    tag: "if",
                    cond: this.lowerExpr(kind.cond),
                    truthy: this.lowerBlock(kind.truthy),
                    falsy: kind.falsy && this.lowerExpr(kind.falsy),
                }, span);
            case "loop":
                return this.ctx.internExpr({
                    tag: "loop",
                    body: this.lowerBlock(kind.body),
                }, span);
            case "while":
                throw new Error("not implemented");
            case "for":
                throw new Error("not implemented");
            case "c_for":
                throw new Error("not implemented");
        }
        exhausted(kind);
    }

    private lowerPat(pat: ast.Pat): Pat {
        const span = pat.span;
        const kind = pat.kind;
        switch (kind.tag) {
            case "error":
                return this.ctx.internPat(kind, span);
            case "bind": {
                const v = this.ctx.internPat(kind, span);
                this.valRib().bindings.set(idKey(kind.ident.id), {
                    tag: "local",
                    id: v.id,
                });
                return v;
            }
            case "path":
                return this.ctx.internPat({
                    tag: "path",
                    path: this.lowerPath(kind.path, kind.qty),
                }, span);
        }
        exhausted(kind);
    }

    private lowerTy(ty: ast.Ty): Ty {
        const span = ty.span;
        const kind = ty.kind;
        switch (kind.tag) {
            case "error":
                return this.ctx.internTy(kind, span);
            case "null":
                return this.ctx.internTy(kind, span);
            case "int":
                return this.ctx.internTy(kind, span);
            case "bool":
                return this.ctx.internTy(kind, span);
            case "str":
                return this.ctx.internTy(kind, span);
            case "path":
                return this.ctx.internTy({
                    tag: "path",
                    path: this.lowerPath(kind.path, kind.qty),
                }, span);
            case "ref":
                return this.ctx.internTy({
                    tag: "ref",
                    ty: this.lowerTy(kind.ty),
                    mut: kind.mut,
                }, span);
            case "ptr":
                return this.ctx.internTy({
                    tag: "ptr",
                    ty: this.lowerTy(kind.ty),
                    mut: kind.mut,
                }, span);
            case "slice":
                return this.ctx.internTy({
                    tag: "slice",
                    ty: this.lowerTy(kind.ty),
                }, span);
            case "array":
                return this.ctx.internTy({
                    tag: "array",
                    ty: this.lowerTy(kind.ty),
                    length: this.lowerExpr(kind.length),
                }, span);
            case "anon_struct":
                return this.ctx.internTy({
                    tag: "anon_struct",
                    fields: kind.fields.map(({ ident, ty, span }) => ({
                        ident,
                        ty: this.lowerTy(ty),
                        span,
                    })),
                }, span);
        }
        exhausted(kind);
    }

    private lowerBlock(block: ast.Block): Block {
        const point = this.ribPoint();
        this.pushRib({ tag: "block" });
        const stmts = block.stmts.map((stmt) => this.lowerStmt(stmt));
        const expr = block.expr && this.lowerExpr(block.expr);
        this.returnToRibPoint(point);
        return this.ctx.internBlock({ stmts, expr, span: block.span });
    }

    private lowerPath(path: ast.Path, qty?: ast.Ty): QPath {
        if (qty) {
            const ty = this.lowerTy(qty);
            if (path.segments.length !== 1) {
                throw new Error();
            }
            const seg = path.segments[0];
            return {
                tag: "type_relative",
                ty,
                seg: {
                    ident: seg.ident,
                    res: todo(),
                    genericArgs: seg.genericArgs &&
                        seg.genericArgs.map((ty) => this.lowerTy(ty)),
                    inferArgs: false,
                    span: path.span,
                },
            };
        }
        const [res, segments] = this.resolvePathSegs(path.segments);
        return {
            tag: "resolved",
            path: { segments, res, span: path.span },
        };
    }

    private resolvePathSegs(segs: ast.PathSegment[]): [Res, PathSegment[]] {
        if (segs.length <= 1) {
            const seg = segs[0];
            const res = this.resolveTyIdent(seg.ident);
            return [res, [{
                ident: seg.ident,
                res,
                genericArgs: seg.genericArgs &&
                    seg.genericArgs.map((ty) => this.lowerTy(ty)),
                inferArgs: false,
                span: seg.span,
            }]];
        }
        const seg = segs.at(-1)!;
        const [innerRes, resSegs] = this.resolvePathSegs(
            segs.slice(0, segs.length - 1),
        );
        switch (innerRes.tag) {
            case "error":
                return [innerRes, [...resSegs, {
                    ident: seg.ident,
                    res: innerRes,
                    inferArgs: false,
                    span: seg.span,
                }]];
            case "def": {
                const error = (): [
                    Res,
                    PathSegment[],
                ] => [{ tag: "error" }, [...resSegs, {
                    ident: seg.ident,
                    res: innerRes,
                    inferArgs: false,
                    span: seg.span,
                }]];

                const irk = innerRes.kind;
                switch (irk.tag) {
                    case "mod": {
                        const mod = this.ctx.getMod(innerRes.id);
                        const res = mod.defs.get(seg.ident.id);
                        if (!res) {
                            this.ctx.report({
                                severity: "error",
                                file: this.currentFile,
                                msg: `module does not contain definition for '${
                                    this.ctx.identText(seg.ident.id)
                                }'`,
                                span: seg.span,
                            });
                            return error();
                        }
                        return [res, [...resSegs, {
                            ident: seg.ident,
                            res: innerRes,
                            inferArgs: false,
                            span: seg.span,
                        }]];
                    }
                    case "struct":
                        return todo();
                    case "enum":
                        return todo();
                    case "ty_alias":
                        return todo();
                    case "ty_param":
                        return todo();
                    case "use":
                        return todo();
                    case "fn":
                    case "variant":
                    case "ctor":
                    case "field":
                        this.ctx.report({
                            severity: "error",
                            file: this.currentFile,
                            msg: `${irk.tag} contains zero members`,
                            span: seg.span,
                        });
                        return error();
                }
                exhausted(irk);
                throw new Error();
            }
            case "local":
                throw new Error("should not be possible");
        }
        exhausted(innerRes);
    }

    private resolveTyIdent(ident: Ident): Res {
        return this.findTyRibIdent(this.tyRibs.length - 1, ident);
    }

    private findTyRibIdent(ribIdx: number, ident: Ident): Res {
        const rib = this.tyRibs[ribIdx];
        if (rib.bindings.has(idKey(ident.id))) {
            return rib.bindings.get(idKey(ident.id))!;
        }
        if (ribIdx === 0) {
            const text = this.ctx.identText(ident.id);
            this.ctx.report({
                severity: "error",
                file: this.currentFile,
                span: ident.span,
                msg: `no type with name '${text}' in module`,
            });
            return { tag: "error" };
        }
        const res = this.findTyRibIdent(ribIdx - 1, ident);
        const kind = rib.kind;
        switch (kind.tag) {
            case "normal":
                return res;
            case "fn":
                return res;
            case "item":
                if (res.tag === "local") {
                    const text = this.ctx.identText(ident.id);
                    this.ctx.report({
                        severity: "error",
                        file: this.currentFile,
                        span: ident.span,
                        msg: `cannot use local '${text}' here`,
                    });
                    return { tag: "error" };
                }
                return res;
            case "mod": {
                const text = this.ctx.identText(ident.id);
                this.ctx.report({
                    severity: "error",
                    file: this.currentFile,
                    span: ident.span,
                    msg: `no type with name '${text}' in module`,
                });
                return { tag: "error" };
            }
            case "block":
                return res;
        }
        exhausted(kind);
    }

    private resolveValIdent(ident: Ident): Res {
        return this.findValRibIdent(this.valRibs.length - 1, ident);
    }

    private findValRibIdent(ribIdx: number, ident: Ident): Res {
        const rib = this.valRibs[ribIdx];
        if (rib.bindings.has(idKey(ident.id))) {
            return rib.bindings.get(idKey(ident.id))!;
        }
        if (ribIdx === 0) {
            const text = this.ctx.identText(ident.id);
            this.ctx.report({
                severity: "error",
                file: this.currentFile,
                span: ident.span,
                msg: `no value with name '${text}' in module`,
            });
            return { tag: "error" };
        }
        const res = this.findValRibIdent(ribIdx - 1, ident);
        const kind = rib.kind;
        switch (kind.tag) {
            case "normal":
                return res;
            case "fn":
                return res;
            case "item":
                if (res.tag === "local") {
                    const text = this.ctx.identText(ident.id);
                    this.ctx.report({
                        severity: "error",
                        file: this.currentFile,
                        span: ident.span,
                        msg: `cannot use local '${text}' here`,
                    });
                    return { tag: "error" };
                }
                return res;
            case "mod": {
                const text = this.ctx.identText(ident.id);
                this.ctx.report({
                    severity: "error",
                    file: this.currentFile,
                    span: ident.span,
                    msg: `no value with name '${text}' in module`,
                });
                return { tag: "error" };
            }
            case "block":
                return res;
        }
        exhausted(kind);
    }
    private tyRib(): Rib {
        return this.tyRibs.at(-1)!;
    }

    private valRib(): Rib {
        return this.valRibs.at(-1)!;
    }

    private pushRib(kind: RibKind) {
        this.tyRibs.push({ kind, bindings: new Map() });
        this.valRibs.push({ kind, bindings: new Map() });
    }

    private popRib() {
        this.tyRibs.pop();
        this.valRibs.pop();
    }

    private ribPoint(): number {
        return this.valRibs.length;
    }

    private returnToRibPoint(point: number) {
        this.tyRibs = this.tyRibs.slice(0, point);
        this.valRibs = this.valRibs.slice(0, point);
    }
}