import type { AssignStmt, Expr, Stmt } from "../ast"
import { R8 } from "../sm83/cpu"
import { Loc, LocType } from "./loc"
import type { Operand, SSA } from "./ssa"

type IRState = {
    nextID: number
    ssa_stmts: Array<SSA>
}

export const convertIR = (stmts: Array<Stmt>): Array<SSA> => {
    const state: IRState = {
        nextID: 0,
        ssa_stmts: [],
    }

    stmts.forEach(stmt => {
        const ssa_stmts = convertAssignStmt(state, stmt)
        state.ssa_stmts.push(...ssa_stmts)
    })

    return state.ssa_stmts
}

export const convertAssignStmt = (state: IRState, stmt: AssignStmt): Array<SSA> => {
    return convertExpr(state, Loc.vari(stmt.args.name), stmt.args.expr)
}

export const convertExpr = (state: IRState, dest: Loc, expr: Expr): Array<SSA> => {
    const threeAddress = convertExpr3(state, dest, expr)
    const output: Array<SSA> = []

    threeAddress.forEach(ssa => {
        // skip copies
        if (!("op" in ssa)) {
            output.push(ssa)
            return
        } else if ("source1" in ssa) {
            let source1 = ssa.source1
            if (typeof source1 !== 'number' && source1.type === LocType.VARIABLE) {
                source1 = Loc.temp(`t${state.nextID++}`)
                output.push(
                    { dest: Loc.reg(R8.A), source: ssa.source1 },
                    { dest: source1, source: Loc.reg(R8.A) },
                )
            }

            output.push(
                { dest: Loc.reg(R8.A), source: ssa.source },
                { dest: Loc.reg(R8.A), source: Loc.reg(R8.A), op: ssa.op, source1 },
                { dest: ssa.dest, source: Loc.reg(R8.A) },
            )
        } else {
            output.push(
                { dest: Loc.reg(R8.A), source: ssa.source },
                { dest: Loc.reg(R8.A), source: Loc.reg(R8.A), op: ssa.op },
                { dest: ssa.dest, source: Loc.reg(R8.A) },
            )
        }
    })

    return output
}

export const convertExpr3 = (state: IRState, dest: Loc, expr: Expr): Array<SSA> => {
    let expr_stmts: Array<SSA> = []

    if (typeof expr === "number") {
        expr_stmts = [{ dest, source: expr }]
    } else if (typeof expr === "string") {
        expr_stmts = [{ dest, source: Loc.vari(expr) }]
    } else if (expr.type === "unary") {
        const [source, stmts] = getSource(state, expr.arg)
        stmts.push({
            dest,
            op: expr.op,
            source,
        })

        expr_stmts = stmts
    } else {
        const [left_source, left_stmts] = getSource(state, expr.left)
        const [right_source, right_stmts] = getSource(state, expr.right)

        const stmts = [...left_stmts, ...right_stmts]
        stmts.push({
            dest,
            op: expr.op,
            source: left_source,
            source1: right_source,
        })

        expr_stmts = stmts
    }

    return expr_stmts
}

const getSource = (state: IRState, expr: Expr): [Operand, Array<SSA>] => {
    if (typeof expr === "number") {
        return [expr, []]
    } else if (typeof expr === "string") {
        return [Loc.vari(expr), []]
    }

    const source = Loc.temp(`t${state.nextID++}`)
    const stmts = convertExpr3(state, source, expr)
    return [source, stmts]
}