import { Ops, 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 class Assembler {
    private lines: Line[] = [];
    private addedLabels: string[] = [];
    private labelCounter = 0;

    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 concat(assembler: Assembler) {
        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 makeLabel(): Label {
        return { label: `.L${(this.labelCounter++).toString()}` };
    }

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

    public assemble(): number[] {
        console.log("Assembling...");
        let ip = 0;
        const output: number[] = [];
        const locs: { [key: string]: number } = {};
        const refs: { [key: number]: string } = {};

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

    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
            ).reduce((acc, curr) => acc + curr, 0);
        }
    }
}