use crate::constants_table::ConstantsTable;
use crate::ir::assemble::{Assemble, InstructionsBuilder};
use crate::ir::ir_block::IrBlock;
use crate::ir::ir_parameter::IrParameter;
use crate::ir::ir_variable::IrVrVariableDescriptor;
use crate::ir::register_allocation::HasVrUsers;
use crate::symbol::Symbol;
use crate::symbol::function_symbol::FunctionSymbol;
use crate::type_info::TypeInfo;
use dvm_lib::vm::function::Function;
use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt::Display;
use std::rc::Rc;

pub struct IrFunction {
    function_symbol: Rc<RefCell<FunctionSymbol>>,
    parameters: Vec<Rc<IrParameter>>,
    return_type_info: TypeInfo,
    entry: Rc<RefCell<IrBlock>>,
}

impl IrFunction {
    pub fn new(
        function_symbol: Rc<RefCell<FunctionSymbol>>,
        parameters: &[Rc<IrParameter>],
        return_type_info: &TypeInfo,
        entry: Rc<RefCell<IrBlock>>,
    ) -> Self {
        Self {
            function_symbol,
            parameters: parameters.to_vec(),
            return_type_info: return_type_info.clone(),
            entry,
        }
    }

    pub fn assign_registers(
        &mut self,
        register_count: usize,
    ) -> (HashMap<IrVrVariableDescriptor, usize>, isize) {
        self.entry.borrow_mut().assign_registers(register_count)
    }

    pub fn assemble(&self, stack_size: isize, constants_table: &mut ConstantsTable) -> Function {
        let mut builder = InstructionsBuilder::new();
        self.entry.borrow().assemble(&mut builder, constants_table);
        let instructions = builder.take_instructions();
        Function::new(
            self.function_symbol.borrow().declared_name_owned(),
            self.parameters.len(),
            stack_size,
            instructions,
        )
    }
}

impl Display for IrFunction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "fn {}(", self.function_symbol.borrow().declared_name())?;
        for (i, parameter) in self.parameters.iter().enumerate() {
            write!(f, "{}: {}", parameter, parameter.type_info())?;
            if i < self.parameters.len() - 1 {
                write!(f, ", ")?;
            }
        }
        write!(f, ") -> {}\n{}", self.return_type_info, self.entry.borrow())?;
        Ok(())
    }
}