use std::{collections::HashMap, fmt};

use crate::{
    backend::dsa::instruction::{InsBlock, Instruction},
    model::CompilerError,
};

/// Register allocator for DSA assembly generation
/// Manages general-purpose registers (rg0-rgf) and handles stack spilling
pub struct RegisterAllocator {
    /// Available general-purpose registers
    /// Maps variable names to their current location (register or stack offset)
    variable_locations: HashMap<String, Location>,

    /// Maps registers to the variables they currently hold
    register_contents: HashMap<Register, String>,

    /// Current stack offset for local variables (relative to bpr)
    /// Starts at -4 (going downward from base pointer)
    stack_offset: i32,

    /// Track which registers are currently in use
    in_use: Vec<(Register, bool)>,
}

#[derive(Debug, Clone)]
pub struct Location {
    register: Option<Register>,
    stack: Option<i32>,
}

impl Location {
    pub fn stack(offset: i32) -> Self {
        Location {
            register: None,
            stack: Some(offset),
        }
    }

    pub fn register(register: Register) -> Self {
        Location {
            register: Some(register),
            stack: None,
        }
    }
}

impl RegisterAllocator {
    pub fn new() -> Self {
        // Initialize with available GP registers (rg0-rgf = 16 registers)
        let in_use = vec![
            Register::Rg0,
            Register::Rg1,
            Register::Rg2,
            Register::Rg3,
            Register::Rg4,
            Register::Rg5,
            Register::Rg6,
            Register::Rg7,
            Register::Rg8,
            Register::Rg9,
            Register::Rga,
            Register::Rgb,
            Register::Rgc,
            Register::Rgd,
            Register::Rge,
            Register::Rgf,
        ]
        .iter()
        .map(|&reg| (reg, false))
        .collect();

        RegisterAllocator {
            // available_registers: registers,
            variable_locations: HashMap::new(),
            register_contents: HashMap::new(),
            stack_offset: -4, // Start at -4 (first local below saved bpr)
            in_use,
        }
    }

    /// Allocate a temporary register for expression evaluation
    /// Returns the register name and optionally assembly code to save it
    pub fn alloc_temp(&mut self) -> Result<(Register, InsBlock), CompilerError> {
        // Try to find an unused register

        // println!("finding! {:#?}", self.in_use);

        if let Some(reg) = self.find_free_register() {
            self.in_use[reg as usize].1 = true;
            return Ok((reg, InsBlock::new()));
        }

        // All registers in use - need to spill one
        // Choose the first register with a variable we can spill
        // Find a register to spill

        // let reg_to_spill = self
        //     .available_registers
        //     .iter()
        //     .find(|reg| self.register_contents.contains_key(*reg))
        //     .cloned();

        // if let Some(reg) = reg_to_spill {
        //     // Spill this variable to stack
        //     let spill_code = self.spill_register(&reg)?;
        //     code.extend(spill_code);

        //     self.in_use.insert(reg.clone(), true);
        //     return Ok((reg, code));
        // }

        todo!("an efficient stack spilling algorithm. needs scope awareness.");

        Err(CompilerError::Generic(
            "All registers are used up yet there are no variables to spill to the stack"
                .to_string(),
        ))
    }

    // fn set_in_use(&mut self, reg: Register, in_use: bool) {
    //     self.in_use[reg as usize].1 = in_use;
    // }

    /// Free a temporary register after use
    /// NOTE: This will NOT free registers that contain variables!
    /// Variables persist throughout their scope and must not be freed
    pub fn free_temp(&mut self, reg: Register) {
        // Check if this register contains a variable
        if self.register_contents.contains_key(&reg) {
            // This register holds a variable - don't free it!
            // Variables are only freed when they go out of scope via free_var()
            return;
        }

        // This is a true temporary - safe to free
        if !matches!(reg, Register::Zero | Register::Null) {
            self.in_use[reg as usize].1 = false;
        }
    }

