import type { SSA } from "./ir"

const reg8 = ["B", "C", "D", "E"] as const

type R8 = typeof reg8[number]

class ASMState {
    regs: { [r in R8]: string | null }
    syms: { [s: string]: R8 }
    stack_offsets: { [s: string]: number }
    cur_offset: number
    insns: Array<string>

    constructor() {
        this.regs = { "B": null, "C": null, "D": null, "E": null }
        this.syms = {}
        this.stack_offsets = {}
        this.cur_offset = 0
        this.insns = []
    }

    // TODO: Generalize for u16
    get_sym_loc(sym: string): R8 | number | null {
        // return register if already allocated
        if (typeof this.syms[sym] !== "undefined") {
            return this.syms[sym]
        }

        // return stack offset if already allocated
        if (typeof this.stack_offsets[sym] !== "undefined") {
            return this.cur_offset - this.stack_offsets[sym]
        }

        return null
    }

    alloc_sym_loc(sym: string): R8 | number {
        const existing_loc = this.get_sym_loc(sym)
        if (existing_loc !== null) {
            return existing_loc
        }

        // look for free register
        for (const r8 of reg8) {
            if (this.regs[r8] === null) {
                this.regs[r8] = sym
                this.syms[sym] = r8
                return r8
            }
        }

        // allocate space on stack
        this.cur_offset++
        this.stack_offsets[sym] = this.cur_offset

        this.insns.push("ADD SP, -1")
        return this.stack_offsets[sym]
    }
}

export const convertASM = (ir: Array<SSA>): Array<string> => {
    const state = new ASMState()

    ir.forEach(ssa => convertASM_SSA(state, ssa))
    return state.insns
}

const convertASM_SSA = (state: ASMState, ssa: SSA): void => {
    const dest = state.alloc_sym_loc(ssa.dest)

    if ("source1" in ssa) { // Binary
    } else if ("op" in ssa) { // Unary
    } else { // Copy
        let source = ""
        if (typeof ssa.source == "number") {
            source = ssa.source.toString()
        } else {
            const loc = state.get_sym_loc(ssa.source)

            if (loc === null) {
                state.insns.push(`LD A, (${ssa.source})`)
                source = "A"
            } else if (typeof loc === "string") {
                source = loc
            } else {
                state.insns.push(
                    `LD HL, SP + (${loc})`,
                    `LD A, (HL)`
                )
                source = "A"
            }
        }

        if (typeof dest === "string") {
            state.insns.push(`LD ${dest}, ${source}`)
        } else {
            state.insns.push(
                `LD HL, SP + (${dest})`,
                `LD (HL), ${source}`
            )
        }
    }
}