//! Topological Sort with Cycle Detection

use crate::algorithms::astar::CfgGraphNode;
use std::collections::{HashMap, HashSet, VecDeque};

#[derive(Debug, Clone)]
pub enum TopoError {
    CycleDetected {
        cycle: Vec<u64>,
        explanation: String,
    },
}

impl std::fmt::Display for TopoError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TopoError::CycleDetected { cycle, explanation } => {
                write!(
                    f,
                    "Cycle detected: {}. Cycle path: {:?}",
                    explanation, cycle
                )
            }
        }
    }
}

impl std::error::Error for TopoError {}

#[derive(Debug, Clone)]
pub struct TopoResult {
    pub ordering: Vec<u64>,
    pub levels: HashMap<u64, usize>,
}

pub fn topological_sort(nodes: &[CfgGraphNode]) -> Result<TopoResult, TopoError> {
    let node_map: HashMap<u64, &CfgGraphNode> = nodes.iter().map(|n| (n.id, n)).collect();

    let mut in_degree: HashMap<u64, usize> = HashMap::new();
    for node in nodes {
        in_degree.entry(node.id).or_insert(0);
        for &succ in &node.successors {
            *in_degree.entry(succ).or_insert(0) += 1;
        }
    }

    let mut queue: VecDeque<u64> = VecDeque::new();
    for (&node_id, &deg) in &in_degree {
        if deg == 0 {
            queue.push_back(node_id);
        }
    }

    let mut ordering: Vec<u64> = Vec::new();
    let mut levels: HashMap<u64, usize> = HashMap::new();

    while let Some(node_id) = queue.pop_front() {
        ordering.push(node_id);

        let node_level = *levels.get(&node_id).unwrap_or(&0);

        if let Some(node) = node_map.get(&node_id) {
            for &succ_id in &node.successors {
                let succ_in_degree = in_degree.get_mut(&succ_id).unwrap();
                *succ_in_degree -= 1;

                levels
                    .entry(succ_id)
                    .and_modify(|l| *l = (*l).max(node_level + 1))
                    .or_insert(node_level + 1);

                if *succ_in_degree == 0 {
                    queue.push_back(succ_id);
                }
            }
        }
    }

    if ordering.len() != nodes.len() {
        let cycle = find_cycle(nodes, &node_map);
        let explanation = format!(
            "Topological sort requires a DAG. Found {} nodes that cannot be ordered.",
            nodes.len() - ordering.len()
        );
        return Err(TopoError::CycleDetected { cycle, explanation });
    }

    Ok(TopoResult { ordering, levels })
}

fn find_cycle(nodes: &[CfgGraphNode], node_map: &HashMap<u64, &CfgGraphNode>) -> Vec<u64> {
    let mut visited: HashSet<u64> = HashSet::new();
    let mut rec_stack: HashSet<u64> = HashSet::new();
    let mut parent: HashMap<u64, u64> = HashMap::new();

    for node in nodes {
        if !visited.contains(&node.id) {
            if let Some(cycle) =
                dfs_cycle(node.id, node_map, &mut visited, &mut rec_stack, &mut parent)
            {
                return cycle;
            }
        }
    }
    vec![]
}

fn dfs_cycle(
    node_id: u64,
    node_map: &HashMap<u64, &CfgGraphNode>,
    visited: &mut HashSet<u64>,
    rec_stack: &mut HashSet<u64>,
    parent: &mut HashMap<u64, u64>,
) -> Option<Vec<u64>> {
    visited.insert(node_id);
    rec_stack.insert(node_id);

    if let Some(node) = node_map.get(&node_id) {
        for &succ_id in &node.successors {
            if !visited.contains(&succ_id) {
                parent.insert(succ_id, node_id);
                if let Some(cycle) = dfs_cycle(succ_id, node_map, visited, rec_stack, parent) {
                    return Some(cycle);
                }
            } else if rec_stack.contains(&succ_id) {
                let mut cycle = vec![succ_id];
                let mut current = node_id;
                while current != succ_id {
                    cycle.push(current);
                    current = *parent.get(&current).unwrap_or(&succ_id);
                }
                cycle.push(succ_id);
                cycle.reverse();
                return Some(cycle);
            }
        }
    }

    rec_stack.remove(&node_id);
    None
}

pub fn is_dag(nodes: &[CfgGraphNode]) -> bool {
    topological_sort(nodes).is_ok()
}

pub fn critical_path_length(nodes: &[CfgGraphNode]) -> Option<usize> {
    topological_sort(nodes)
        .ok()
        .and_then(|r| r.levels.values().copied().max())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_topo_sort_linear() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![2],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        let result = topological_sort(&nodes).unwrap();
        assert_eq!(result.ordering, vec![0, 1, 2]);
    }

    #[test]
    fn test_topo_sort_diamond() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1, 2],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 3,
                x: 3.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        let result = topological_sort(&nodes).unwrap();
        assert_eq!(result.ordering.len(), 4);
    }

    #[test]
    fn test_topo_sort_cycle() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![2],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![0],
            },
        ];
        assert!(topological_sort(&nodes).is_err());
    }

    #[test]
    fn test_is_dag() {
        let dag = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        assert!(is_dag(&dag));
        let cyclic = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![0],
            },
        ];
        assert!(!is_dag(&cyclic));
    }

    #[test]
    fn test_critical_path_length() {
        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1, 2],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
            CfgGraphNode {
                id: 3,
                x: 3.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        assert_eq!(critical_path_length(&nodes), Some(2));
    }
}
