use std::collections::{BTreeMap, BTreeSet, HashSet};

use petgraph::dot::Dot;
use petgraph::graph::Graph;

use crate::ir::types::IRInstr;

pub type CfgGraph = Graph<String, String>;

#[derive(Clone)]
pub struct BasicBlock {
    pub id: usize,
    pub instrs: Vec<IRInstr>,
}

pub struct ControlFlowGraph {
    pub func_name: String,
    pub blocks: Vec<BasicBlock>,
    pub edges: Vec<(usize, usize, String)>,
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum EdgeKind {
    Branch,
    Fallthrough,
}

impl ControlFlowGraph {
    pub fn from_body(func_name: impl Into<String>, body: &[IRInstr]) -> Self {
        let blocks = split_basic_blocks(body);
        let edges = collect_edges(&blocks)
            .into_iter()
            .map(|(from, to, label, _)| (from, to, label))
            .collect();
        Self {
            func_name: func_name.into(),
            blocks,
            edges,
        }
    }

    fn add_to_graph(&self, graph: &mut CfgGraph) {
        let mut nodes = Vec::new();
        for block in &self.blocks {
            nodes.push(graph.add_node(block_dot_label(&self.func_name, block)));
        }
        for (from, to, label) in &self.edges {
            if let (Some(from), Some(to)) = (nodes.get(*from), nodes.get(*to)) {
                graph.add_edge(*from, *to, label.clone());
            }
        }
    }

    fn refresh_edges(&mut self) {
        self.edges = collect_edges(&self.blocks)
            .into_iter()
            .map(|(from, to, label, _)| (from, to, label))
            .collect();
    }

    fn remove_unreachable(&mut self) -> bool {
        if self.blocks.is_empty() {
            return false;
        }

        let edges = collect_edges(&self.blocks);
        let mut adjacency = vec![Vec::new(); self.blocks.len()];
        for (from, to, _, _) in edges {
            adjacency[from].push(to);
        }

        let mut reachable = vec![false; self.blocks.len()];
        let mut stack = vec![0];
        while let Some(index) = stack.pop() {
            if reachable[index] {
                continue;
            }
            reachable[index] = true;
            for next in &adjacency[index] {
                stack.push(*next);
            }
        }

        let old_len = self.blocks.len();
        self.blocks = self
            .blocks
            .drain(..)
            .enumerate()
            .filter_map(|(index, block)| reachable[index].then_some(block))
            .collect();
        let changed = self.blocks.len() != old_len;
        if changed {
            self.refresh_edges();
        }
        changed
    }

    fn remove_trivial_blocks(&mut self) -> bool {
        if self.blocks.is_empty() {
            return false;
        }

        let edges = collect_edges(&self.blocks);
        let explicit_pred_count = explicit_predecessor_counts(self.blocks.len(), &edges);
        let mut redirect = BTreeMap::new();
        let mut remove = HashSet::new();
        let mut fallthrough_gotos = Vec::new();

        for (index, block) in self.blocks.iter().enumerate() {
            if let Some((label, target)) = jump_only_block(block) {
                let target = resolve_label(target, &redirect);
                redirect.insert(label, target);
                for (from, to, _, kind) in &edges {
                    if *to == index && *kind == EdgeKind::Fallthrough {
                        fallthrough_gotos.push((*from, target));
                    }
                }
                remove.insert(index);
                continue;
            }

            if let Some(label) = label_only_block(block) {
                if let Some(next_index) = next_kept_index(index, self.blocks.len(), &remove) {
                    if let Some(next_label) = first_label(&self.blocks[next_index]) {
                        redirect.insert(label, resolve_label(next_label, &redirect));
                        remove.insert(index);
                    } else if explicit_pred_count[index] == 0 {
                        remove.insert(index);
                    }
                } else if explicit_pred_count[index] == 0 {
                    remove.insert(index);
                }
            }
        }

        if redirect.is_empty() && remove.is_empty() && fallthrough_gotos.is_empty() {
            return false;
        }

        for (from, target) in fallthrough_gotos {
            if !remove.contains(&from) && !ends_with_terminal(&self.blocks[from]) {
                self.blocks[from].instrs.push(IRInstr::Goto(target));
            }
        }

        for block in &mut self.blocks {
            rewrite_targets(&mut block.instrs, &redirect);
        }

        self.blocks = self
            .blocks
            .drain(..)
            .enumerate()
            .filter_map(|(index, block)| (!remove.contains(&index)).then_some(block))
            .collect();
        self.refresh_edges();
        true
    }

