import type { AssignStmt, Expr, Stmt } from "../ast"
import { R8 } from "../sm83/cpu"
import { Loc, LocType } from "./loc"
import { Operand, AbsInsn2, AbsInsn3, CopyInsn, GotoInsn, UnaryInsn, BinaryInsn, LabelInsn } from "./insn"

export class IRState {
    nextID: number
    insns: Array<AbsInsn2>

    constructor() {
        this.nextID = 0
        this.insns = []
    }

    temp(): Loc {
        return Loc.temp(`t${this.nextID++}`)
    }
}

export const convertIR = (stmts: Array<Stmt>): Array<AbsInsn2> => {
    const state: IRState = new IRState()

    stmts.forEach(stmt => {
        switch (stmt.stmt_type) {
            case 'assign':
                convertExpr(state, Loc.vari(stmt.args.name), stmt.args.expr)
                break
            case 'label':
                state.insns.push(LabelInsn(stmt.args.name))
                break
            case 'goto':
                state.insns.push(GotoInsn(stmt.args.name))
                break
        }
    })

    return state.insns
}

export const convertExpr = (state: IRState, dest: Loc, expr: Expr): void => {
    const threeAddress = convertExpr3(state, dest, expr)

    threeAddress.forEach(ssa => {
        switch (ssa.type) {
            case 'copy':
                state.insns.push(ssa)
                break

            case 'unary':
                state.insns.push(
                    CopyInsn(Loc.reg(R8.A), ssa.source),
                    UnaryInsn(Loc.reg(R8.A), ssa.op, Loc.reg(R8.A)),
                    CopyInsn(ssa.dest, Loc.reg(R8.A)),
                )
                break

            case 'binary':
                let source1 = ssa.source1
                if (typeof source1 !== 'number' && source1.type === LocType.VARIABLE) {
                    source1 = state.temp()
                    state.insns.push(
                        CopyInsn(Loc.reg(R8.A), ssa.source1),
                        CopyInsn(source1, Loc.reg(R8.A)),
                    )
                }

                state.insns.push(
                    CopyInsn(Loc.reg(R8.A), ssa.source),
                    UnaryInsn(Loc.reg(R8.A), ssa.op, source1),
                    CopyInsn(ssa.dest, Loc.reg(R8.A)),
                )
                break
        }
    })
}

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

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

    if (typeof expr === "number") {
        expr_stmts.push(CopyInsn(dest, expr))
    } else if (typeof expr === "string") {
        expr_stmts.push(CopyInsn(dest, Loc.vari(expr)))
    } else if (expr.type === "unary") {
        const [source, stmts] = getSource(expr.arg)
        stmts.push(UnaryInsn(dest, expr.op, source))
        expr_stmts = stmts
    } else {
        const [left_source, left_stmts] = getSource(expr.left)
        const [right_source, right_stmts] = getSource(expr.right)
        expr_stmts = [...left_stmts, ...right_stmts, BinaryInsn(dest, left_source, expr.op, right_source)]
    }

    return expr_stmts
}