compliance-scanner-agent/compliance-graph/src/graph/impact.rs

use std::collections::{HashSet, VecDeque};

use compliance_core::models::graph::ImpactAnalysis;
use petgraph::graph::NodeIndex;
use petgraph::visit::EdgeRef;
use petgraph::Direction;

use super::engine::CodeGraph;

/// Analyzes the impact/blast radius of findings within a code graph
pub struct ImpactAnalyzer<'a> {
    code_graph: &'a CodeGraph,
}

impl<'a> ImpactAnalyzer<'a> {
    pub fn new(code_graph: &'a CodeGraph) -> Self {
        Self { code_graph }
    }

    /// Compute impact analysis for a finding at the given file path and line number
    pub fn analyze(
        &self,
        repo_id: &str,
        finding_id: &str,
        graph_build_id: &str,
        file_path: &str,
        line_number: Option<u32>,
    ) -> ImpactAnalysis {
        let mut analysis = ImpactAnalysis::new(
            repo_id.to_string(),
            finding_id.to_string(),
            graph_build_id.to_string(),
        );

        // Find the node containing the finding
        let target_node = self.find_node_at_location(file_path, line_number);
        let target_idx = match target_node {
            Some(idx) => idx,
            None => return analysis,
        };

        // BFS forward: compute blast radius (what this node affects)
        let forward_reachable = self.bfs_reachable(target_idx, Direction::Outgoing);
        analysis.blast_radius = forward_reachable.len() as u32;

        // BFS backward: find entry points that reach this node
        let backward_reachable = self.bfs_reachable(target_idx, Direction::Incoming);

        // Find affected entry points
        for &idx in &backward_reachable {
            if let Some(node) = self.get_node_by_index(idx) {
                if node.is_entry_point {
                    analysis
                        .affected_entry_points
                        .push(node.qualified_name.clone());
                }
            }
        }

        // Extract call chains from entry points to the target (limited depth)
        for entry_name in &analysis.affected_entry_points.clone() {
            if let Some(&entry_idx) = self.code_graph.node_map.get(entry_name) {
                if let Some(chain) = self.find_path(entry_idx, target_idx, 10) {
                    analysis.call_chains.push(chain);
                }
            }
        }

        // Direct callers (incoming edges to target)
        for edge in self
            .code_graph
            .graph
            .edges_directed(target_idx, Direction::Incoming)
        {
            if let Some(node) = self.get_node_by_index(edge.source()) {
                analysis.direct_callers.push(node.qualified_name.clone());
            }
        }

        // Direct callees (outgoing edges from target)
        for edge in self.code_graph.graph.edges(target_idx) {
            if let Some(node) = self.get_node_by_index(edge.target()) {
                analysis.direct_callees.push(node.qualified_name.clone());
            }
        }

        // Affected communities
        let mut affected_comms: HashSet<u32> = HashSet::new();
        for &idx in forward_reachable.iter().chain(std::iter::once(&target_idx)) {
            if let Some(node) = self.get_node_by_index(idx) {
                if let Some(cid) = node.community_id {
                    affected_comms.insert(cid);
                }
            }
        }
        analysis.affected_communities = affected_comms.into_iter().collect();
        analysis.affected_communities.sort();

        analysis
    }

    /// Find the graph node at a given file/line location
    fn find_node_at_location(
        &self,
        file_path: &str,
        line_number: Option<u32>,
    ) -> Option<NodeIndex> {
        let mut best: Option<(NodeIndex, u32)> = None; // (index, line_span)

        for node in &self.code_graph.nodes {
            if node.file_path != file_path {
                continue;
            }

            if let Some(line) = line_number {
                if line >= node.start_line && line <= node.end_line {
                    let span = node.end_line - node.start_line;
                    // Prefer the narrowest containing node
                    if best.is_none() || span < best.as_ref().map(|b| b.1).unwrap_or(u32::MAX) {
                        if let Some(gi) = node.graph_index {
                            best = Some((NodeIndex::new(gi as usize), span));
                        }
                    }
                }
            } else {
                // No line number, use file node
                if node.kind == compliance_core::models::graph::CodeNodeKind::File {
                    if let Some(gi) = node.graph_index {
                        return Some(NodeIndex::new(gi as usize));
                    }
                }
            }
        }

        best.map(|(idx, _)| idx)
    }

