import * as mir from "../../mir.ts";
import { Block, Fn, Inst, Reg, Regs } from "./module.ts";

export function selectFnInstructions(fn: mir.Fn): Fn {
    return new InstructionSelector(fn).generateFn();
}

class InstructionSelector {
    private localsSize = 0;
    private locals = new Map<mir.Inst, number>();
    private exitBlock!: Block;

    constructor(
        private fn: mir.Fn,
    ) {}

    generateFn(): Fn {
        this.generateFrameLayout();

        const entry = new Block();
        const exit = new Block();

        entry.push(new Inst({ tag: "PushR", src: Reg.reg(Regs.bp) }));
        entry.push(
            new Inst({
                tag: "MovRR",
                dst: Reg.reg(Regs.bp),
                src: Reg.reg(Regs.sp),
            }),
        );
        entry.push(
            new Inst({
                tag: "SubRI",
                dst: Reg.reg(Regs.sp),
                imm: this.localsSize,
            }),
        );

        this.exitBlock = exit;

        const blocks: Block[] = [];
        for (const bb of this.fn.bbs) {
            blocks.push(this.generateBasicBlock(bb));
        }

        entry.push(new Inst({ tag: "Jmp", bb: blocks[0] ?? exit }));

        exit.push(
            new Inst({
                tag: "MovRR",
                dst: Reg.reg(Regs.sp),
                src: Reg.reg(Regs.bp),
            }),
        );
        exit.push(new Inst({ tag: "PopR", dst: Reg.reg(Regs.bp) }));
        exit.push(new Inst({ tag: "Ret" }));

        return new Fn(this.fn, [entry, ...blocks, exit]);
    }

    private generateFrameLayout() {
        let offset = 0;

        const align = (alignment: number) => {
            if (offset % alignment != 0) {
                offset -= alignment + offset % alignment;
            }
        };

        const pushLocal = (inst: mir.Inst, size: number, alignment: number) => {
            align(alignment);
            this.locals.set(inst, offset);
            offset -= size;
        };

        this.fn.visit({
            visitInst(inst) {
                if (inst.kind.tag === "Alloca") {
                    const ty = inst.ty.as("PtrMut").kind.ty;
                    switch (ty.kind.tag) {
                        case "Int": {
                            switch (ty.kind.intTy) {
                                case "i32":
                                    pushLocal(inst, 4, 4);
                                    break;
                                case "i64":
                                    pushLocal(inst, 8, 8);
                                    break;
                                default:
                                    throw new Error(
                                        `not implemented (${ty.kind.intTy})`,
                                    );
                            }
                            break;
                        }
                        default:
                            throw new Error(
                                `not implemented (${ty.kind.tag})`,
                            );
                    }
                }
            },
        });

        // c abi requires 16 bit stack alignment,
        // might aswell do it here.
        align(16);

        this.localsSize = -offset;
    }

    private generateBasicBlock(bb: mir.BasicBlock): Block {
        const regs = new Map<mir.Inst, Reg>();

        const block = new Block();
        for (const inst of bb.insts) {
            const save = (r: Reg): Reg => {
                regs.set(inst, r);
                return r;
            };

            const k = inst.kind;
            switch (k.tag) {
                case "Alloca":
                    save(Reg.temp());
                    break;
                case "Void":
                    block.push(
                        new Inst({
                            tag: "MovRI",
                            dst: save(Reg.temp()),
                            imm: 0,
                        }),
                    );
                    break;
                case "Int":
                    block.push(
                        new Inst({
                            tag: "MovRI",
                            dst: save(Reg.temp()),
                            imm: k.value,
                        }),
                    );
                    break;
                case "Store":
                    block.push(
                        new Inst({
                            tag: "MovMRR",
                            dst: Reg.reg(Regs.bp),
                            offset: this.locals.get(k.target)!,
                            src: regs.get(k.source)!,
                        }),
                    );
                    break;
                case "Load":
                    block.push(
                        new Inst({
                            tag: "MovRMR",
                            dst: save(Reg.temp()),
                            offset: this.locals.get(k.source)!,
                            src: Reg.reg(Regs.bp),
                        }),
                    );
                    break;
                case "Add":
                    block.push(
                        new Inst({
                            tag: "AddR",
                            dst_op1: save(regs.get(k.left)!),
                            op2: regs.get(k.right)!,
                        }),
                    );
                    break;

                case "Return":
                    block.push(
                        new Inst({
                            tag: "MovRR",
                            dst: Reg.reg(Regs.a),
                            src: regs.get(k.source)!,
                        }),
                    );
                    block.push(
                        new Inst({
                            tag: "Jmp",
                            bb: this.exitBlock,
                        }),
                    );
                    break;
                default:
                    console.error(`not implemented (${k.tag})`);
                    // throw new Error(
                    //     `not implemented (${k.tag})`,
                    // );
            }
        }
        return block;
    }
}