import { Reporter } from "../info.ts";
import { Pos } from "../token.ts";
import { createCfg } from "./cfg.ts";
import { Cfg } from "./cfg.ts";
import { Block, BlockId, Fn, Local, LocalId, Mir, RValue } from "./mir.ts";

export function checkBorrows(
    mir: Mir,
    reporter: Reporter,
) {
    for (const fn of mir.fns) {
        new BorrowCheckerFnPass(fn, reporter).pass();
    }
}

class BorrowCheckerFnPass {
    private cfg: Cfg;

    public constructor(
        private fn: Fn,
        private reporter: Reporter,
    ) {
        this.cfg = createCfg(this.fn);
    }

    public pass() {
        for (const local of this.fn.locals) {
            new LocalChecker(local, this.fn, this.cfg, this.reporter).check();
        }
    }
}

class LocalChecker {
    private visitedBlocks = new Set<BlockId>();

    private assignedTo = false;
    private moved = false;
    private borrowed = false;
    private borrowedMut = false;

    private movedPos?: Pos;
    private borrowedPos?: Pos;

    public constructor(
        private local: Local,
        private fn: Fn,
        private cfg: Cfg,
        private reporter: Reporter,
    ) {}

    public check() {
        this.checkBlock(this.cfg.entry());
    }

    private checkBlock(block: Block) {
        if (this.visitedBlocks.has(block.id)) {
            return;
        }
        this.visitedBlocks.add(block.id);
        for (const op of block.ops) {
            const ok = op.kind;
            switch (ok.type) {
                case "error":
                    break;
                case "assign":
                    this.markDst(ok.dst);
                    this.markSrc(ok.src);
                    break;
                case "ref":
                case "ptr":
                    this.markDst(ok.dst);
                    this.markBorrow(ok.src);
                    break;
                case "ref_mut":
                case "ptr_mut":
                    this.markDst(ok.dst);
                    this.markBorrowMut(ok.src);
                    break;
                case "deref":
                    this.markDst(ok.dst);
                    this.markSrc(ok.src);
                    break;
                case "assign_deref":
                    this.markSrc(ok.subject);
                    this.markSrc(ok.src);
                    break;
                case "field":
                    this.markDst(ok.dst);
                    this.markSrc(ok.subject);
                    break;
                case "assign_field":
                    this.markSrc(ok.subject);
                    this.markSrc(ok.src);
                    break;
                case "index":
                    this.markDst(ok.dst);
                    this.markSrc(ok.subject);
                    this.markSrc(ok.index);
                    break;
                case "assign_index":
                    this.markSrc(ok.subject);
                    this.markSrc(ok.index);
                    this.markSrc(ok.src);
                    break;
                case "call_val":
                    this.markDst(ok.dst);
                    this.markSrc(ok.subject);
                    for (const arg of ok.args) {
                        this.markSrc(arg);
                    }
                    break;
                case "binary":
                    this.markDst(ok.dst);
                    this.markSrc(ok.left);
                    this.markSrc(ok.right);
                    break;
            }
        }
        const tk = block.ter.kind;
        switch (tk.type) {
            case "error":
                break;
            case "return":
                break;
            case "jump":
                break;
            case "if":
                this.markSrc(tk.cond);
                break;
        }
        for (const child of this.cfg.children(block)) {
            this.checkBlock(child);
        }
    }

    private markDst(localId: LocalId) {
        if (localId !== this.local.id) {
            return;
        }
        if (!this.assignedTo) {
            this.assignedTo = true;
            return;
        }
        if (!this.local.mut) {
            this.reportReassignToNonMut();
            return;
        }
    }

    private markBorrow(localId: LocalId) {
        if (localId !== this.local.id) {
            return;
        }

        if (!this.assignedTo) {
            this.assignedTo = true;
            return;
        }
    }

    private markBorrowMut(localId: LocalId) {
        if (localId !== this.local.id) {
            return;
        }

        if (!this.assignedTo) {
            this.assignedTo = true;
            return;
        }
    }

    private markSrc(src: RValue) {
        if (src.type === "local") {
            throw new Error("should be 'copy' or 'move'");
        }
        if (
            (src.type !== "copy" && src.type !== "move") ||
            src.id !== this.local.id
        ) {
            return;
        }
        if (src.type === "move") {
            if (this.moved) {
                this.reportUseMoved();
                return;
            }
            if (this.borrowed) {
                this.reportUseBorrowed();
                return;
            }
            this.moved = true;
        }
    }

    private reportReassignToNonMut() {
        const ident = this.local.sym!.ident;
        this.reporter.reportError({
            reporter: "borrow checker",
            msg: `cannot re-assign to '${ident}' as it was not declared mutable`,
            pos: this.local.sym!.pos!,
        });
        this.reporter.addNote({
            reporter: "borrow checker",
            msg: `declared here`,
            pos: this.local.sym!.pos!,
        });
    }

    private reportUseMoved() {
        const ident = this.local.sym!.ident;
        this.reporter.reportError({
            reporter: "borrow checker",
            msg: `cannot use '${ident}' as it has been moved`,
            pos: this.local.sym!.pos!,
        });
        if (this.movedPos) {
            this.reporter.addNote({
                reporter: "borrow checker",
                msg: `moved here`,
                pos: this.movedPos,
            });
        }
    }

    private reportUseBorrowed() {
        const ident = this.local.sym!.ident;
        this.reporter.reportError({
            reporter: "borrow checker",
            msg: `cannot use '${ident}' as it has been borrowed`,
            pos: this.local.sym!.pos!,
        });
        if (this.borrowedPos) {
            this.reporter.addNote({
                reporter: "borrow checker",
                msg: `borrowed here`,
                pos: this.movedPos,
            });
        }
    }
}