use std::collections::HashMap;

use petgraph::graph::NodeIndex;
use petgraph::visit::EdgeRef;
use tracing::info;

use super::engine::CodeGraph;

/// Run Louvain community detection on the code graph.
/// Returns the number of communities detected.
/// Mutates node community_id in place.
pub fn detect_communities(code_graph: &CodeGraph) -> u32 {
    let graph = &code_graph.graph;
    let node_count = graph.node_count();

    if node_count == 0 {
        return 0;
    }

    // Initialize: each node in its own community
    let mut community: HashMap<NodeIndex, u32> = HashMap::new();
    for idx in graph.node_indices() {
        community.insert(idx, idx.index() as u32);
    }

    // Compute total edge weight (all edges weight 1.0)
    let total_edges = graph.edge_count() as f64;
    if total_edges == 0.0 {
        // All nodes are isolated, each is its own community
        return node_count as u32;
    }

    let m2 = 2.0 * total_edges;

    // Pre-compute node degrees
    let mut degree: HashMap<NodeIndex, f64> = HashMap::new();
    for idx in graph.node_indices() {
        let d = graph.edges(idx).count() as f64;
        degree.insert(idx, d);
    }

    // Louvain phase 1: local moves
    let mut improved = true;
    let mut iterations = 0;
    let max_iterations = 50;

    while improved && iterations < max_iterations {
        improved = false;
        iterations += 1;

        for node in graph.node_indices() {
            let current_comm = community[&node];
            let node_deg = degree[&node];

            // Compute edges to each neighboring community
            let mut comm_edges: HashMap<u32, f64> = HashMap::new();
            for edge in graph.edges(node) {
                let neighbor = edge.target();
                let neighbor_comm = community[&neighbor];
                *comm_edges.entry(neighbor_comm).or_insert(0.0) += 1.0;
            }
            // Also check incoming edges (undirected treatment)
            for edge in graph.edges_directed(node, petgraph::Direction::Incoming) {
                let neighbor = edge.source();
                let neighbor_comm = community[&neighbor];
                *comm_edges.entry(neighbor_comm).or_insert(0.0) += 1.0;
            }

            // Compute community totals (sum of degrees in each community)
            let mut comm_totals: HashMap<u32, f64> = HashMap::new();
            for (n, &c) in &community {
                *comm_totals.entry(c).or_insert(0.0) += degree[n];
            }

            // Find best community
            let current_total = comm_totals.get(&current_comm).copied().unwrap_or(0.0);
            let edges_to_current = comm_edges.get(&current_comm).copied().unwrap_or(0.0);

            // Modularity gain from removing node from current community
            let remove_cost = edges_to_current - (current_total - node_deg) * node_deg / m2;

            let mut best_comm = current_comm;
            let mut best_gain = 0.0;

            for (&candidate_comm, &edges_to_candidate) in &comm_edges {
                if candidate_comm == current_comm {
                    continue;
                }
                let candidate_total = comm_totals.get(&candidate_comm).copied().unwrap_or(0.0);

                // Modularity gain from adding node to candidate community
                let add_gain = edges_to_candidate - candidate_total * node_deg / m2;
                let gain = add_gain - remove_cost;

                if gain > best_gain {
                    best_gain = gain;
                    best_comm = candidate_comm;
                }
            }

            if best_comm != current_comm {
                community.insert(node, best_comm);
                improved = true;
            }
        }
    }

    // Renumber communities to be contiguous
    let mut comm_remap: HashMap<u32, u32> = HashMap::new();
    let mut next_id: u32 = 0;
    for &c in community.values() {
        if let std::collections::hash_map::Entry::Vacant(e) = comm_remap.entry(c) {
            e.insert(next_id);
            next_id += 1;
        }
    }

    // Apply to community map
    for c in community.values_mut() {
        if let Some(&new_id) = comm_remap.get(c) {
            *c = new_id;
        }
    }

    let num_communities = next_id;
    info!(
        communities = num_communities,
        iterations, "Community detection complete"
    );

    // NOTE: community IDs are stored in the HashMap but need to be applied
    // back to the CodeGraph nodes by the caller (engine) if needed for persistence.
    // For now we return the count; the full assignment is available via the map.

    num_communities
}

/// Apply community assignments back to code nodes
pub fn apply_communities(code_graph: &mut CodeGraph) -> u32 {
    detect_communities_with_assignment(code_graph)
}

