use crate::constants_table::ConstantsTable;
use crate::ir::assemble::InstructionsBuilder;
use crate::ir::ir_parameter::IrParameter;
use crate::ir::ir_variable::{
    IrStackVariableDescriptor, IrVariable, IrVariableDescriptor, IrVrVariableDescriptor,
};
use crate::ir::register_allocation::{OffsetCounter, VrUser};
use crate::type_info::TypeInfo;
use dvm_lib::instruction::{
    AddOperand, Location, MoveOperand, MultiplyOperand, PushOperand, ReturnOperand,
    SetFieldOperand, SubtractOperand,
};
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};
use std::rc::Rc;

pub enum IrExpression {
    Field(Rc<RefCell<IrVariable>>, usize),
    Parameter(Rc<IrParameter>),
    Variable(Rc<RefCell<IrVariable>>),
    Int(i32),
    Double(f64),
    String(Rc<str>),
}

impl IrExpression {
    pub fn move_operand(&self, constants_table: &mut ConstantsTable) -> MoveOperand {
        match self {
            IrExpression::Field(self_variable, field_index) => {
                match self_variable.borrow().descriptor() {
                    IrVariableDescriptor::VirtualRegister(register_self_variable) => {
                        MoveOperand::Location(Location::RegisterField(
                            register_self_variable.assigned_register(),
                            *field_index,
                        ))
                    }
                    IrVariableDescriptor::Stack(stack_self_variable) => MoveOperand::Location(
                        Location::StackFrameOffsetField(stack_self_variable.offset(), *field_index),
                    ),
                }
            }
            IrExpression::Parameter(ir_parameter) => {
                MoveOperand::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
            }
            IrExpression::Variable(ir_variable) => match ir_variable.borrow().descriptor() {
                IrVariableDescriptor::VirtualRegister(register_variable) => {
                    MoveOperand::Location(Location::Register(register_variable.assigned_register()))
                }
                IrVariableDescriptor::Stack(stack_variable) => {
                    MoveOperand::Location(Location::StackFrameOffset(stack_variable.offset()))
                }
            },
            IrExpression::Int(i) => MoveOperand::Int(*i),
            IrExpression::Double(d) => MoveOperand::Double(*d),
            IrExpression::String(s) => {
                let constant_name = constants_table.get_or_insert(s);
                MoveOperand::String(constant_name)
            }
        }
    }

    pub fn push_operand(&self, constants_table: &mut ConstantsTable) -> PushOperand {
        match self {
            IrExpression::Field(self_variable, _) => {
                todo!()
            }
            IrExpression::Parameter(ir_parameter) => {
                PushOperand::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
            }
            IrExpression::Variable(ir_variable) => match ir_variable.borrow().descriptor() {
                IrVariableDescriptor::VirtualRegister(register_variable) => {
                    PushOperand::Location(Location::Register(register_variable.assigned_register()))
                }
                IrVariableDescriptor::Stack(stack_variable) => {
                    PushOperand::Location(Location::StackFrameOffset(stack_variable.offset()))
                }
            },
            IrExpression::Int(i) => PushOperand::Int(*i),
            IrExpression::Double(d) => PushOperand::Double(*d),
            IrExpression::String(s) => {
                let constant_name = constants_table.get_or_insert(s);
                PushOperand::String(constant_name)
            }
        }
    }

    pub fn add_operand(&self, constants_table: &mut ConstantsTable) -> AddOperand {
        match self {
            IrExpression::Field(self_variable, _) => {
                todo!()
            }
            IrExpression::Parameter(ir_parameter) => match ir_parameter.type_info() {
                TypeInfo::Integer | TypeInfo::Double | TypeInfo::String => {
                    AddOperand::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
                }
                _ => panic!(
                    "Attempt to add non- integer/double/string (found: {})",
                    ir_parameter.type_info()
                ),
            },
            IrExpression::Variable(ir_variable) => match ir_variable.borrow().type_info() {
                TypeInfo::Integer | TypeInfo::Double | TypeInfo::String => {
                    match ir_variable.borrow().descriptor() {
                        IrVariableDescriptor::VirtualRegister(register_variable) => {
                            AddOperand::Location(Location::Register(
                                register_variable.assigned_register(),
                            ))
                        }
                        IrVariableDescriptor::Stack(stack_variable) => AddOperand::Location(
                            Location::StackFrameOffset(stack_variable.offset()),
                        ),
                    }
                }
                _ => panic!(
                    "Attempt to add non-integer/non-string (found: {})",
                    ir_variable.borrow().type_info()
                ),
            },
            IrExpression::Int(i) => AddOperand::Int(*i),
            IrExpression::Double(d) => AddOperand::Double(*d),
            IrExpression::String(s) => {
                let constant_name = constants_table.get_or_insert(s);
                AddOperand::String(constant_name)
            }
        }
    }

