pub mod assemble_ir;

use std::fmt::{Display, Formatter};
use std::rc::Rc;

pub type ControlUnitId = usize;

pub type VirtualRegisterId = u32;

pub struct AsmFunction {
    fqn: Rc<str>,
    control_units: Vec<ControlUnit>,
}

impl AsmFunction {
    pub fn new(fqn: Rc<str>, control_units: Vec<ControlUnit>) -> Self {
        Self { fqn, control_units }
    }
}

impl Display for AsmFunction {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "fn {}", self.fqn)?;
        for control_unit in &self.control_units {
            writeln!(f, "    #{}:", control_unit.id)?;
            for instruction in &control_unit.instructions {
                writeln!(f, "    {}", instruction)?;
            }
        }
        Ok(())
    }
}

pub struct ControlUnit {
    id: ControlUnitId,
    successor_ids: Vec<ControlUnitId>,
    instructions: Vec<AsmInstruction>,
}

impl ControlUnit {
    pub fn new(id: ControlUnitId) -> Self {
        Self {
            id,
            successor_ids: vec![],
            instructions: vec![],
        }
    }

    pub fn append_instruction(&mut self, instruction: AsmInstruction) {
        self.instructions.push(instruction);
    }

    pub fn append_instructions(&mut self, instructions: Vec<AsmInstruction>) {
        self.instructions.extend(instructions);
    }
}

pub enum AsmInstruction {
    Goto {
        control_unit_id: ControlUnitId,
    },
    Move {
        source: AsmOperand,
        destination: VirtualRegisterId,
    },
    Push {
        source: AsmOperand,
    },
    Pop {
        destination: VirtualRegisterId,
    },
    Allocate {
        class_name: Rc<str>,
        destination: VirtualRegisterId,
    },
    InvokeStatic {
        fqn: Rc<str>,
    },
    InvokeObject {
        fqn: Rc<str>,
    },
    Return {
        operand: AsmOperand,
    },
    BinaryOperation {
        destination: VirtualRegisterId,
        left: AsmOperand,
        right: AsmOperand,
        operator: AsmBinaryOperator,
    },
}

impl Display for AsmInstruction {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            AsmInstruction::Goto { control_unit_id } => {
                write!(f, "goto #{}", control_unit_id)
            }
            AsmInstruction::Move {
                source,
                destination,
            } => {
                write!(f, "move vr{}, {}", destination, source)
            }
            AsmInstruction::Push { source } => {
                write!(f, "push {}", source)
            }
            AsmInstruction::Pop { destination } => {
                write!(f, "pop vr{}", destination)
            }
            AsmInstruction::Allocate {
                class_name,
                destination,
            } => {
                write!(f, "allocate r{}, {}", destination, class_name)
            }
            AsmInstruction::InvokeStatic { fqn } => {
                write!(f, "invoke_static {}", fqn)
            }
            AsmInstruction::InvokeObject { fqn } => {
                write!(f, "invoke_object {}", fqn)
            }
            AsmInstruction::Return { operand } => {
                write!(f, "return {}", operand)
            }
            AsmInstruction::BinaryOperation {
                destination,
                left,
                right,
                operator,
            } => {
                write!(f, "op{} vr{}, {}, {}", operator, destination, left, right)
            }
        }
    }
}

pub enum AsmOperand {
    I8(i8),
    I16(i16),
    I32(i32),
    I64(i64),
    I128(i128),
    ISize(isize),
    U8(u8),
    U16(u16),
    U32(u32),
    U64(u64),
    U128(u128),
    USize(usize),
    Float(f32),
    Double(f64),
    Boolean(bool),
    Constant { name: Rc<str> },
    VirtualRegister(VirtualRegisterId),
}

impl Display for AsmOperand {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            AsmOperand::I8(i) => {
                write!(f, "{}i8", i)
            }
            AsmOperand::I16(i) => {
                write!(f, "{}i16", i)
            }
            AsmOperand::I32(i) => {
                write!(f, "{}i32", i)
            }
            AsmOperand::I64(i) => {
                write!(f, "{}i64", i)
            }
            AsmOperand::I128(i) => {
                write!(f, "{}i128", i)
            }
            AsmOperand::ISize(i) => {
                write!(f, "{}isize", i)
            }
            AsmOperand::U8(u) => {
                write!(f, "{}u8", u)
            }
            AsmOperand::U16(u) => {
                write!(f, "{}u16", u)
            }
            AsmOperand::U32(u) => {
                write!(f, "{}u32", u)
            }
            AsmOperand::U64(u) => {
                write!(f, "{}u64", u)
            }
            AsmOperand::U128(u) => {
                write!(f, "{}u128", u)
            }
            AsmOperand::USize(u) => {
                write!(f, "{}usize", u)
            }
            AsmOperand::Float(float) => {
                write!(f, "{}f32", float)
            }
            AsmOperand::Double(d) => {
                write!(f, "{}f64", d)
            }
            AsmOperand::Boolean(b) => {
                write!(f, "{}", b)
            }
            AsmOperand::Constant { name } => {
                write!(f, "<{}>", name)
            }
            AsmOperand::VirtualRegister(id) => {
                write!(f, "vr{}", id)
            }
        }
    }
}

pub enum AsmBinaryOperator {
    Add,
    Subtract,
    Multiply,
    Divide,
}

impl Display for AsmBinaryOperator {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            AsmBinaryOperator::Add => {
                write!(f, "+")
            }
            AsmBinaryOperator::Subtract => {
                write!(f, "-")
            }
            AsmBinaryOperator::Multiply => {
                write!(f, "*")
            }
            AsmBinaryOperator::Divide => {
                write!(f, "/")
            }
        }
    }
}