/// Detect communities and write assignments into the nodes
fn detect_communities_with_assignment(code_graph: &mut CodeGraph) -> u32 {
    let graph = &code_graph.graph;
    let node_count = graph.node_count();

    if node_count == 0 {
        return 0;
    }

    let mut community: HashMap<NodeIndex, u32> = HashMap::new();
    for idx in graph.node_indices() {
        community.insert(idx, idx.index() as u32);
    }

    let total_edges = graph.edge_count() as f64;
    if total_edges == 0.0 {
        for node in &mut code_graph.nodes {
            if let Some(gi) = node.graph_index {
                node.community_id = Some(gi);
            }
        }
        return node_count as u32;
    }

    let m2 = 2.0 * total_edges;

    let mut degree: HashMap<NodeIndex, f64> = HashMap::new();
    for idx in graph.node_indices() {
        let d = (graph.edges(idx).count()
            + graph
                .edges_directed(idx, petgraph::Direction::Incoming)
                .count()) as f64;
        degree.insert(idx, d);
    }

    let mut improved = true;
    let mut iterations = 0;
    let max_iterations = 50;

    while improved && iterations < max_iterations {
        improved = false;
        iterations += 1;

        for node in graph.node_indices() {
            let current_comm = community[&node];
            let node_deg = degree[&node];

            let mut comm_edges: HashMap<u32, f64> = HashMap::new();
            for edge in graph.edges(node) {
                let neighbor_comm = community[&edge.target()];
                *comm_edges.entry(neighbor_comm).or_insert(0.0) += 1.0;
            }
            for edge in graph.edges_directed(node, petgraph::Direction::Incoming) {
                let neighbor_comm = community[&edge.source()];
                *comm_edges.entry(neighbor_comm).or_insert(0.0) += 1.0;
            }

            let mut comm_totals: HashMap<u32, f64> = HashMap::new();
            for (n, &c) in &community {
                *comm_totals.entry(c).or_insert(0.0) += degree[n];
            }

            let current_total = comm_totals.get(&current_comm).copied().unwrap_or(0.0);
            let edges_to_current = comm_edges.get(&current_comm).copied().unwrap_or(0.0);
            let remove_cost = edges_to_current - (current_total - node_deg) * node_deg / m2;

            let mut best_comm = current_comm;
            let mut best_gain = 0.0;

            for (&candidate_comm, &edges_to_candidate) in &comm_edges {
                if candidate_comm == current_comm {
                    continue;
                }
                let candidate_total = comm_totals.get(&candidate_comm).copied().unwrap_or(0.0);
                let add_gain = edges_to_candidate - candidate_total * node_deg / m2;
                let gain = add_gain - remove_cost;

                if gain > best_gain {
                    best_gain = gain;
                    best_comm = candidate_comm;
                }
            }

            if best_comm != current_comm {
                community.insert(node, best_comm);
                improved = true;
            }
        }
    }

    // Renumber
    let mut comm_remap: HashMap<u32, u32> = HashMap::new();
    let mut next_id: u32 = 0;
    for &c in community.values() {
        if let std::collections::hash_map::Entry::Vacant(e) = comm_remap.entry(c) {
            e.insert(next_id);
            next_id += 1;
        }
    }

    // Apply to nodes
    for node in &mut code_graph.nodes {
        if let Some(gi) = node.graph_index {
            let idx = NodeIndex::new(gi as usize);
            if let Some(&comm) = community.get(&idx) {
                let remapped = comm_remap.get(&comm).copied().unwrap_or(comm);
                node.community_id = Some(remapped);
            }
        }
    }

    next_id
}

#[cfg(test)]
mod tests {
    use super::*;
    use compliance_core::models::graph::{CodeEdgeKind, CodeNode, CodeNodeKind};
    use petgraph::graph::DiGraph;

    fn make_node(qualified_name: &str, graph_index: u32) -> CodeNode {
        CodeNode {
            id: None,
            repo_id: "test".to_string(),
            graph_build_id: "build1".to_string(),
            qualified_name: qualified_name.to_string(),
            name: qualified_name.to_string(),
            kind: CodeNodeKind::Function,
            file_path: "test.rs".to_string(),
            start_line: 1,
            end_line: 10,
            language: "rust".to_string(),
            community_id: None,
            is_entry_point: false,
            graph_index: Some(graph_index),
        }
    }

    fn make_empty_code_graph() -> CodeGraph {
        CodeGraph {
            graph: DiGraph::new(),
            node_map: HashMap::new(),
            nodes: Vec::new(),
            edges: Vec::new(),
        }
    }

    #[test]
    fn test_detect_communities_empty_graph() {
        let cg = make_empty_code_graph();
        assert_eq!(detect_communities(&cg), 0);
    }

    #[test]
    fn test_detect_communities_single_node_no_edges() {
        let mut graph = DiGraph::new();
        let idx = graph.add_node("a".to_string());
        let mut node_map = HashMap::new();
        node_map.insert("a".to_string(), idx);

        let cg = CodeGraph {
            graph,
            node_map,
            nodes: vec![make_node("a", 0)],
            edges: Vec::new(),
        };
        // Single node with no edges => 1 community (itself)
        assert_eq!(detect_communities(&cg), 1);
    }

    #[test]
    fn test_detect_communities_isolated_nodes() {
        let mut graph = DiGraph::new();
        let a = graph.add_node("a".to_string());
        let b = graph.add_node("b".to_string());
        let c = graph.add_node("c".to_string());
        let mut node_map = HashMap::new();
        node_map.insert("a".to_string(), a);
        node_map.insert("b".to_string(), b);
        node_map.insert("c".to_string(), c);

        let cg = CodeGraph {
            graph,
            node_map,
            nodes: vec![make_node("a", 0), make_node("b", 1), make_node("c", 2)],
            edges: Vec::new(),
        };
        // 3 isolated nodes => 3 communities
        assert_eq!(detect_communities(&cg), 3);
    }

