lore_engine/engine/

graph.rs

//! In-memory directed graph of wiki pages and their links.
//!
//! Backed by `petgraph::DiGraph`. Nodes are pages (real or placeholder),
//! edges are wiki links. Used for graph visualization, backlink queries,
//! and orphan cleanup.

use petgraph::graph::{DiGraph, NodeIndex};
use petgraph::Direction;
use serde::Serialize;
use std::collections::{HashMap, HashSet};

/// A node in the wiki graph.
#[derive(Debug, Serialize, Clone)]
pub struct GraphNode {
    pub slug: String,
    pub title: String,
    pub is_placeholder: bool,
}

/// A directed edge in the wiki graph (source page links to target page).
#[derive(Debug, Serialize, Clone)]
pub struct GraphEdge {
    pub source: String,
    pub target: String,
}

/// In-memory directed graph of wiki pages and their `[[wikilink]]` connections.
pub struct WikiGraph {
    graph: DiGraph<GraphNode, ()>,
    index_map: HashMap<String, NodeIndex>,
}

impl Default for WikiGraph {
    fn default() -> Self {
        Self::new()
    }
}

impl WikiGraph {
    pub fn new() -> Self {
        Self {
            graph: DiGraph::new(),
            index_map: HashMap::new(),
        }
    }

    /// Add or update a node. Returns its `NodeIndex`.
    pub fn upsert_node(&mut self, slug: &str, title: &str, is_placeholder: bool) -> NodeIndex {
        if let Some(&idx) = self.index_map.get(slug) {
            if let Some(node) = self.graph.node_weight_mut(idx) {
                node.title = title.to_string();
                node.is_placeholder = is_placeholder;
            }
            idx
        } else {
            let idx = self.graph.add_node(GraphNode {
                slug: slug.to_string(),
                title: title.to_string(),
                is_placeholder,
            });
            self.index_map.insert(slug.to_string(), idx);
            idx
        }
    }

    /// Remove a node and all its edges.
    pub fn remove_node(&mut self, slug: &str) {
        if let Some(idx) = self.index_map.remove(slug) {
            self.graph.remove_node(idx);
            // petgraph may reuse indices after removal, so rebuild the index map
            self.rebuild_index_map();
        }
    }

    fn rebuild_index_map(&mut self) {
        self.index_map.clear();
        for idx in self.graph.node_indices() {
            if let Some(node) = self.graph.node_weight(idx) {
                self.index_map.insert(node.slug.clone(), idx);
            }
        }
    }

    /// Set outgoing edges FROM a source node.
    ///
    /// Removes all existing outgoing edges first, then adds new ones.
    /// Creates placeholder nodes for targets that don't exist.
    /// Duplicate targets are deduplicated — at most one edge per (source, target) pair.
    pub fn set_outgoing_edges(&mut self, source_slug: &str, target_slugs: &[String]) {
        let source_idx = match self.index_map.get(source_slug) {
            Some(&idx) => idx,
            None => return,
        };

        // Remove existing outgoing edges
        let outgoing: Vec<_> = self
            .graph
            .neighbors_directed(source_idx, Direction::Outgoing)
            .collect();
        for target_idx in outgoing {
            if let Some(edge) = self.graph.find_edge(source_idx, target_idx) {
                self.graph.remove_edge(edge);
            }
        }

        // Deduplicate targets — page-level graph, not occurrence graph
        let mut seen = HashSet::new();
        for target_slug in target_slugs {
            if !seen.insert(target_slug.as_str()) {
                continue;
            }

            let target_idx = if let Some(&idx) = self.index_map.get(target_slug.as_str()) {
                idx
            } else {
                self.upsert_node(target_slug, target_slug, true)
            };

            self.graph.add_edge(source_idx, target_idx, ());
        }
    }

    /// Get all slugs that link TO a given slug (backlinks / incoming).
    /// Results are sorted by slug for deterministic output.
    pub fn get_backlinks(&self, slug: &str) -> Vec<String> {
        let idx = match self.index_map.get(slug) {
            Some(&idx) => idx,
            None => return Vec::new(),
        };

        let mut result: Vec<String> = self.graph
            .neighbors_directed(idx, Direction::Incoming)
            .filter_map(|n| self.graph.node_weight(n).map(|w| w.slug.clone()))
            .collect();
        result.sort();
        result
    }

    /// Get all slugs that a given slug links TO (forward links / outgoing).
    /// Results are sorted by slug for deterministic output.
    pub fn get_forward_links(&self, slug: &str) -> Vec<String> {
        let idx = match self.index_map.get(slug) {
            Some(&idx) => idx,
            None => return Vec::new(),
        };

        let mut result: Vec<String> = self.graph
            .neighbors_directed(idx, Direction::Outgoing)
            .filter_map(|n| self.graph.node_weight(n).map(|w| w.slug.clone()))
            .collect();
        result.sort();
        result
    }

