use crate::constants_table::ConstantsTable;
use crate::ir::assemble::{Assemble, InstructionsBuilder};
use crate::ir::ir_operation::IrOperation;
use crate::ir::ir_variable::{IrVariable, IrVariableDescriptor, IrVrVariableDescriptor};
use crate::ir::register_allocation::{OffsetCounter, VrUser};
use dvm_lib::instruction::Instruction;
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};
use std::rc::Rc;

pub struct IrAssign {
    destination: Rc<RefCell<IrVariable>>,
    initializer: Box<IrOperation>,
}

impl IrAssign {
    pub fn new(destination: Rc<RefCell<IrVariable>>, initializer: IrOperation) -> Self {
        Self {
            destination,
            initializer: initializer.into(),
        }
    }
}

impl VrUser for IrAssign {
    fn vr_definitions(&self) -> HashSet<IrVrVariableDescriptor> {
        match self.destination.borrow().descriptor() {
            IrVariableDescriptor::VirtualRegister(vr_descriptor) => {
                HashSet::from([vr_descriptor.clone()])
            }
            IrVariableDescriptor::Stack(_) => HashSet::new(),
        }
    }

    fn vr_uses(&self) -> HashSet<IrVrVariableDescriptor> {
        self.initializer.vr_uses()
    }

    fn propagate_spills(&mut self, spills: &HashSet<IrVrVariableDescriptor>) {
        let borrowed_destination = self.destination.borrow();
        if let IrVariableDescriptor::VirtualRegister(vr_variable) =
            borrowed_destination.descriptor()
        {
            if spills.contains(vr_variable) {
                let replacement = IrVariable::new_stack(
                    vr_variable.name().into(),
                    vr_variable.block_id(),
                    borrowed_destination.type_info().clone(),
                );
                drop(borrowed_destination);
                self.destination.replace(replacement);
            }
        }
    }

    fn propagate_register_assignments(
        &mut self,
        assignments: &HashMap<IrVrVariableDescriptor, usize>,
    ) {
        // Have to isolate the following, because when the destination and the value are the same,
        // we get a BorrowMutError
        {
            let mut borrowed_destination = self.destination.borrow_mut();
            if let IrVariableDescriptor::VirtualRegister(vr_variable) =
                borrowed_destination.descriptor_mut()
            {
                if assignments.contains_key(vr_variable) {
                    vr_variable.set_assigned_register(assignments[vr_variable]);
                }
            }
        }
        self.initializer.propagate_register_assignments(assignments);
    }

    fn propagate_stack_offsets(&mut self, counter: &mut OffsetCounter) {
        let mut borrowed_destination = self.destination.borrow_mut();
        if let IrVariableDescriptor::Stack(stack_variable) = borrowed_destination.descriptor_mut() {
            stack_variable.set_offset(counter.next());
        }
    }
}

impl Assemble for IrAssign {
    fn assemble(&self, builder: &mut InstructionsBuilder, constants_table: &mut ConstantsTable) {
        let destination = self.destination.borrow().descriptor().as_location();

        match self.initializer.as_ref() {
            IrOperation::GetFieldRef(ir_get_field_ref) => {
                let self_location = ir_get_field_ref.self_parameter_or_variable().as_location();
                builder.push(Instruction::GetFieldPointer(
                    self_location,
                    ir_get_field_ref.field_index(),
                    destination,
                ));
            }
            IrOperation::GetFieldRefMut(ir_get_field_ref_mut) => {
                let self_location = ir_get_field_ref_mut
                    .self_parameter_or_variable()
                    .as_location();
                builder.push(Instruction::GetFieldPointerMut(
                    self_location,
                    ir_get_field_ref_mut.field_index(),
                    destination,
                ));
            }
            IrOperation::ReadField(ir_read_field) => {
                let field_ref_location = ir_read_field
                    .field_ref_variable()
                    .borrow()
                    .descriptor()
                    .as_location();
                builder.push(Instruction::ReadField(field_ref_location, destination));
            }
            IrOperation::Load(ir_expression) => {
                let move_operand = ir_expression.move_operand(constants_table);
                builder.push(Instruction::Move(move_operand, destination));
            }
            IrOperation::Add(ir_add) => {
                let (left, right) = ir_add.operand_pair(constants_table);
                builder.push(Instruction::Add(left, right, destination));
            }
            IrOperation::Subtract(ir_subtract) => {
                let (left, right) = ir_subtract.operand_pair();
                builder.push(Instruction::Subtract(left, right, destination));
            }
            IrOperation::Multiply(ir_multiply) => {
                let (left, right) = ir_multiply.operand_pair();
                builder.push(Instruction::Multiply(left, right, destination));
            }
            IrOperation::Call(ir_call) => {
                ir_call.assemble(builder, constants_table);
                builder.push(Instruction::Pop(Some(destination)));
            }
            IrOperation::Allocate(ir_allocate) => builder.push(Instruction::Allocate(
                ir_allocate.class_fqn_owned(),
                destination,
            )),
        }
    }
}

impl Display for IrAssign {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}: {} = {}",
            self.destination.borrow(),
            self.destination.borrow().type_info(),
            self.initializer
        )
    }
}