    #[test]
    fn test_detect_communities_fully_connected() {
        let mut graph = DiGraph::new();
        let a = graph.add_node("a".to_string());
        let b = graph.add_node("b".to_string());
        let c = graph.add_node("c".to_string());
        graph.add_edge(a, b, CodeEdgeKind::Calls);
        graph.add_edge(b, c, CodeEdgeKind::Calls);
        graph.add_edge(c, a, CodeEdgeKind::Calls);

        let mut node_map = HashMap::new();
        node_map.insert("a".to_string(), a);
        node_map.insert("b".to_string(), b);
        node_map.insert("c".to_string(), c);

        let cg = CodeGraph {
            graph,
            node_map,
            nodes: vec![make_node("a", 0), make_node("b", 1), make_node("c", 2)],
            edges: Vec::new(),
        };
        let num = detect_communities(&cg);
        // Fully connected triangle should converge to 1 community
        assert!(num >= 1);
        assert!(num <= 3);
    }

    #[test]
    fn test_detect_communities_two_clusters() {
        // Two separate triangles connected by a single weak edge
        let mut graph = DiGraph::new();
        let a = graph.add_node("a".to_string());
        let b = graph.add_node("b".to_string());
        let c = graph.add_node("c".to_string());
        let d = graph.add_node("d".to_string());
        let e = graph.add_node("e".to_string());
        let f = graph.add_node("f".to_string());

        // Cluster 1: a-b-c fully connected
        graph.add_edge(a, b, CodeEdgeKind::Calls);
        graph.add_edge(b, a, CodeEdgeKind::Calls);
        graph.add_edge(b, c, CodeEdgeKind::Calls);
        graph.add_edge(c, b, CodeEdgeKind::Calls);
        graph.add_edge(a, c, CodeEdgeKind::Calls);
        graph.add_edge(c, a, CodeEdgeKind::Calls);

        // Cluster 2: d-e-f fully connected
        graph.add_edge(d, e, CodeEdgeKind::Calls);
        graph.add_edge(e, d, CodeEdgeKind::Calls);
        graph.add_edge(e, f, CodeEdgeKind::Calls);
        graph.add_edge(f, e, CodeEdgeKind::Calls);
        graph.add_edge(d, f, CodeEdgeKind::Calls);
        graph.add_edge(f, d, CodeEdgeKind::Calls);

        let mut node_map = HashMap::new();
        node_map.insert("a".to_string(), a);
        node_map.insert("b".to_string(), b);
        node_map.insert("c".to_string(), c);
        node_map.insert("d".to_string(), d);
        node_map.insert("e".to_string(), e);
        node_map.insert("f".to_string(), f);

        let cg = CodeGraph {
            graph,
            node_map,
            nodes: vec![
                make_node("a", 0),
                make_node("b", 1),
                make_node("c", 2),
                make_node("d", 3),
                make_node("e", 4),
                make_node("f", 5),
            ],
            edges: Vec::new(),
        };
        let num = detect_communities(&cg);
        // Two disconnected clusters should yield 2 communities
        assert_eq!(num, 2);
    }

    #[test]
    fn test_apply_communities_assigns_ids() {
        let mut graph = DiGraph::new();
        let a = graph.add_node("a".to_string());
        let b = graph.add_node("b".to_string());
        graph.add_edge(a, b, CodeEdgeKind::Calls);

        let mut node_map = HashMap::new();
        node_map.insert("a".to_string(), a);
        node_map.insert("b".to_string(), b);

        let mut cg = CodeGraph {
            graph,
            node_map,
            nodes: vec![make_node("a", 0), make_node("b", 1)],
            edges: Vec::new(),
        };
        let count = apply_communities(&mut cg);
        assert!(count >= 1);
        // All nodes should have a community_id assigned
        for node in &cg.nodes {
            assert!(node.community_id.is_some());
        }
    }

    #[test]
    fn test_apply_communities_empty() {
        let mut cg = make_empty_code_graph();
        assert_eq!(apply_communities(&mut cg), 0);
    }

    #[test]
    fn test_apply_communities_isolated_nodes_get_own_community() {
        let mut graph = DiGraph::new();
        let a = graph.add_node("a".to_string());
        let b = graph.add_node("b".to_string());
        let mut node_map = HashMap::new();
        node_map.insert("a".to_string(), a);
        node_map.insert("b".to_string(), b);

        let mut cg = CodeGraph {
            graph,
            node_map,
            nodes: vec![make_node("a", 0), make_node("b", 1)],
            edges: Vec::new(),
        };
        let count = apply_communities(&mut cg);
        assert_eq!(count, 2);
        // Each isolated node should be in a different community
        let c0 = cg.nodes[0].community_id.unwrap();
        let c1 = cg.nodes[1].community_id.unwrap();
        assert_ne!(c0, c1);
    }
}
