use crate::ir::ir_variable::IrVariableId;
use crate::ir::register_allocation::RegisterAssignment;
use std::collections::HashMap;

pub type StackVariableOffset = isize;

pub enum VariableLocation {
    Register(RegisterAssignment),
    Stack(StackVariableOffset),
}

pub struct VariableLocations {
    register_variables: HashMap<IrVariableId, RegisterAssignment>,
    stack_variables: HashMap<IrVariableId, StackVariableOffset>,
    next_stack_variable_offset: StackVariableOffset,
}

impl VariableLocations {
    pub fn new() -> Self {
        Self {
            register_variables: HashMap::new(),
            stack_variables: HashMap::new(),
            next_stack_variable_offset: 0,
        }
    }

    pub fn push_all_register_variables(
        &mut self,
        register_variables: &HashMap<IrVariableId, RegisterAssignment>,
    ) {
        self.register_variables.extend(register_variables);
    }

    pub fn get_variable_location(&self, id: &IrVariableId) -> VariableLocation {
        match self.register_variables.get(id) {
            Some(register_assignment) => VariableLocation::Register(*register_assignment),
            None => VariableLocation::Stack(self.stack_variables[id]),
        }
    }

    pub fn register_variables_count(&self) -> usize {
        self.register_variables.len()
    }

    pub fn stack_variables_count(&self) -> usize {
        self.stack_variables.len()
    }

    pub fn extend(&mut self, other: &VariableLocations) {
        // todo: add logic to transpose the `other`'s stack assignments on top of the current ones.

        self.register_variables
            .extend(other.register_variables.clone());
        self.stack_variables.extend(other.stack_variables.clone());
    }

    pub fn push_stack_variable(&mut self, ir_variable_id: IrVariableId) {
        self.stack_variables
            .insert(ir_variable_id, self.next_stack_variable_offset);
        self.next_stack_variable_offset += 1;
    }
}