    /// BFS to find all reachable nodes in a given direction
    fn bfs_reachable(&self, start: NodeIndex, direction: Direction) -> HashSet<NodeIndex> {
        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(start);

        while let Some(current) = queue.pop_front() {
            if !visited.insert(current) {
                continue;
            }

            let neighbors: Vec<NodeIndex> = match direction {
                Direction::Outgoing => self
                    .code_graph
                    .graph
                    .edges(current)
                    .map(|e| e.target())
                    .collect(),
                Direction::Incoming => self
                    .code_graph
                    .graph
                    .edges_directed(current, Direction::Incoming)
                    .map(|e| e.source())
                    .collect(),
            };

            for neighbor in neighbors {
                if !visited.contains(&neighbor) {
                    queue.push_back(neighbor);
                }
            }
        }

        visited.remove(&start);
        visited
    }

    /// Find a path from source to target (BFS, limited depth)
    fn find_path(&self, from: NodeIndex, to: NodeIndex, max_depth: usize) -> Option<Vec<String>> {
        let mut visited = HashSet::new();
        let mut queue: VecDeque<(NodeIndex, Vec<NodeIndex>)> = VecDeque::new();
        queue.push_back((from, vec![from]));

        while let Some((current, path)) = queue.pop_front() {
            if current == to {
                return Some(
                    path.iter()
                        .filter_map(|&idx| {
                            self.get_node_by_index(idx)
                                .map(|n| n.qualified_name.clone())
                        })
                        .collect(),
                );
            }

            if path.len() >= max_depth {
                continue;
            }

            if !visited.insert(current) {
                continue;
            }

            for edge in self.code_graph.graph.edges(current) {
                let next = edge.target();
                if !visited.contains(&next) {
                    let mut new_path = path.clone();
                    new_path.push(next);
                    queue.push_back((next, new_path));
                }
            }
        }

        None
    }

