use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_variable::IrVirtualRegisterVariable;
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>,
}

type LivenessMapByStatement = HashMap<usize, HashSet<Rc<IrVirtualRegisterVariable>>>;

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
    }

    fn vr_definitions(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
        let mut set = HashSet::new();
        for statement in &self.statements {
            set.extend(statement.vr_definitions());
        }
        set
    }

    fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
        let mut set = HashSet::new();
        for statement in &self.statements {
            set.extend(statement.vr_uses());
        }
        set
    }

    fn live_in_live_out(&self) -> (LivenessMapByStatement, LivenessMapByStatement) {
        let mut live_in: LivenessMapByStatement = HashMap::new();
        let mut live_out: LivenessMapByStatement = HashMap::new();

        loop {
            let mut did_work = false;
            for (statement_index, statement) in self.statements.iter().enumerate().rev() {
                // init if necessary
                if !live_in.contains_key(&statement_index) {
                    live_in.insert(statement_index, HashSet::new());
                }
                if !live_out.contains_key(&statement_index) {
                    live_out.insert(statement_index, HashSet::new());
                }

                // out (union of successors ins)
                // for now, a statement can only have one successor
                // this will need to be updated when we add jumps
                if let Some(successor_live_in) = live_in.get(&(statement_index + 1)) {
                    let statement_live_out = live_out.get_mut(&statement_index).unwrap();
                    for vr_variable in successor_live_in {
                        if statement_live_out.insert(vr_variable.clone()) {
                            did_work = true;
                        }
                    }
                }

                // in: use(s) U ( out(s) - def(s) )
                let mut new_ins = statement.vr_uses();
                let statement_live_out = live_out.get(&statement_index).unwrap();
                let defs = statement.vr_definitions();
                let rhs = statement_live_out - &defs;
                new_ins.extend(rhs);

                let statement_live_in = live_in.get_mut(&statement_index).unwrap();
                for new_in in new_ins {
                    if statement_live_in.insert(new_in) {
                        did_work = true;
                    }
                }
            }
            if !did_work {
                break;
            }
        }
        (live_in, live_out)
    }

    fn interference_graph(
        &self,
    ) -> HashMap<Rc<IrVirtualRegisterVariable>, HashSet<Rc<IrVirtualRegisterVariable>>> {
        let mut all_vr_variables: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
        for statement in &self.statements {
            all_vr_variables.extend(statement.vr_definitions());
            all_vr_variables.extend(statement.vr_uses());
        }

        let mut graph: HashMap<
            Rc<IrVirtualRegisterVariable>,
            HashSet<Rc<IrVirtualRegisterVariable>>,
        > = HashMap::new();
        for variable in all_vr_variables {
            graph.insert(variable, HashSet::new());
        }

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

        for (statement_index, statement) in self.statements.iter().enumerate() {
            let statement_live_out = live_out.get(&statement_index).unwrap();
            for definition_vr_variable in statement.vr_definitions() {
                for live_out_variable in statement_live_out {
                    // we do the following check to avoid adding an edge to itself
                    if definition_vr_variable != *live_out_variable {
                        graph
                            .get_mut(&definition_vr_variable)
                            .unwrap()
                            .insert(live_out_variable.clone());
                        graph
                            .get_mut(live_out_variable)
                            .unwrap()
                            .insert(definition_vr_variable.clone());
                    }
                }
            }
        }

        graph
    }

    pub fn assign_registers(&mut self) {
        let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
        loop {
            let (registers, new_spills) = register_assignment::assign_registers(self);
            if spills != new_spills {
                spills = new_spills;
                // mutate all IrVirtualRegisters to constituent statements
                for statement in &mut self.statements {
                    statement.propagate_spills(&spills);
                }
            } else {
                println!("{:?}", registers);
                break;
            }
        }
    }
}

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(())
    }
}

mod register_assignment {
    use super::*;

    #[derive(Debug)]
    struct WorkItem {
        vr: Rc<IrVirtualRegisterVariable>,
        edges: HashSet<Rc<IrVirtualRegisterVariable>>,
        color: bool,
    }

