import type { AssignStmt, Expr, Stmt } from "./ast"

type Operand = string | number

type ASSA<Data> = { dest: string, source: Operand } & Data

export type SSACopy = ASSA<{ source: Operand }> // dest = source

export type SSAUnary = ASSA<{ op: string }> // dest = op(source)

export type SSABinary = ASSA<{ op: string, source1: Operand }> // dest = op(source, source1)

export type SSA = SSACopy | SSAUnary | SSABinary

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> => {
    const dest = stmt.args.name
    const [source, expr_stmts] = convertExpr(state, stmt.args.expr)

    expr_stmts.push({
        dest,
        source,
    })

    return expr_stmts
}

export const convertExpr = (state: IRState, expr: Expr): [string | number, Array<SSA>] => {
    if (typeof expr === "number" || typeof expr === "string") {
        return [expr, []]
    } else if (expr.type === "unary") {
        const [source, expr_stmts] = convertExpr(state, expr.arg)

        const name = `temp${state.nextID++}`
        expr_stmts.push({
            dest: name,
            source,
            op: expr.op,
        })

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

        const name = `temp${state.nextID++}`
        expr_stmts.push({
            dest: name,
            op: expr.op,
            source: left_source,
            source1: right_source,
        })

        return [name, expr_stmts]
    }
}

export const prettyPrintIR = (ssa: SSA): string => {
    let output = ""

    if ("source1" in ssa) {
        output = `${ssa.dest} = ${ssa.source} ${ssa.op} ${ssa.source1}`
    } else if ("op" in ssa) {
        output = `${ssa.dest} = ${ssa.op} ${ssa.source}`
    } else {
        output = `${ssa.dest} = ${ssa.source}`
    }

    return output
}