    fn get_node_by_index(
        &self,
        idx: NodeIndex,
    ) -> Option<&compliance_core::models::graph::CodeNode> {
        let target_gi = idx.index() as u32;
        self.code_graph
            .nodes
            .iter()
            .find(|n| n.graph_index == Some(target_gi))
    }
}

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

    fn make_node(
        qualified_name: &str,
        file_path: &str,
        start: u32,
        end: u32,
        graph_index: u32,
        is_entry: bool,
        kind: CodeNodeKind,
    ) -> CodeNode {
        CodeNode {
            id: None,
            repo_id: "test".to_string(),
            graph_build_id: "build1".to_string(),
            qualified_name: qualified_name.to_string(),
            name: qualified_name
                .split("::")
                .last()
                .unwrap_or(qualified_name)
                .to_string(),
            kind,
            file_path: file_path.to_string(),
            start_line: start,
            end_line: end,
            language: "rust".to_string(),
            community_id: None,
            is_entry_point: is_entry,
            graph_index: Some(graph_index),
        }
    }

    fn make_fn_node(
        qualified_name: &str,
        file_path: &str,
        start: u32,
        end: u32,
        gi: u32,
    ) -> CodeNode {
        make_node(
            qualified_name,
            file_path,
            start,
            end,
            gi,
            false,
            CodeNodeKind::Function,
        )
    }

    /// Build a simple linear graph: A -> B -> C
    fn build_linear_graph() -> CodeGraph {
        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);

        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);

        CodeGraph {
            graph,
            node_map,
            nodes: vec![
                make_fn_node("a", "src/main.rs", 1, 5, 0),
                make_fn_node("b", "src/main.rs", 7, 12, 1),
                make_fn_node("c", "src/main.rs", 14, 20, 2),
            ],
            edges: Vec::new(),
        }
    }

    #[test]
    fn test_bfs_reachable_outgoing_linear() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let start = cg.node_map["a"];
        let reachable = analyzer.bfs_reachable(start, Direction::Outgoing);
        // From a, we can reach b and c
        assert_eq!(reachable.len(), 2);
        assert!(reachable.contains(&cg.node_map["b"]));
        assert!(reachable.contains(&cg.node_map["c"]));
    }

    #[test]
    fn test_bfs_reachable_incoming_linear() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let start = cg.node_map["c"];
        let reachable = analyzer.bfs_reachable(start, Direction::Incoming);
        // c is reached by a and b
        assert_eq!(reachable.len(), 2);
        assert!(reachable.contains(&cg.node_map["a"]));
        assert!(reachable.contains(&cg.node_map["b"]));
    }

    #[test]
    fn test_bfs_reachable_no_neighbors() {
        let mut graph = DiGraph::new();
        let a = graph.add_node("a".to_string());
        let cg = CodeGraph {
            graph,
            node_map: [("a".to_string(), a)].into_iter().collect(),
            nodes: vec![make_fn_node("a", "src/main.rs", 1, 5, 0)],
            edges: Vec::new(),
        };
        let analyzer = ImpactAnalyzer::new(&cg);
        let reachable = analyzer.bfs_reachable(a, Direction::Outgoing);
        assert!(reachable.is_empty());
    }

    #[test]
    fn test_bfs_reachable_cycle() {
        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);
        graph.add_edge(b, a, CodeEdgeKind::Calls);

        let cg = CodeGraph {
            graph,
            node_map: [("a".to_string(), a), ("b".to_string(), b)]
                .into_iter()
                .collect(),
            nodes: vec![
                make_fn_node("a", "f.rs", 1, 5, 0),
                make_fn_node("b", "f.rs", 6, 10, 1),
            ],
            edges: Vec::new(),
        };
        let analyzer = ImpactAnalyzer::new(&cg);
        let reachable = analyzer.bfs_reachable(a, Direction::Outgoing);
        // Should handle cycle without infinite loop
        assert_eq!(reachable.len(), 1);
        assert!(reachable.contains(&b));
    }

    #[test]
    fn test_find_path_exists() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let path = analyzer.find_path(cg.node_map["a"], cg.node_map["c"], 10);
        assert!(path.is_some());
        let names = path.unwrap();
        assert_eq!(names, vec!["a", "b", "c"]);
    }

    #[test]
    fn test_find_path_direct() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let path = analyzer.find_path(cg.node_map["a"], cg.node_map["b"], 10);
        assert!(path.is_some());
        let names = path.unwrap();
        assert_eq!(names, vec!["a", "b"]);
    }

    #[test]
    fn test_find_path_same_node() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let path = analyzer.find_path(cg.node_map["a"], cg.node_map["a"], 10);
        assert!(path.is_some());
        let names = path.unwrap();
        assert_eq!(names, vec!["a"]);
    }

    #[test]
    fn test_find_path_no_connection() {
        let mut graph = DiGraph::new();
        let a = graph.add_node("a".to_string());
        let b = graph.add_node("b".to_string());
        // No edge between a and b
        let cg = CodeGraph {
            graph,
            node_map: [("a".to_string(), a), ("b".to_string(), b)]
                .into_iter()
                .collect(),
            nodes: vec![
                make_fn_node("a", "f.rs", 1, 5, 0),
                make_fn_node("b", "f.rs", 6, 10, 1),
            ],
            edges: Vec::new(),
        };
        let analyzer = ImpactAnalyzer::new(&cg);
        let path = analyzer.find_path(a, b, 10);
        assert!(path.is_none());
    }

    #[test]
    fn test_find_path_depth_limited() {
        // Build a long chain: a -> b -> c -> d -> e
        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());
        graph.add_edge(a, b, CodeEdgeKind::Calls);
        graph.add_edge(b, c, CodeEdgeKind::Calls);
        graph.add_edge(c, d, CodeEdgeKind::Calls);
        graph.add_edge(d, e, 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);

        let cg = CodeGraph {
            graph,
            node_map,
            nodes: vec![
                make_fn_node("a", "f.rs", 1, 2, 0),
                make_fn_node("b", "f.rs", 3, 4, 1),
                make_fn_node("c", "f.rs", 5, 6, 2),
                make_fn_node("d", "f.rs", 7, 8, 3),
                make_fn_node("e", "f.rs", 9, 10, 4),
            ],
            edges: Vec::new(),
        };
        let analyzer = ImpactAnalyzer::new(&cg);
        // Depth 3 won't reach e from a (path length 5)
        let path = analyzer.find_path(a, e, 3);
        assert!(path.is_none());
        // Depth 5 should reach
        let path = analyzer.find_path(a, e, 5);
        assert!(path.is_some());
    }

    #[test]
    fn test_find_node_at_location_exact_line() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        // Node "b" is at lines 7-12
        let result = analyzer.find_node_at_location("src/main.rs", Some(9));
        assert!(result.is_some());
        assert_eq!(result.unwrap(), cg.node_map["b"]);
    }

    #[test]
    fn test_find_node_at_location_narrowest_match() {
        // Outer function 1-20, inner nested 5-10
        let mut graph = DiGraph::new();
        let outer = graph.add_node("outer".to_string());
        let inner = graph.add_node("inner".to_string());

        let cg = CodeGraph {
            graph,
            node_map: [("outer".to_string(), outer), ("inner".to_string(), inner)]
                .into_iter()
                .collect(),
            nodes: vec![
                make_fn_node("outer", "src/main.rs", 1, 20, 0),
                make_fn_node("inner", "src/main.rs", 5, 10, 1),
            ],
            edges: Vec::new(),
        };
        let analyzer = ImpactAnalyzer::new(&cg);
        // Line 7 is inside both, but inner is narrower
        let result = analyzer.find_node_at_location("src/main.rs", Some(7));
        assert!(result.is_some());
        assert_eq!(result.unwrap(), inner);
    }

    #[test]
    fn test_find_node_at_location_no_line_returns_file_node() {
        let mut graph = DiGraph::new();
        let file_node = graph.add_node("src/main.rs".to_string());
        let fn_node = graph.add_node("src/main.rs::foo".to_string());

        let cg = CodeGraph {
            graph,
            node_map: [
                ("src/main.rs".to_string(), file_node),
                ("src/main.rs::foo".to_string(), fn_node),
            ]
            .into_iter()
            .collect(),
            nodes: vec![
                make_node(
                    "src/main.rs",
                    "src/main.rs",
                    1,
                    100,
                    0,
                    false,
                    CodeNodeKind::File,
                ),
                make_fn_node("src/main.rs::foo", "src/main.rs", 5, 10, 1),
            ],
            edges: Vec::new(),
        };
        let analyzer = ImpactAnalyzer::new(&cg);
        let result = analyzer.find_node_at_location("src/main.rs", None);
        assert!(result.is_some());
        assert_eq!(result.unwrap(), file_node);
    }

    #[test]
    fn test_find_node_at_location_wrong_file() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let result = analyzer.find_node_at_location("nonexistent.rs", Some(5));
        assert!(result.is_none());
    }

    #[test]
    fn test_find_node_at_location_line_out_of_range() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let result = analyzer.find_node_at_location("src/main.rs", Some(999));
        assert!(result.is_none());
    }

    #[test]
    fn test_analyze_basic() {
        // A (entry) -> B -> C
        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);

        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", "src/main.rs", 1, 5, 0, true, CodeNodeKind::Function),
                make_fn_node("b", "src/main.rs", 7, 12, 1),
                make_fn_node("c", "src/main.rs", 14, 20, 2),
            ],
            edges: Vec::new(),
        };

        let analyzer = ImpactAnalyzer::new(&cg);
        let result = analyzer.analyze("repo1", "finding1", "build1", "src/main.rs", Some(9));
        // B's blast radius: C is reachable forward
        assert_eq!(result.blast_radius, 1);
        // B has A as direct caller
        assert_eq!(result.direct_callers, vec!["a"]);
        // B calls C
        assert_eq!(result.direct_callees, vec!["c"]);
        // A is an entry point that reaches B
        assert_eq!(result.affected_entry_points, vec!["a"]);
    }

    #[test]
    fn test_analyze_no_matching_node() {
        let cg = build_linear_graph();
        let analyzer = ImpactAnalyzer::new(&cg);
        let result = analyzer.analyze("repo1", "f1", "b1", "nonexistent.rs", Some(1));
        assert_eq!(result.blast_radius, 0);
        assert!(result.affected_entry_points.is_empty());
        assert!(result.direct_callers.is_empty());
        assert!(result.direct_callees.is_empty());
    }
}