    pub fn subtract_operand(&self) -> SubtractOperand {
        match self {
            IrExpression::Field(self_variable, _) => {
                todo!()
            }
            IrExpression::Parameter(ir_parameter) => match ir_parameter.type_info() {
                TypeInfo::Integer | TypeInfo::Double => SubtractOperand::Location(
                    Location::StackFrameOffset(ir_parameter.stack_offset()),
                ),
                _ => panic!(
                    "Attempt to subtract with non-integer type (found {})",
                    ir_parameter.type_info()
                ),
            },
            IrExpression::Variable(ir_variable) => match ir_variable.borrow().type_info() {
                TypeInfo::Integer | TypeInfo::Double => match ir_variable.borrow().descriptor() {
                    IrVariableDescriptor::VirtualRegister(vr_variable) => {
                        SubtractOperand::Location(Location::Register(
                            vr_variable.assigned_register(),
                        ))
                    }
                    IrVariableDescriptor::Stack(stack_variable) => SubtractOperand::Location(
                        Location::StackFrameOffset(stack_variable.offset()),
                    ),
                },
                _ => panic!(
                    "Attempt to subtract with non-integer type (found {})",
                    ir_variable.borrow().type_info()
                ),
            },
            IrExpression::Int(i) => SubtractOperand::Int(*i),
            IrExpression::Double(d) => SubtractOperand::Double(*d),
            IrExpression::String(_) => {
                panic!("Attempt to subtract with a string");
            }
        }
    }

    pub fn multiply_operand(&self) -> MultiplyOperand {
        match self {
            IrExpression::Field(self_variable, _) => {
                todo!()
            }
            IrExpression::Parameter(ir_parameter) => {
                MultiplyOperand::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
            }
            IrExpression::Variable(ir_variable) => match ir_variable.borrow().descriptor() {
                IrVariableDescriptor::VirtualRegister(vr_variable) => {
                    MultiplyOperand::Location(Location::Register(vr_variable.assigned_register()))
                }
                IrVariableDescriptor::Stack(stack_variable) => {
                    MultiplyOperand::Location(Location::StackFrameOffset(stack_variable.offset()))
                }
            },
            IrExpression::Int(i) => MultiplyOperand::Int(*i),
            IrExpression::Double(d) => MultiplyOperand::Double(*d),
            IrExpression::String(_) => {
                panic!("Attempt to multiply with a string");
            }
        }
    }

    pub fn return_operand(&self, constants_table: &mut ConstantsTable) -> ReturnOperand {
        match self {
            IrExpression::Field(self_variable, field_index) => {
                match self_variable.borrow().descriptor() {
                    IrVariableDescriptor::VirtualRegister(self_register_variable) => {
                        ReturnOperand::Location(Location::RegisterField(
                            self_register_variable.assigned_register(),
                            *field_index,
                        ))
                    }
                    IrVariableDescriptor::Stack(self_stack_variable) => ReturnOperand::Location(
                        Location::StackFrameOffsetField(self_stack_variable.offset(), *field_index),
                    ),
                }
            }
            IrExpression::Parameter(ir_parameter) => {
                ReturnOperand::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
            }
            IrExpression::Variable(ir_variable) => match ir_variable.borrow().descriptor() {
                IrVariableDescriptor::VirtualRegister(register_variable) => {
                    ReturnOperand::Location(Location::Register(
                        register_variable.assigned_register(),
                    ))
                }
                IrVariableDescriptor::Stack(stack_variable) => {
                    ReturnOperand::Location(Location::StackFrameOffset(stack_variable.offset()))
                }
            },
            IrExpression::Int(i) => ReturnOperand::Int(*i),
            IrExpression::Double(d) => ReturnOperand::Double(*d),
            IrExpression::String(s) => {
                let constant_name = constants_table.get_or_insert(s);
                ReturnOperand::String(constant_name)
            }
        }
    }

