import type { SSA, SSABinary, SSACopy, SSAUnary } 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))

    for (let sym of Object.keys(state.syms)) {
        if (sym.startsWith("temp")) {
            continue
        }

        state.insns.push(
            `LD A, ${state.syms[sym]}`,
            `LD (${sym}), A`
        )
    }

    for (let sym of Object.keys(state.stack_offsets)) {
        if (sym.startsWith("temp")) {
            continue
        }

        state.insns.push(
            `LD HL, SP + (${state.cur_offset - state.stack_offsets[sym]})`,
            `LD A, (HL)`,
            `LD (${sym}), A`
        )
    }

    if (state.cur_offset !== 0) {
        state.insns.push(`ADD SP, ${state.cur_offset}`)
    }

    return state.insns
}

// TODO: Track liveness
const convertASM_SSA = (state: ASMState, ssa: SSA): void => {
    if ("source1" in ssa) { // Binary
        convertASM_SSA_Binary(state, ssa)
    } else if ("op" in ssa) { // Unary
        convertASM_SSA_Unary(state, ssa)
    } else { // Copy
        convertASM_SSA_Copy(state, ssa)
    }
}

export const convertASM_SSA_Binary = (state: ASMState, ssa: SSABinary): void => {
    const dest = state.alloc_sym_loc(ssa.dest)
}

export const convertASM_SSA_Unary = (state: ASMState, ssa: SSAUnary): void => {
    const dest = state.alloc_sym_loc(ssa.dest)

    if (typeof ssa.source == "number") {
        state.insns.push(`LD A, ${ssa.source}`)
    } else {
        const loc = state.get_sym_loc(ssa.source)

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

    switch (ssa.op) {
        case '-':
            state.insns.push(
                `CPL`,
                `INC A`,
            )
            break

        default:
            throw new Error(`unsupported unary op \`'${ssa.op}`)
    }

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


export const convertASM_SSA_Copy = (state: ASMState, ssa: SSACopy): void => {
    const dest = state.alloc_sym_loc(ssa.dest)

    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}`
        )
    }
}