import { Ops, opToString } 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;
    }
}

export type Line2 = { labels?: number[]; ins: Ins2 };

export type Ins2 = [Ops, ...Lit[]];

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

export class Assembler2 {
    private lines: Line2[] = [];
    private addedLabels: number[] = [];
    private labelCounter = 0;

    public add(ins: Ins2): Assembler2 {
        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: this.labelCounter++ };
    }

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

    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) {
            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") {
                    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() {
        for (const line of this.lines) {
            for (const label of line.labels ?? []) {
                console.log(`.L${label}:`);
            }
            const op = opToString(line.ins[0] as unknown 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 `.L${lit.label}`;
                }
            }).join(", ");
            console.log(`    ${op} ${args}`);
        }
    }
}