    fn merge_linear_blocks(&mut self) -> bool {
        let edges = collect_edges(&self.blocks);
        let pred_count = predecessor_counts(self.blocks.len(), &edges);

        let mut index = 0;
        while index + 1 < self.blocks.len() {
            let successors = edges
                .iter()
                .filter_map(|(from, to, _, _)| (*from == index).then_some(*to))
                .collect::<BTreeSet<_>>();

            if successors.len() == 1 {
                let successor = *successors.iter().next().unwrap();
                if successor == index + 1 && pred_count[successor] == 1 {
                    if matches!(self.blocks[index].instrs.last(), Some(IRInstr::Goto(_))) {
                        self.blocks[index].instrs.pop();
                    }
                    let successor_block = self.blocks.remove(successor);
                    self.blocks[index].instrs.extend(successor_block.instrs);
                    self.refresh_edges();
                    return true;
                }
            }
            index += 1;
        }

        false
    }

    fn into_instrs(self) -> Vec<IRInstr> {
        self.blocks
            .into_iter()
            .flat_map(|block| block.instrs)
            .collect()
    }
}

pub fn optimize_ir(ir: Vec<IRInstr>) -> Vec<IRInstr> {
    ir.into_iter()
        .map(|instr| match instr {
            IRInstr::DefineFunc(func, args, body) => {
                let func_name = func.name.clone();
                IRInstr::DefineFunc(func, args, optimize_body(&func_name, body))
            }
            other => other,
        })
        .collect()
}

pub fn ir_to_cfg_dot(ir: &[IRInstr]) -> String {
    let mut graph = Graph::new();
    for instr in ir {
        if let IRInstr::DefineFunc(func, _, body) = instr {
            ControlFlowGraph::from_body(func.name.clone(), body).add_to_graph(&mut graph);
        }
    }
    format!("{}", Dot::new(&graph))
}

pub fn optimize_body(func_name: &str, body: Vec<IRInstr>) -> Vec<IRInstr> {
    let (mut prologue, cfg_body) = split_function_prologue(body);
    if cfg_body.is_empty() {
        return prologue;
    }

    let mut cfg = ControlFlowGraph::from_body(func_name, &cfg_body);
    cfg.remove_unreachable();
    loop {
        let mut changed = false;
        changed |= cfg.remove_trivial_blocks();
        changed |= cfg.remove_unreachable();
        changed |= cfg.merge_linear_blocks();
        if !changed {
            break;
        }
    }

    let mut body = cfg.into_instrs();
    remove_redundant_fallthrough_gotos(&mut body);
    remove_unused_labels(&mut body);
    prologue.extend(body);
    prologue
}

fn split_function_prologue(mut body: Vec<IRInstr>) -> (Vec<IRInstr>, Vec<IRInstr>) {
    if let Some(entry_index) = body.iter().position(|instr| matches!(instr, IRInstr::Entry)) {
        let cfg_body = body.split_off(entry_index + 1);
        (body, cfg_body)
    } else {
        (Vec::new(), body)
    }
}

fn split_basic_blocks(instrs: &[IRInstr]) -> Vec<BasicBlock> {
    if instrs.is_empty() {
        return Vec::new();
    }

    let mut leaders = BTreeSet::new();
    leaders.insert(0);
    for (index, instr) in instrs.iter().enumerate() {
        if matches!(instr, IRInstr::Label(_)) {
            leaders.insert(index);
        }
        if is_terminal(instr) && index + 1 < instrs.len() {
            leaders.insert(index + 1);
        }
    }

    let leader_list = leaders.into_iter().collect::<Vec<_>>();
    let mut blocks = Vec::new();
    for (block_id, start) in leader_list.iter().enumerate() {
        let end = leader_list
            .get(block_id + 1)
            .copied()
            .unwrap_or(instrs.len());
        blocks.push(BasicBlock {
            id: block_id,
            instrs: instrs[*start..end].to_vec(),
        });
    }
    blocks
}

fn collect_edges(blocks: &[BasicBlock]) -> Vec<(usize, usize, String, EdgeKind)> {
    let label_to_block = label_to_block(blocks);
    let mut edges = Vec::new();

    for (index, block) in blocks.iter().enumerate() {
        match block.instrs.last() {
            Some(IRInstr::Goto(label)) => {
                if let Some(target) = label_to_block.get(label) {
                    edges.push((index, *target, "br".to_string(), EdgeKind::Branch));
                }
            }
            Some(IRInstr::CondGoto(_, true_label, false_label)) => {
                if let Some(target) = label_to_block.get(true_label) {
                    edges.push((index, *target, "true".to_string(), EdgeKind::Branch));
                }
                if let Some(target) = label_to_block.get(false_label) {
                    edges.push((index, *target, "false".to_string(), EdgeKind::Branch));
                }
            }
            Some(IRInstr::Exit(_)) => {}
            Some(_) | None => {
                if index + 1 < blocks.len() {
                    edges.push((index, index + 1, "fallthrough".to_string(), EdgeKind::Fallthrough));
                }
            }
        }
    }

    edges
}

fn label_to_block(blocks: &[BasicBlock]) -> BTreeMap<usize, usize> {
    let mut map = BTreeMap::new();
    for (index, block) in blocks.iter().enumerate() {
        for instr in &block.instrs {
            if let IRInstr::Label(label) = instr {
                map.insert(*label, index);
            }
        }
    }
    map
}

fn predecessor_counts(block_count: usize, edges: &[(usize, usize, String, EdgeKind)]) -> Vec<usize> {
    let mut counts = vec![0; block_count];
    for (_, to, _, _) in edges {
        counts[*to] += 1;
    }
    counts
}

fn explicit_predecessor_counts(block_count: usize, edges: &[(usize, usize, String, EdgeKind)]) -> Vec<usize> {
    let mut counts = vec![0; block_count];
    for (_, to, _, kind) in edges {
        if *kind == EdgeKind::Branch {
            counts[*to] += 1;
        }
    }
    counts
}

fn next_kept_index(index: usize, block_count: usize, remove: &HashSet<usize>) -> Option<usize> {
    (index + 1..block_count).find(|candidate| !remove.contains(candidate))
}

fn first_label(block: &BasicBlock) -> Option<usize> {
    match block.instrs.first() {
        Some(IRInstr::Label(label)) => Some(*label),
        _ => None,
    }
}

fn label_only_block(block: &BasicBlock) -> Option<usize> {
    match block.instrs.as_slice() {
        [IRInstr::Label(label)] => Some(*label),
        _ => None,
    }
}

fn jump_only_block(block: &BasicBlock) -> Option<(usize, usize)> {
    match block.instrs.as_slice() {
        [IRInstr::Label(label), IRInstr::Goto(target)] => Some((*label, *target)),
        _ => None,
    }
}

fn is_terminal(instr: &IRInstr) -> bool {
    matches!(instr, IRInstr::Goto(_) | IRInstr::CondGoto(_, _, _) | IRInstr::Exit(_))
}

fn ends_with_terminal(block: &BasicBlock) -> bool {
    block.instrs.last().is_some_and(is_terminal)
}

fn rewrite_targets(instrs: &mut [IRInstr], redirect: &BTreeMap<usize, usize>) {
    for instr in instrs {
        match instr {
            IRInstr::Goto(label) => *label = resolve_label(*label, redirect),
            IRInstr::CondGoto(_, true_label, false_label) => {
                *true_label = resolve_label(*true_label, redirect);
                *false_label = resolve_label(*false_label, redirect);
            }
            _ => {}
        }
    }
}

fn resolve_label(label: usize, redirect: &BTreeMap<usize, usize>) -> usize {
    let mut current = label;
    let mut visited = HashSet::new();
    while let Some(next) = redirect.get(&current) {
        if !visited.insert(current) {
            break;
        }
        current = *next;
    }
    current
}

fn remove_unused_labels(instrs: &mut Vec<IRInstr>) {
    let mut used = BTreeSet::new();
    for instr in instrs.iter() {
        match instr {
            IRInstr::Goto(label) => {
                used.insert(*label);
            }
            IRInstr::CondGoto(_, true_label, false_label) => {
                used.insert(*true_label);
                used.insert(*false_label);
            }
            _ => {}
        }
    }
    instrs.retain(|instr| match instr {
        IRInstr::Label(label) => used.contains(label),
        _ => true,
    });
}

fn remove_redundant_fallthrough_gotos(instrs: &mut Vec<IRInstr>) {
    let mut index = 0;
    while index + 1 < instrs.len() {
        let remove = match (&instrs[index], &instrs[index + 1]) {
            (IRInstr::Goto(target), IRInstr::Label(label)) => target == label,
            _ => false,
        };
        if remove {
            instrs.remove(index);
        } else {
            index += 1;
        }
    }
}

fn block_dot_label(func_name: &str, block: &BasicBlock) -> String {
    let instrs = block
        .instrs
        .iter()
        .map(|instr| instr.to_string())
        .collect::<Vec<_>>()
        .join("\\l");
    format!("{}::B{}\\l{}\\l", func_name, block.id, instrs)
}