    pub fn set_field_operand(&self, constants_table: &mut ConstantsTable) -> SetFieldOperand {
        match self {
            IrExpression::Field(self_variable, _) => {
                todo!()
            }
            IrExpression::Parameter(ir_parameter) => {
                SetFieldOperand::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
            }
            IrExpression::Variable(ir_variable) => match ir_variable.borrow().descriptor() {
                IrVariableDescriptor::VirtualRegister(vr_variable) => {
                    SetFieldOperand::Location(Location::Register(vr_variable.assigned_register()))
                }
                IrVariableDescriptor::Stack(stack_variable) => {
                    SetFieldOperand::Location(Location::StackFrameOffset(stack_variable.offset()))
                }
            },
            IrExpression::Int(i) => SetFieldOperand::Int(*i),
            IrExpression::Double(d) => SetFieldOperand::Double(*d),
            IrExpression::String(s) => {
                let constant_name = constants_table.get_or_insert(s);
                SetFieldOperand::String(constant_name)
            }
        }
    }
}

impl Display for IrExpression {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            IrExpression::Field(self_variable, field_index) => {
                write!(f, "{}.{}", self_variable.borrow(), field_index)
            }
            IrExpression::Parameter(ir_parameter) => {
                write!(f, "{}", ir_parameter)
            }
            IrExpression::Variable(ir_variable) => {
                write!(f, "{}", ir_variable.borrow())
            }
            IrExpression::Int(i) => {
                write!(f, "{}", i)
            }
            IrExpression::Double(d) => {
                write!(f, "{}", d)
            }
            IrExpression::String(s) => {
                write!(f, "\"{}\"", s)
            }
        }
    }
}

impl VrUser for IrExpression {
    fn vr_definitions(&self) -> HashSet<IrVrVariableDescriptor> {
        // no defs for an expression
        HashSet::new()
    }

    fn vr_uses(&self) -> HashSet<IrVrVariableDescriptor> {
        match self {
            IrExpression::Field(self_variable, field_index) => {
                if let IrVariableDescriptor::VirtualRegister(vr_variable) =
                    self_variable.borrow().descriptor()
                {
                    HashSet::from([vr_variable.clone()])
                } else {
                    HashSet::new()
                }
            }
            IrExpression::Parameter(_) => HashSet::new(),
            IrExpression::Variable(ir_variable) => {
                if let IrVariableDescriptor::VirtualRegister(vr_variable) =
                    ir_variable.borrow().descriptor()
                {
                    HashSet::from([vr_variable.clone()])
                } else {
                    HashSet::new()
                }
            }
            IrExpression::Int(_) => HashSet::new(),
            IrExpression::Double(_) => HashSet::new(),
            IrExpression::String(_) => HashSet::new(),
        }
    }

    fn propagate_spills(&mut self, spills: &HashSet<IrVrVariableDescriptor>) {
        match self {
            IrExpression::Field(self_variable, field_index) => {
                // no-op, no defs
            }
            IrExpression::Parameter(_) => {
                // no-op
            }
            IrExpression::Variable(ir_variable) => {
                if let IrVariableDescriptor::VirtualRegister(vr_variable) =
                    ir_variable.borrow().descriptor()
                {
                    if spills.contains(&vr_variable) {
                        ir_variable
                            .borrow_mut()
                            .set_descriptor(IrVariableDescriptor::Stack(
                                IrStackVariableDescriptor::new(
                                    vr_variable.name().into(),
                                    vr_variable.block_id(),
                                ),
                            ));
                    }
                }
            }
            IrExpression::Int(_) => {
                // no-op
            }
            IrExpression::Double(_) => {
                // no-op
            }
            IrExpression::String(_) => {
                // no-op
            }
        }
    }

    fn propagate_register_assignments(
        &mut self,
        assignments: &HashMap<IrVrVariableDescriptor, usize>,
    ) {
        match self {
            IrExpression::Field(self_variable, field_index) => {
                // no-op
            }
            IrExpression::Parameter(_) => {
                // no-op
            }
            IrExpression::Variable(ir_variable) => {
                let mut borrowed = ir_variable.borrow_mut();
                if let IrVariableDescriptor::VirtualRegister(vr_variable) =
                    borrowed.descriptor_mut()
                {
                    if assignments.contains_key(vr_variable) {
                        vr_variable.set_assigned_register(assignments[vr_variable]);
                    }
                }
            }
            IrExpression::Int(_) => {
                // no-op
            }
            IrExpression::Double(_) => {
                // no-op
            }
            IrExpression::String(_) => {
                // no-op
            }
        }
    }

    fn propagate_stack_offsets(&mut self, _counter: &mut OffsetCounter) {
        // no-op
    }
}