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

export type Line = {
    labels?: number[];
} & LineKind;

export type Label = { label: number };

export type LineKind =
    | { type: "op"; op: number }
    | { type: "lit"; val: number }
    | { type: "ref"; label: number };

export class Assembler {
    private lines: Line[] = [];
    private labelCounter = 0;

    public push(...values: number[]) {
        for (const value of values) {
            this.addLit(value);
        }
    }

    public addOp(op: number): Assembler {
        this.lines.push({ type: "op", op });
        return this;
    }

    public addLit(val: number): Assembler {
        this.lines.push({ type: "lit", val });
        return this;
    }

    public addRef({ label }: Label): Assembler {
        this.lines.push({ type: "ref", label });
        return this;
    }

    public setLabel({ label }: Label): Assembler {
        const line = this.lines.at(-1);
        if (line === undefined) {
            return this;
        }
        if (line.labels === undefined) {
            line.labels = [];
        }
        line.labels.push(label);
        return this;
    }

    public makeLabel(): Label {
        const label = this.labelCounter;
        this.labelCounter += 1;
        return { label };
    }

    public concat(assembler: Assembler) {
        this.lines.push(...assembler.lines);
    }

    public assemble(): number[] {
        let ip = 0;
        const output: number[] = [];
        const locs: { [key: number]: number } = {};
        const refs: { [key: number]: number } = {};
        for (const line of this.lines) {
            switch (line.type) {
                case "op":
                    output.push(line.op);
                    ip += 1;
                    break;
                case "lit":
                    output.push(line.val);
                    ip += 1;
                    break;
                case "ref":
                    output.push(0);
                    refs[ip] = line.label;
                    ip += 1;
                    break;
            }
        }
        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;
    }
}