use crate::constants_table::ConstantsTable;
use crate::ir::assemble::{Assemble, InstructionsBuilder};
use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_variable::IrVrVariableDescriptor;
use crate::ir::register_allocation::{HasVrUsers, OffsetCounter, VrUser};
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};
use std::rc::Rc;

pub struct IrBlock {
    id: usize,
    debug_label: String,
    predecessors: Vec<Rc<RefCell<IrBlock>>>,
    successors: Vec<Rc<RefCell<IrBlock>>>,
    statements: Vec<IrStatement>,
}

impl IrBlock {
    pub fn new(id: usize, debug_label: &str, statements: Vec<IrStatement>) -> Self {
        Self {
            id,
            debug_label: debug_label.into(),
            predecessors: vec![],
            successors: vec![],
            statements,
        }
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn statements(&self) -> &[IrStatement] {
        &self.statements
    }
}

impl HasVrUsers for IrBlock {
    fn vr_users(&self) -> Vec<&dyn VrUser> {
        self.statements.iter().map(|s| s as &dyn VrUser).collect()
    }

    fn vr_users_mut(&mut self) -> Vec<&mut dyn VrUser> {
        self.statements
            .iter_mut()
            .map(|s| s as &mut dyn VrUser)
            .collect()
    }
}

impl VrUser for IrBlock {
    fn vr_definitions(&self) -> HashSet<IrVrVariableDescriptor> {
        self.statements
            .iter()
            .flat_map(|s| s.vr_definitions())
            .collect()
    }

    fn vr_uses(&self) -> HashSet<IrVrVariableDescriptor> {
        self.statements.iter().flat_map(|s| s.vr_uses()).collect()
    }

    fn propagate_spills(&mut self, spills: &HashSet<IrVrVariableDescriptor>) {
        for statement in &mut self.statements {
            statement.propagate_spills(spills);
        }
    }

    fn propagate_register_assignments(
        &mut self,
        assignments: &HashMap<IrVrVariableDescriptor, usize>,
    ) {
        for statement in &mut self.statements {
            statement.propagate_register_assignments(assignments);
        }
    }

    fn propagate_stack_offsets(&mut self, counter: &mut OffsetCounter) {
        for statement in &mut self.statements {
            statement.propagate_stack_offsets(counter);
        }
    }
}

impl Assemble for IrBlock {
    fn assemble(&self, builder: &mut InstructionsBuilder, constants_table: &mut ConstantsTable) {
        for statement in &self.statements {
            statement.assemble(builder, constants_table);
        }
    }
}

impl Display for IrBlock {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "  {}:", self.debug_label)?;

        let (live_in, live_out) = self.live_in_live_out();

        for (statement_index, statement) in self.statements.iter().enumerate() {
            let statement_live_in = live_in.get(&statement_index).unwrap();
            let statement_live_out = live_out.get(&statement_index).unwrap();

            writeln!(
                f,
                "    {} // live_in: {:?}, live_out: {:?}",
                statement, statement_live_in, statement_live_out
            )?;
        }
        writeln!(f, "    // ---- {} meta ----", self.debug_label)?;
        writeln!(f, "    // definitions: {:?}", self.vr_definitions())?;
        writeln!(f, "    // uses: {:?}", self.vr_uses())?;
        writeln!(
            f,
            "    // interference graph: {:?}",
            self.interference_graph()
        )?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::ast::module_level_declaration::ModuleLevelDeclaration;
    use crate::parser::parse_compilation_unit;
    use crate::symbol_table::SymbolTable;

    #[test]
    fn overlapping_assignments_bug_when_k_2() {
        let mut compilation_unit = parse_compilation_unit(
            "
            fn main()
                let a = 1
                let b = 2
                let c = 3
                let x = a + b + c
            end
        ",
        )
        .unwrap();
        let mut symbol_table = SymbolTable::new();
        compilation_unit.gather_declared_names(&mut symbol_table);
        compilation_unit.check_name_usages(&mut symbol_table);
        compilation_unit.type_check(&mut symbol_table);

        let main = compilation_unit
            .declarations()
            .iter()
            .find(|d| matches!(d, ModuleLevelDeclaration::Function(_)))
            .unwrap();

        if let ModuleLevelDeclaration::Function(main) = main {
            let mut main_ir = main.to_ir(&symbol_table);
            let (register_assignments, _) = main_ir.assign_registers(2);
            assert_eq!(register_assignments.len(), 4);
        } else {
            unreachable!()
        }
    }
}