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

type LocationType = "register" | "temporary" | "variable"

type Location = {
    type: LocationType,
    name: string,
}

type Operand = Location | number

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

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

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

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

export type SSA = SSACopy | SSAUnary | SSABinary

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

const temp = (state: IRState): Location => ({
    type: "temporary",
    name: `t${state.nextID++}`,
})

const vari = (name: string): Location => ({
    type: "variable",
    name,
})

const reg = (name: string): Location => ({
    type: "register",
    name,
})

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, vari(stmt.args.name), stmt.args.expr)
}

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

    if (typeof expr === "number") {
        expr_stmts = [{ dest, source: expr }]
    } else if (typeof expr === "string") {
        expr_stmts = [{ dest, source: 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 [vari(expr), []]
    }

    const source = temp(state)
    const stmts = convertExpr(state, source, expr)
    return [source, stmts]
}

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

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

    return output
}

const ppOp = (op: Operand): string =>
    typeof op === "number"
        ? op.toString()
        : ppLoc(op)

const ppLoc = (loc: Location): string => {
    let output = ''

    switch (loc.type) {
        case "register":
            output = `%${loc.name}`
            break

        case "temporary":
            output = `#${loc.name}`
            break

        case "variable":
            output = loc.name
            break
    }

    return output
}