import { opToString } from "./arch.ts";

export type Line = { labels?: string[]; ins: Ins };

export type Ins = Lit[];

export type Label = { label: string };

export type Lit = number | string | boolean | Label;

export type Locs = { [key: string]: number };
export type Refs = { [key: number]: string };

export class Assembler {
    private lines: Line[] = [];
    private addedLabels: string[] = [];

    private constructor(private labelCounter: number) {}

    public static newRoot(): Assembler {
        return new Assembler(0);
    }

    public fork(): Assembler {
        return new Assembler(this.labelCounter);
    }

    public join(assembler: Assembler) {
        this.labelCounter = assembler.labelCounter;
        if (assembler.lines.length < 0) {
            return;
        }
        if (assembler.lines[0].labels !== undefined) {
            this.addedLabels.push(...assembler.lines[0].labels);
        }
        this.add(...assembler.lines[0].ins);
        this.lines.push(...assembler.lines.slice(1));
    }

    public add(...ins: Ins): Assembler {
        if (this.addedLabels.length > 0) {
            this.lines.push({ ins, labels: this.addedLabels });
            this.addedLabels = [];
            return this;
        }
        this.lines.push({ ins });
        return this;
    }

    public makeLabel(): Label {
        return { label: `.L${(this.labelCounter++).toString()}` };
    }

    public setLabel({ label }: Label) {
        this.addedLabels.push(label);
    }

    public assemble(): { program: number[]; locs: Locs } {
        let ip = 0;
        const program: number[] = [];
        const locs: Locs = {};
        const refs: Refs = {};

        let selectedLabel = "";
        for (const line of this.lines) {
            for (const label of line.labels ?? []) {
                const isAbsLabel = !label.startsWith(".");
                if (isAbsLabel) {
                    selectedLabel = label;
                    locs[label] = ip;
                } else {
                    locs[`${selectedLabel}${label}`] = ip;
                }
            }
            for (const lit of line.ins as Lit[]) {
                if (typeof lit === "number") {
                    program.push(lit);
                    ip += 1;
                } else if (typeof lit === "boolean") {
                    program.push(lit ? 1 : 0);
                    ip += 1;
                } else if (typeof lit === "string") {
                    program.push(lit.length);
                    ip += 1;
                    for (let i = 0; i < lit.length; ++i) {
                        program.push(lit.charCodeAt(i));
                        ip += 1;
                    }
                } else {
                    program.push(0);
                    refs[ip] = lit.label.startsWith(".")
                        ? `${selectedLabel}${lit.label}`
                        : refs[ip] = lit.label;
                    ip += 1;
                }
            }
        }
        for (let i = 0; i < program.length; ++i) {
            if (!(i in refs)) {
                continue;
            }
            if (!(refs[i] in locs)) {
                console.error(
                    `Assembler: label '${refs[i]}' used at ${i} not defined`,
                );
                continue;
            }
            program[i] = locs[refs[i]];
        }
        return { program, locs };
    }

    public printProgram() {
        let ip = 0;
        for (const line of this.lines) {
            for (const label of line.labels ?? []) {
                console.log(`        ${label}:`);
            }
            const op = opToString(line.ins[0] as number)
                .padEnd(13, " ");
            const args = (line.ins.slice(1) as Lit[]).map((lit) => {
                if (typeof lit === "number") {
                    return lit;
                } else if (typeof lit === "boolean") {
                    return lit.toString();
                } else if (typeof lit === "string") {
                    return '"' +
                        lit.replaceAll("\\", "\\\\").replaceAll("\0", "\\0")
                            .replaceAll("\n", "\\n").replaceAll("\t", "\\t")
                            .replaceAll("\r", "\\r") +
                        '"';
                } else {
                    return lit.label;
                }
            }).join(", ");
            console.log(`${ip.toString().padStart(8, " ")}:    ${op} ${args}`);
            ip += line.ins.map((lit) =>
                typeof lit === "string" ? lit.length + 1 : 1
            ).reduce((acc, curr) => acc + curr, 0);
        }
    }
}