import * as ty from "./ty.ts";
import * as mir from "./mir.ts";
import * as codegenX8664 from "./codegen/x86_64/module.ts";

export type PrettyColors = {
    punctuation: string;
    operator: string;
    literal: string;
    string: string;
    keyword: string;
    varIdent: string;
    fnIdent: string;
    typeIdent: string;
};

export const noColors: PrettyColors = {
    punctuation: "",
    operator: "",
    literal: "",
    string: "",
    keyword: "",
    varIdent: "",
    fnIdent: "",
    typeIdent: "",
};

export const consoleColors: PrettyColors = {
    punctuation: "\x1bpunctuation",
    operator: "\x1boperator",
    literal: "\x1bliteral",
    string: "\x1bstring",
    keyword: "\x1bkeyword",
    varIdent: "\x1bvarIdent",
    fnIdent: "\x1bfnIdent",
    typeIdent: "\x1btypeIdent",
};

export function printWithConsoleColors(str: string) {
    const keys = Object.keys(consoleColors).join("|");
    const pat = new RegExp(`\x1b(${keys})`, "g");
    const matches = str.matchAll(pat).toArray();

    const replacers: Record<string, string> = {
        punctuation: "color: gray;",
        operator: "color: #fb4934;",
        literal: "color: pink;",
        string: "color: green;",
        keyword: "color: blue;",
        varIdent: "color: lightblue;",
        fnIdent: "color: green; font-weight: bold;",
        typeIdent: "color: yellow;",
    };

    const colorArgs: string[] = [];

    for (const match of matches) {
        str = str.replace(match[0], "%c");
        colorArgs.push(replacers[match[1]]);
    }
    str += "%c";
    colorArgs.push("");

    console.log(str, ...colorArgs);
}

export function tyPretty(ty: ty.Ty, colors = noColors): string {
    const c = colors;
    switch (ty.kind.tag) {
        case "Error":
            return `<error>`;
        case "Void":
            return `${c.typeIdent}void`;
        case "Int":
            return `${c.typeIdent}${ty.kind.intTy}`;
        case "Bool":
            return `${c.typeIdent}bool`;
        case "Ptr":
            return `${c.punctuation}*${ty.kind.ty.pretty(c)}`;
        case "PtrMut":
            return `${c.punctuation}*${c.keyword}mut ${ty.kind.ty.pretty(c)}`;
        case "Array":
            return `${c.punctuation}[${
                ty.kind.ty.pretty(c)
            }${c.punctuation}; ${c.literal}${ty.kind.length}${c.punctuation}]`;
        case "Slice":
            return `${c.punctuation}[${ty.kind.ty.pretty(c)}${c.punctuation}]`;
        case "Range":
            return `${c.typeIdent}Range`;
        case "Fn":
            return `${c.keyword}fn ${c.punctuation}(${
                ty.kind.params.map((param) => param.pretty(c)).join(
                    `${c.punctuation}, `,
                )
            }${c.punctuation}) -> ${ty.kind.retTy.pretty(c)}`;
        case "FnStmt":
            if (!ty.kind.ty.is("Fn")) {
                throw new Error();
            }
            return `${c.keyword}fn ${ty.kind.stmt.kind.ident}(${
                ty.kind.ty.kind.params.map((param) => param.pretty(c))
                    .join(
                        `${c.punctuation}, `,
                    )
            }${c.punctuation}) -> ${ty.kind.ty.kind.retTy.pretty(c)}`;
    }
    return `{${ty.kind.tag}}`;
}

export function mirFnPretty(fn: mir.Fn, colors = noColors): string {
    return new MirFnPrettyStringifier(fn, colors).stringify();
}

class MirFnPrettyStringifier {
    private bbIds = new Map<mir.BasicBlock, number>();
    private instIds = new Map<mir.Inst, number>();
    private result = "";

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

    stringify(): string {
        const c = this.colors;

        const fnTy = this.fn.ty.is("FnStmt") && this.fn.ty.kind.ty.is("Fn")
            ? this.fn.ty.kind.ty
            : null;
        if (!fnTy) {
            throw new Error();
        }

        const ident = this.fn.stmt.kind.ident;

        const params = fnTy.kind.params
            .map((ty, idx) =>
                `${c.varIdent}${idx}${c.punctuation}: ${ty.pretty(c)}`
            )
            .join(`${c.punctuation}, `);
        const retTy = fnTy.kind.retTy.pretty(c);

        this.result +=
            `${c.keyword}fn ${c.fnIdent}${ident}${c.punctuation}(${params}${c.punctuation}) -> ${retTy}\n${c.punctuation}{\n`;

        for (const bb of this.fn.bbs) {
            this.basicBlock(bb);
        }

        this.result += `${c.punctuation}}`;

        return this.result;
    }