    pub fn free_var(&mut self, var: &str) {
        // Check if this variable is in a register
        if let Some(location) = self.variable_locations.get(var).cloned() {
            if let Some(reg) = location.register
                && !matches!(reg, Register::Zero | Register::Null)
            {
                self.register_contents.remove(&reg);
                self.in_use[reg as usize].1 = false;
            }

            self.variable_locations.remove(var);
        }
    }

    /// Allocate a register for a named variable
    /// Returns the register and any necessary assembly code
    pub fn alloc_var(
        &mut self,
        var_name: &str,
    ) -> Result<(Register, InsBlock), CompilerError> {
        if let Some(mut location) = self.variable_locations.get(var_name).cloned() {
            // if the var is in a register we can use it already.
            if let Some(reg) = location.register {
                return Ok((reg, InsBlock::new()));
            }

            // if the variable is on the stack only, we need to get it in a register.
            if let Some(offset) = location.stack {
                // Variable was pushed, need to calculate actual position and update its
                // location.
                let (reg, mut code) = self.alloc_temp()?;

                // acknowledge var is now in a reg as well.
                location.register = Some(reg);

                // Load from bpr + offset (offset is negative)
                // code.push(format!("\tsubi bpr {} {}", -(offset + 4), reg));

                code.push(Instruction::ldw_reg_offset(
                    Register::Spr,
                    reg,
                    offset - self.stack_offset,
                ));

                // Update location to register
                self.variable_locations
                    .insert(var_name.to_string(), location);
                self.register_contents.insert(reg, var_name.to_string());

                return Ok((reg, code));
            }
        }

        // Variable doesn't have a location yet, allocate a new register
        let (reg, code) = self.alloc_temp()?;
        self.variable_locations
            .insert(var_name.to_string(), Location::register(reg));
        self.register_contents.insert(reg, var_name.to_string());

        Ok((reg, code))
    }

    /// Get the current location of a variable
    pub fn _get_var_location(&self, var_name: &str) -> Option<&Location> {
        self.variable_locations.get(var_name)
    }

    /// Load a variable into a register (allocating if necessary)
    /// Returns the register and assembly code to load it
    pub fn load_var(
        &mut self,
        var_name: &str,
    ) -> Result<(Register, InsBlock), CompilerError> {
        self.alloc_var(var_name)
    }

    /// Store a value from a register into a variable
    /// Updates tracking and returns any necessary assembly code
    pub fn store_var(&mut self, var_name: &str, source_reg: &Register) -> InsBlock {
        let mut block = InsBlock::new();

        // Check if variable already has a location
        if let Some(location) = self.variable_locations.get(var_name) {
            // if the variable exists in a register we write to that.
            match location.register {
                Some(reg) if reg == *source_reg => {
                    block.push(Instruction::mov(*source_reg, reg));
                    return block;
                }
                _ => (),
            }

            // if the variable exists on the stack but not a register we write here.
            if let Some(offset) = location.stack {
                block.push(Instruction::stw_reg_offset(
                    *source_reg,
                    Register::Spr,
                    offset - self.stack_offset,
                ));
                return block;
            }
        }

        // Variable doesn't exist yet, we can just use the same reg.
        // if we can avoid a move, absolutely do that.

        // if this is true then there's no permanent variable here so it's safe to use.
        if !self.register_contents.contains_key(source_reg) {
            self.variable_locations
                .insert(var_name.to_string(), Location::register(*source_reg));
            self.register_contents
                .insert(*source_reg, var_name.to_string());
            self.in_use[*source_reg as usize].1 = true;

            return block;
        }

        // if current register isn't free, (eg is another variable) we assign somewhere
        // else.
        if let Some(free_reg) = self.find_free_register() {
            self.variable_locations
                .insert(var_name.to_string(), Location::register(free_reg));
            self.register_contents
                .insert(free_reg, var_name.to_string());
            self.in_use[free_reg as usize].1 = true;

            block.push(Instruction::mov(*source_reg, free_reg));
            return block;
        }

        // No free registers - allocate on stack
        // code.push(format!("\tstw {}, bpr, {}", source_reg, self.stack_offset));
        // self.variable_locations
        //     .insert(var_name.to_string(), Location::Stack(self.stack_offset));
        // self.stack_offset -= 4; // Move to next stack slot
        //
        todo!("an efficient stack spilling algorithm. needs scope awareness.");
    }