    pub fn assign_registers(
        block: &IrBlock,
    ) -> (
        HashMap<Rc<IrVirtualRegisterVariable>, usize>,
        HashSet<Rc<IrVirtualRegisterVariable>>,
    ) {
        let k = 8;
        let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();

        loop {
            // 1. get interference graph
            let mut interference_graph = block.interference_graph();
            let mut work_stack: Vec<WorkItem> = vec![];

            loop {
                // 2. coloring by simplification
                //    try to find a node (virtual register) with less than k outgoing edges,
                //    and mark as color
                //    if not, pick any, and mark as spill for step 3
                let register_lt_k = interference_graph.iter().find_map(|(vr, neighbors)| {
                    if neighbors.len() < k { Some(vr) } else { None }
                });
                if let Some(vr) = register_lt_k {
                    let vr = vr.clone();
                    // remove both outgoing and incoming edges; save either set for WorkItem
                    // first, outgoing:
                    let outgoing_edges = interference_graph.remove(&vr).unwrap();

                    // second, incoming
                    interference_graph.iter_mut().for_each(|(_, neighbors)| {
                        neighbors.remove(&vr);
                    });

                    // push to work stack
                    work_stack.push(WorkItem {
                        vr,
                        edges: outgoing_edges,
                        color: true,
                    })
                } else {
                    // pick any
                    let vr = interference_graph.iter().last().unwrap().0.clone();

                    // first, outgoing
                    let outgoing_edges = interference_graph.remove(&vr).unwrap();

                    // second, incoming
                    interference_graph.iter_mut().for_each(|(_, neighbors)| {
                        neighbors.remove(&vr);
                    });

                    work_stack.push(WorkItem {
                        vr,
                        edges: outgoing_edges,
                        color: false, // spill
                    });
                }

                if interference_graph.is_empty() {
                    break;
                }
            }

            // 3. assign colors to registers
            let mut rebuilt_graph: HashMap<
                Rc<IrVirtualRegisterVariable>,
                HashSet<Rc<IrVirtualRegisterVariable>>,
            > = HashMap::new();
            let mut register_assignments: HashMap<Rc<IrVirtualRegisterVariable>, usize> =
                HashMap::new();
            let mut new_spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();

            while let Some(work_item) = work_stack.pop() {
                if work_item.color {
                    assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
                } else {
                    // first, see if we can optimistically color
                    // find how many assignments have been made for the outgoing edges
                    // if it's less than k, we can do it
                    let mut number_of_assigned_edges = 0;
                    for edge in &work_item.edges {
                        if register_assignments.contains_key(edge) {
                            number_of_assigned_edges += 1;
                        }
                    }

                    if number_of_assigned_edges < k {
                        // optimistically color
                        assign_register(
                            &work_item,
                            &mut rebuilt_graph,
                            k,
                            &mut register_assignments,
                        );
                    } else {
                        // spill
                        new_spills.insert(work_item.vr.clone());
                    }
                }
            }

            if spills.eq(&new_spills) {
                return (register_assignments, spills);
            } else {
                spills = new_spills;
            }
        }
    }

    fn assign_register(
        work_item: &WorkItem,
        rebuilt_graph: &mut HashMap<
            Rc<IrVirtualRegisterVariable>,
            HashSet<Rc<IrVirtualRegisterVariable>>,
        >,
        k: usize,
        register_assignments: &mut HashMap<Rc<IrVirtualRegisterVariable>, usize>,
    ) {
        let this_vertex_vr = &work_item.vr;
        // init the vertex
        rebuilt_graph.insert(this_vertex_vr.clone(), HashSet::new());

        // add edges, both outgoing and incoming
        let neighbors = rebuilt_graph.get_mut(this_vertex_vr).unwrap();
        // outgoing
        for edge in &work_item.edges {
            neighbors.insert(edge.clone());
        }
        // incoming
        for neighbor in neighbors.clone() {
            if rebuilt_graph.contains_key(&neighbor) {
                rebuilt_graph
                    .get_mut(&neighbor)
                    .unwrap()
                    .insert(this_vertex_vr.clone());
            }
        }

        // find a register which is not yet shared by all outgoing edges' vertices
        // I think the bug is somewhere here
        'outer: for i in 0..k {
            for edge in rebuilt_graph.get_mut(this_vertex_vr).unwrap().iter() {
                if register_assignments.contains_key(edge) {
                    let assignment = register_assignments.get(edge).unwrap();
                    if *assignment == i {
                        continue 'outer;
                    }
                }
            }
            register_assignments.insert(this_vertex_vr.clone(), i);
            break;
        }
    }
}