    private basicBlock(bb: mir.BasicBlock) {
        const c = this.colors;

        this.result += `${c.varIdent}bb${this.bbId(bb)}${c.punctuation}:\n`;

        for (const inst of bb.insts) {
            this.inst(inst);
        }
    }

    private instsPrinted = new Set<mir.Inst>();

    private inst(inst: mir.Inst) {
        if (this.instsPrinted.has(inst)) {
            return;
        }

        const c = this.colors;

        // inst is a value
        const expr = (tail: string) => {
            this.result += `    ${c.varIdent}%${
                this.instId(inst)
            }${c.punctuation}: ${
                inst.ty.pretty(c)
            }${c.punctuation} = ${c.operator}${inst.kind.tag}${tail}\n`;
        };
        // inst is a statement
        const stmt = (tail: string) => {
            this.result += `    ${c.operator}${inst.kind.tag}${tail}\n`;
        };

        this.instTail(inst, expr, stmt);

        this.instsPrinted.add(inst);
    }

    private instArg(inst: mir.Inst): string {
        const c = this.colors;

        const inline: mir.InstKind["tag"][] = ["Void", "Int", "Bool", "Array"];
        if (
            inline.includes(inst.kind.tag) &&
            !this.instsPrinted.has(inst)
        ) {
            let str = `${c.operator}${inst.kind.tag}`;
            this.instTail(inst, (s) => {
                str += s;
            }, () => {});
            return str;
        }
        this.inst(inst);
        return `${c.varIdent}%${this.instId(inst)}`;
    }

    private instTail(
        inst: mir.Inst,
        expr: (str: string) => void,
        stmt: (str: string) => void,
    ) {
        const c = this.colors;
        const r = (inst: mir.Inst) => this.instArg(inst);

        const k = inst.kind;
        switch (k.tag) {
            case "Error":
                return expr(``);
            case "Void":
                return expr(``);
            case "Int":
                return expr(
                    ` ${c.literal}${k.value}${
                        k.intTy !== "any" ? k.intTy : ""
                    }`,
                );
            case "Bool":
                return expr(` ${c.literal}${k.value}`);
            case "Str":
                return expr(` ${c.string}${JSON.stringify(k.value)}`);
            case "Array":
                return expr(
                    ` ${c.punctuation}[${
                        k.values.map(r).join(`${c.punctuation}, `)
                    }${c.punctuation}]`,
                );
            case "Fn":
                return expr(` ${c.fnIdent}${k.fn.stmt.kind.ident}`);
            case "Param":
                return expr(` ${c.varIdent}${k.idx}`);
            case "GetElemPtr":
                return expr(
                    ` ${c.punctuation}&(${c.keyword}ptr ${
                        r(k.base)
                    }${c.punctuation})[${r(k.offset)}${c.punctuation}]`,
                );
            case "Slice":
                return expr(
                    ` ${c.punctuation}&(${c.keyword}ptr ${
                        r(k.value)
                    }${c.punctuation})[${
                        k.begin ? r(k.begin) : ""
                    }${c.punctuation}..${
                        k.end ? r(k.end) : ""
                    }${c.punctuation}]`,
                );
            case "Call":
                return expr(
                    ` ${r(k.callee)} ${c.punctuation}(${
                        k.args.map(r).join(`${c.punctuation}, `)
                    }${c.punctuation})`,
                );
            case "Alloca":
                return expr(``);
            case "Load":
                return expr(
                    ` ${c.punctuation}[${c.keyword}ptr ${
                        r(k.source)
                    }${c.punctuation}]`,
                );
            case "Store":
                return stmt(
                    ` ${c.punctuation}[${c.keyword}ptr ${
                        r(k.target)
                    }${c.punctuation}] = ${r(k.source)}`,
                );
            case "Jump":
                return stmt(` ${c.varIdent}bb${this.bbId(k.target)}`);
            case "Branch":
                return stmt(
                    ` ${c.keyword}if ${
                        r(k.cond)
                    }${c.punctuation}: ${c.varIdent}bb${
                        this.bbId(k.truthy)
                    }${c.punctuation}, ${c.keyword}else${c.punctuation}: ${c.varIdent}bb${
                        this.bbId(k.falsy)
                    }`,
                );
            case "Return":
                return stmt(` ${r(k.source)}`);
            case "Not":
            case "Negate":
                return expr(`${r(k.source)}`);
            case "Eq":
            case "Ne":
            case "Lt":
            case "Gt":
            case "Lte":
            case "Gte":
            case "BitOr":
            case "BitXor":
            case "BitAnd":
            case "Shl":
            case "Shr":
            case "Add":
            case "Sub":
            case "Mul":
            case "Div":
            case "Rem":
                return expr(` ${r(k.left)}${c.punctuation}, ${r(k.right)}`);
            case "Len":
                return expr(` ${r(k.source)}`);
            case "DebugPrint":
                return stmt(` ${k.args.map(r).join(`${c.punctuation}, `)}`);
        }
    }