    /// spill a register to the stack (WITHOUT FREEING)
    /// DO NOT USE this if it's for a pointer!!!!
    pub fn _spill_register(&mut self, reg: &Register) -> Result<InsBlock, CompilerError> {
        let mut code = InsBlock::new();

        // check if the variable is declared.
        if let Some(var_name) = self.register_contents.get(reg).cloned()
            && let Some(location) = self.variable_locations.get_mut(&var_name)
        {
            // check if var is on the stack
            if let Some(offset) = location.stack {
                code.push(Instruction::stw_reg_offset(
                    *reg,
                    Register::Spr,
                    offset - self.stack_offset,
                ));
                return Ok(code);
            }

            // Track that we pushed one word
            self.stack_offset -= 4;

            // if the variable is not on the stack:
            // push register to stack (spr decrements automatically)
            let offset = self.stack_offset;
            code.push(Instruction::push(*reg));

            // Update variable location - it's now at current spr
            // Note: We track offset from bpr for consistency
            location.stack = Some(offset);

            Ok(code)
        } else {
            Err(CompilerError::Generic(format!(
                "Register {} does not contain a variable to spill!",
                reg
            )))
        }
    }

    /// free a register by spilling it to the stack.
    /// Returns assembly code to perform the spill
    pub fn free_register(
        &mut self,
        reg: &Register,
    ) -> Result<(i32, Instruction), CompilerError> {
        // check if the variable is declared.
        if let Some(var_name) = self.register_contents.get(reg).cloned()
            && let Some(location) = self.variable_locations.get_mut(&var_name)
        {
            // check if var name is on the stack
            if let Some(offset) = location.stack {
                // store current register value in stack location
                let code = Instruction::stw_reg_offset(
                    *reg,
                    Register::Spr,
                    offset - self.stack_offset,
                );

                // free the register.
                location.register = None;
                self.register_contents.remove(reg);
                return Ok((offset, code));
            }

            // Track that we pushed one word
            self.stack_offset -= 4;

            let offset = self.stack_offset;
            let code = Instruction::push(*reg);

            // Update variable location
            // Note: We track offset from bpr for consistency
            location.stack = Some(offset);
            location.register = None;
            self.register_contents.remove(reg);

            Ok((offset, code))
        } else {
            Err(CompilerError::Generic(format!(
                "Register {} does not contain a variable to spill!",
                reg
            )))
        }
    }

    /// Find a free register (not currently in use)
    fn find_free_register(&self) -> Option<Register> {
        self.in_use
            .iter()
            .filter(|(_, in_use)| !*in_use)
            .map(|(reg, _)| *reg)
            .next()
    }

    /// Spill all registers to stack (useful before function calls)
    pub fn _spill_all(&mut self) -> InsBlock {
        let mut code = InsBlock::new();

        let regs_to_spill: Vec<Register> =
            self.register_contents.keys().cloned().collect();

        for reg in regs_to_spill {
            if let Ok(spill_code) = self.free_register(&reg) {
                code.push(spill_code.1);
            }
        }

        code
    }

    /// Get the total stack offset
    pub fn get_stack_offset(&self) -> i32 {
        self.stack_offset
    }

    /// Get the total stack space needed for local variables
    pub fn _get_stack_size(&self) -> i32 {
        -self.stack_offset // Convert negative offset to positive size
    }