    /// Get the degree (total edges in + out) of a node.
    pub fn degree(&self, slug: &str) -> usize {
        let idx = match self.index_map.get(slug) {
            Some(&idx) => idx,
            None => return 0,
        };

        self.graph
            .neighbors_directed(idx, Direction::Incoming)
            .count()
            + self
                .graph
                .neighbors_directed(idx, Direction::Outgoing)
                .count()
    }

    /// Remove placeholder nodes that have zero edges.
    /// Returns the slugs of removed nodes.
    ///
    /// Uses batched removal with a single index rebuild at the end,
    /// avoiding the O(n * k) cost of rebuilding after each removal.
    pub fn cleanup_orphaned_placeholders(&mut self) -> Vec<String> {
        // Collect orphans using direct index access (no map lookup)
        let orphan_indices: Vec<(NodeIndex, String)> = self
            .graph
            .node_indices()
            .filter_map(|idx| {
                let node = self.graph.node_weight(idx)?;
                if node.is_placeholder {
                    let incoming = self.graph.neighbors_directed(idx, Direction::Incoming).count();
                    let outgoing = self.graph.neighbors_directed(idx, Direction::Outgoing).count();
                    if incoming + outgoing == 0 {
                        return Some((idx, node.slug.clone()));
                    }
                }
                None
            })
            .collect();

        let slugs: Vec<String> = orphan_indices.iter().map(|(_, s)| s.clone()).collect();

        // Remove nodes in reverse index order to avoid invalidating earlier indices
        // (petgraph::DiGraph::remove_node swaps with the last node)
        let mut indices: Vec<NodeIndex> = orphan_indices.into_iter().map(|(idx, _)| idx).collect();
        indices.sort_by_key(|idx| std::cmp::Reverse(idx.index()));

        for idx in indices {
            self.index_map.remove(&self.graph[idx].slug.clone());
            self.graph.remove_node(idx);
        }

        // Single rebuild after all removals
        if !slugs.is_empty() {
            self.rebuild_index_map();
        }

        slugs
    }

    /// Get full graph data for serialization.
    /// Nodes sorted by slug, edges sorted by (source, target) for deterministic output.
    pub fn get_full_graph(&self) -> (Vec<GraphNode>, Vec<GraphEdge>) {
        let mut nodes: Vec<GraphNode> = self
            .graph
            .node_indices()
            .filter_map(|idx| self.graph.node_weight(idx).cloned())
            .collect();
        nodes.sort_by(|a, b| a.slug.cmp(&b.slug));

        let mut edges: Vec<GraphEdge> = self
            .graph
            .edge_indices()
            .filter_map(|eidx| {
                let (src, tgt) = self.graph.edge_endpoints(eidx)?;
                let src_slug = self.graph.node_weight(src)?.slug.clone();
                let tgt_slug = self.graph.node_weight(tgt)?.slug.clone();
                Some(GraphEdge {
                    source: src_slug,
                    target: tgt_slug,
                })
            })
            .collect();
        edges.sort_by(|a, b| (&a.source, &a.target).cmp(&(&b.source, &b.target)));

        (nodes, edges)
    }