    private bbId(bb: mir.BasicBlock) {
        const found = this.bbIds.get(bb);
        if (found !== undefined) {
            return found;
        }
        const id = this.bbIds.size;
        this.bbIds.set(bb, id);
        return id;
    }

    private instId(inst: mir.Inst) {
        const found = this.instIds.get(inst);
        if (found !== undefined) {
            return found;
        }
        const id = this.instIds.size;
        this.instIds.set(inst, id);
        return id;
    }
}

export function x86_64FnPretty(fn: codegenX8664.Fn, colors = noColors): string {
    const c = colors;
    const blockIds = new Map<codegenX8664.Block, number>();

    const reg = (reg: codegenX8664.Reg): string => {
        const cl = c.typeIdent;
        switch (reg.id) {
            case codegenX8664.Regs.a:
                return `${cl}rax`;
            case codegenX8664.Regs.b:
                return `${cl}rbx`;
            case codegenX8664.Regs.c:
                return `${cl}rbc`;
            case codegenX8664.Regs.si:
                return `${cl}rsi`;
            case codegenX8664.Regs.di:
                return `${cl}rdi`;
            case codegenX8664.Regs.sp:
                return `${cl}rsp`;
            case codegenX8664.Regs.bp:
                return `${cl}rbp`;
            case codegenX8664.Regs.r8:
                return `${cl}r8`;
            case codegenX8664.Regs.r9:
                return `${cl}r9`;
            case codegenX8664.Regs.r10:
                return `${cl}r10`;
            case codegenX8664.Regs.r11:
                return `${cl}r11`;
            case codegenX8664.Regs.r12:
                return `${cl}r12`;
            case codegenX8664.Regs.r13:
                return `${cl}r13`;
            case codegenX8664.Regs.r14:
                return `${cl}r14`;
            case codegenX8664.Regs.r15:
                return `${cl}r15`;
            default:
                return `${c.varIdent}%${reg.id - 1000}`;
        }
    };

    const imm = (v: number): string => {
        return `${c.literal}${v}`;
    };

    const mem = (r: codegenX8664.Reg, off: number): string => {
        return `${c.keyword}qword ${c.punctuation}[${reg(r)}${c.punctuation}${
            off !== 0 ? `-${c.literal}${-off}${c.punctuation}` : ""
        }]`;
    };

    const line = (op: string, ...ops: string[]): string => {
        return `    ${c.operator}${op} ${ops.join(", ")}\n`;
    };

    for (const block of fn.blocks) {
        const id = blockIds.size;
        blockIds.set(block, id);
    }

    let result = "";
    result += `${c.fnIdent}${fn.mirFn.stmt.kind.ident}${c.punctuation}:\n`;
    for (const block of fn.blocks) {
        const id = blockIds.get(block)!;
        result += `${c.varIdent}.l${id}${c.punctuation}:\n`;
        for (const inst of block.insts) {
            const k = inst.kind;
            switch (k.tag) {
                case "PushR":
                    result += line("push", reg(k.src));
                    break;
                case "PopR":
                    result += line("pop", reg(k.dst));
                    break;
                case "MovRR":
                    result += line("mov", reg(k.dst), reg(k.src));
                    break;
                case "SubRI":
                    result += line("sub", reg(k.dst), imm(k.imm));
                    break;
                case "MovRI":
                    result += line("mov", reg(k.dst), imm(k.imm));
                    break;
                case "MovMRR":
                    result += line("mov", mem(k.dst, k.offset), reg(k.src));
                    break;
                case "MovRMR":
                    result += line("mov", reg(k.dst), mem(k.src, k.offset));
                    break;
                case "AddR":
                    result += line("add", reg(k.dst_op1), reg(k.op2));
                    break;
                case "Jmp":
                    result += line(
                        "jmp",
                        `${c.varIdent}.l${blockIds.get(k.bb)!}`,
                    );
                    break;
                case "Ret":
                    result += line("ret");
                    break;
            }
        }
    }
    return result;
}