    /// Reset allocator for a new function
    pub fn reset(&mut self) {
        self.variable_locations.clear();
        self.register_contents.clear();
        self.stack_offset = -4;
        self.in_use = vec![
            Register::Rg0,
            Register::Rg1,
            Register::Rg2,
            Register::Rg3,
            Register::Rg4,
            Register::Rg5,
            Register::Rg6,
            Register::Rg7,
            Register::Rg8,
            Register::Rg9,
            Register::Rga,
            Register::Rgb,
            Register::Rgc,
            Register::Rgd,
            Register::Rge,
            Register::Rgf,
        ]
        .iter()
        .map(|&reg| (reg, false))
        .collect();
    }

    /// Get list of registers that contain variables and are in use
    /// These need to be saved before function calls
    pub fn get_caller_saved_registers(&self) -> Vec<Register> {
        self.register_contents
            .iter()
            .filter(|(reg, _)| {
                self.in_use
                    .get(**reg as usize)
                    .unwrap_or(&(Register::Null, false))
                    .1
            })
            .map(|(reg, _)| *reg)
            .collect()
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum Register {
    // general purpose
    Rg0 = 0,
    Rg1 = 1,
    Rg2 = 2,
    Rg3 = 3,
    Rg4 = 4,
    Rg5 = 5,
    Rg6 = 6,
    Rg7 = 7,
    Rg8 = 8,
    Rg9 = 9,
    Rga = 10,
    Rgb = 11,
    Rgc = 12,
    Rgd = 13,
    Rge = 14,
    Rgf = 15,

    // special
    Bpr,
    Spr,
    Ret,
    Acc,

    // read only
    Pcx,
    Zero,

    // null
    Null,
}

impl Register {
    pub fn get_gp() -> [Register; 16] {
        [
            Register::Rg0,
            Register::Rg1,
            Register::Rg2,
            Register::Rg3,
            Register::Rg4,
            Register::Rg5,
            Register::Rg6,
            Register::Rg7,
            Register::Rg8,
            Register::Rg9,
            Register::Rga,
            Register::Rgb,
            Register::Rgc,
            Register::Rgd,
            Register::Rge,
            Register::Rgf,
        ]
    }

    pub fn is_gp(&self) -> bool {
        (*self as u8) < 16
    }

    pub fn from_index(idx: usize) -> Register {
        match idx {
            0 => Register::Rg0,
            1 => Register::Rg1,
            2 => Register::Rg2,
            3 => Register::Rg3,
            4 => Register::Rg4,
            5 => Register::Rg5,
            6 => Register::Rg6,
            7 => Register::Rg7,
            8 => Register::Rg8,
            9 => Register::Rg9,
            10 => Register::Rga,
            11 => Register::Rgb,
            12 => Register::Rgc,
            13 => Register::Rgd,
            14 => Register::Rge,
            15 => Register::Rgf,
            _ => unreachable!("this function shouldn't ever be called with idx>15"),
        }
    }
}

impl fmt::Display for Register {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Rg0 => write!(f, "rg0"),
            Self::Rg1 => write!(f, "rg1"),
            Self::Rg2 => write!(f, "rg2"),
            Self::Rg3 => write!(f, "rg3"),
            Self::Rg4 => write!(f, "rg4"),
            Self::Rg5 => write!(f, "rg5"),
            Self::Rg6 => write!(f, "rg6"),
            Self::Rg7 => write!(f, "rg7"),
            Self::Rg8 => write!(f, "rg8"),
            Self::Rg9 => write!(f, "rg9"),
            Self::Rga => write!(f, "rga"),
            Self::Rgb => write!(f, "rgb"),
            Self::Rgc => write!(f, "rgc"),
            Self::Rgd => write!(f, "rgd"),
            Self::Rge => write!(f, "rge"),
            Self::Rgf => write!(f, "rgf"),

            Self::Acc => write!(f, "acc"),
            Self::Ret => write!(f, "ret"),
            Self::Bpr => write!(f, "bpr"),
            Self::Spr => write!(f, "spr"),

            Self::Zero => write!(f, "zero"),
            Self::Pcx => write!(f, "pcx"),

            Self::Null => write!(f, "null"),
        }
    }
}