    /// Get node and edge counts.
    pub fn stats(&self) -> (usize, usize) {
        (self.graph.node_count(), self.graph.edge_count())
    }
}

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

    #[test]
    fn test_add_nodes_and_edges() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "Page A", false);
        g.upsert_node("b", "Page B", false);
        g.upsert_node("c", "Page C", false);

        g.set_outgoing_edges("a", &["b".into(), "c".into()]);

        assert_eq!(g.get_forward_links("a").len(), 2);
        assert_eq!(g.stats(), (3, 2));
    }

    #[test]
    fn test_backlinks() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "Page A", false);
        g.upsert_node("b", "Page B", false);
        g.upsert_node("c", "Page C", false);

        g.set_outgoing_edges("a", &["b".into(), "c".into()]);

        let backlinks_b = g.get_backlinks("b");
        assert_eq!(backlinks_b, vec!["a".to_string()]);

        let backlinks_c = g.get_backlinks("c");
        assert_eq!(backlinks_c, vec!["a".to_string()]);

        assert!(g.get_backlinks("a").is_empty());
    }

    #[test]
    fn test_remove_node_cleans_edges() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "A", false);
        g.upsert_node("b", "B", false);
        g.set_outgoing_edges("a", &["b".into()]);

        assert_eq!(g.stats(), (2, 1));

        g.remove_node("a");
        assert_eq!(g.stats(), (1, 0));
        assert!(g.get_backlinks("b").is_empty());
    }

    #[test]
    fn test_set_outgoing_replaces() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "A", false);
        g.upsert_node("b", "B", false);
        g.upsert_node("c", "C", false);

        g.set_outgoing_edges("a", &["b".into(), "c".into()]);
        assert_eq!(g.get_forward_links("a").len(), 2);

        g.set_outgoing_edges("a", &["c".into()]);
        let fwd = g.get_forward_links("a");
        assert_eq!(fwd.len(), 1);
        assert_eq!(fwd[0], "c");
        assert!(g.get_backlinks("b").is_empty());
    }

    #[test]
    fn test_placeholder_creation() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "A", false);

        g.set_outgoing_edges("a", &["d".into()]);

        let (nodes, _) = g.get_full_graph();
        let d_node = nodes.iter().find(|n| n.slug == "d");
        assert!(d_node.is_some());
        assert!(d_node.expect("d exists").is_placeholder);
    }

    #[test]
    fn test_cleanup_orphaned_placeholders() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "A", false);
        g.upsert_node("b", "B", true);
        g.set_outgoing_edges("a", &["b".into()]);
        assert_eq!(g.stats(), (2, 1));

        g.set_outgoing_edges("a", &[]);
        assert_eq!(g.stats(), (2, 0));

        let orphans = g.cleanup_orphaned_placeholders();
        assert_eq!(orphans, vec!["b".to_string()]);
        assert_eq!(g.stats(), (1, 0));

        g.upsert_node("c", "C", false);
        let orphans2 = g.cleanup_orphaned_placeholders();
        assert!(orphans2.is_empty());
        assert_eq!(g.stats(), (2, 0));
    }

    #[test]
    fn test_degree() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "A", false);
        g.upsert_node("b", "B", false);
        g.upsert_node("c", "C", false);

        g.set_outgoing_edges("a", &["b".into(), "c".into()]);
        g.set_outgoing_edges("b", &["c".into()]);

        assert_eq!(g.degree("a"), 2);
        assert_eq!(g.degree("b"), 2);
        assert_eq!(g.degree("c"), 2);
    }

    #[test]
    fn test_duplicate_targets_deduplicated() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "A", false);
        g.upsert_node("b", "B", false);

        // Same target three times — should produce one edge, not three
        g.set_outgoing_edges("a", &["b".into(), "b".into(), "b".into()]);

        assert_eq!(g.stats(), (2, 1));
        assert_eq!(g.degree("a"), 1);
        assert_eq!(g.degree("b"), 1);
    }

    #[test]
    fn test_backlinks_sorted() {
        let mut g = WikiGraph::new();
        g.upsert_node("target", "Target", false);
        g.upsert_node("charlie", "Charlie", false);
        g.upsert_node("alice", "Alice", false);
        g.upsert_node("bob", "Bob", false);

        g.set_outgoing_edges("charlie", &["target".into()]);
        g.set_outgoing_edges("alice", &["target".into()]);
        g.set_outgoing_edges("bob", &["target".into()]);

        let backlinks = g.get_backlinks("target");
        assert_eq!(backlinks, vec!["alice", "bob", "charlie"]);
    }

    #[test]
    fn test_forward_links_sorted() {
        let mut g = WikiGraph::new();
        g.upsert_node("a", "A", false);
        g.upsert_node("z", "Z", false);
        g.upsert_node("m", "M", false);

        g.set_outgoing_edges("a", &["z".into(), "m".into()]);

        let forward = g.get_forward_links("a");
        assert_eq!(forward, vec!["m", "z"]);
    }

    #[test]
    fn test_full_graph_sorted() {
        let mut g = WikiGraph::new();
        g.upsert_node("c", "C", false);
        g.upsert_node("a", "A", false);
        g.upsert_node("b", "B", false);

        g.set_outgoing_edges("c", &["a".into()]);
        g.set_outgoing_edges("a", &["b".into()]);

        let (nodes, edges) = g.get_full_graph();
        let slugs: Vec<&str> = nodes.iter().map(|n| n.slug.as_str()).collect();
        assert_eq!(slugs, vec!["a", "b", "c"]);

        assert_eq!(edges[0].source, "a");
        assert_eq!(edges[0].target, "b");
        assert_eq!(edges[1].source, "c");
        assert_eq!(edges[1].target, "a");
    }

    #[test]
    fn test_batch_orphan_cleanup() {
        let mut g = WikiGraph::new();
        g.upsert_node("real", "Real", false);

        // Create 10 placeholders linked from real
        let targets: Vec<String> = (0..10).map(|i| format!("ph-{i}")).collect();
        g.set_outgoing_edges("real", &targets);
        assert_eq!(g.stats(), (11, 10));

        // Remove all outgoing edges — placeholders become orphans
        g.set_outgoing_edges("real", &[]);
        assert_eq!(g.stats(), (11, 0));

        let orphans = g.cleanup_orphaned_placeholders();
        assert_eq!(orphans.len(), 10);
        assert_eq!(g.stats(), (1, 0));
    }

    #[test]
    fn test_set_outgoing_edges_missing_source_is_noop() {
        let mut g = WikiGraph::new();
        // Source doesn't exist — should not panic or create nodes
        g.set_outgoing_edges("nonexistent", &["a".into()]);
        assert_eq!(g.stats(), (0, 0));
    }
}