poimen-memory/crates/mem-cli/src/query/bfs_graph_traversal.rs

/// BFS graph traversal with PostgreSQL queries.
///
/// Performs breadth-first search on memory_entity + memory_edge tables,
/// returning a subgraph for visualization.

use std::collections::{HashMap, VecDeque};
use serde::{Deserialize, Serialize};
use chrono::{DateTime, Utc};
use sqlx::{Pool, Postgres, Row};

/// A node in the traversal result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraversalNode {
    pub id: String,
    pub entity_type: String,
    pub name: String,
    pub description: Option<String>,
    pub depth: i32,  // Distance from root (0 = root)
}

/// An edge in the traversal result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraversalEdge {
    pub id: String,
    pub source_id: String,
    pub target_id: String,
    pub relation_type: String,
    pub fact: String,
    pub strength: f32,
}

/// Depth-level breakdown
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DepthBreakdown {
    pub depth: i32,
    pub node_count: usize,
    pub edge_count: usize,
}

/// Graph data from BFS traversal
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphData {
    pub nodes: Vec<TraversalNode>,
    pub edges: Vec<TraversalEdge>,
    pub root_id: String,
    pub requested_depth: i32,           // Depth that was requested
    pub max_depth_reached: i32,         // Actual max depth in result
    pub node_count: usize,
    pub edge_count: usize,
    pub depth_breakdown: Vec<DepthBreakdown>,  // Nodes/edges per depth level
    pub traversal_time_ms: u64,
}

/// BFS traversal configuration
#[derive(Debug, Clone)]
pub struct BfsConfig {
    pub max_depth: i32,           // Max hops from root (1-3)
    pub max_nodes: usize,         // Max nodes to return (default 50)
    pub max_edges_per_node: usize, // Max edges per node (to avoid explosion)
}

impl Default for BfsConfig {
    fn default() -> Self {
        Self {
            max_depth: 2,
            max_nodes: 50,
            max_edges_per_node: 5,
        }
    }
}

/// BFS graph traversal engine
pub struct BfsGraphTraversal {
    pool: Pool<Postgres>,
}

impl BfsGraphTraversal {
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }
    
    /// Traverse graph starting from root entity
    pub async fn traverse(
        &self,
        root_id: &str,
        config: &BfsConfig,
    ) -> Result<GraphData, String> {
        let start_time = std::time::Instant::now();
        
        // 1. Load root entity
        let root = self.load_entity(root_id).await?;
        if root.is_none() {
            return Err(format!("Root entity not found: {}", root_id));
        }
        let root_node = root.unwrap();
        
        // 2. BFS traversal
        let mut nodes = vec![TraversalNode {
            id: root_node.0,
            entity_type: root_node.1,
            name: root_node.2,
            description: root_node.3,
            depth: 0,
        }];
        
        let mut edges = Vec::new();
        let mut visited = std::collections::HashSet::new();
        visited.insert(root_id.to_string());
        
        let mut queue = VecDeque::new();
        queue.push_back((root_id.to_string(), 0));
        
        while let Some((current_id, current_depth)) = queue.pop_front() {
            // Stop if we've reached max depth
            if current_depth >= config.max_depth {
                continue;
            }
            
            // Stop if we've reached max nodes
            if nodes.len() >= config.max_nodes {
                break;
            }
            
            // Load outgoing edges from current node (sampled)
            let out_edges = self.load_edges_from(&current_id, config.max_edges_per_node).await?;
            
            for edge in out_edges {
                let target_id = &edge.1;
                
                // Skip if already visited
                if visited.contains(target_id) {
                    // But still add the edge (creates a cycle in the graph)
                    edges.push(TraversalEdge {
                        id: edge.0,
                        source_id: edge.2.clone(),
                        target_id: target_id.clone(),
                        relation_type: edge.3,
                        fact: edge.4,
                        strength: edge.5,
                    });
                    continue;
                }
                
                // Load target entity
                if let Ok(target_opt) = self.load_entity(target_id).await {
                    if let Some(target) = target_opt {
                        // Add node to result
                        nodes.push(TraversalNode {
                            id: target.0.clone(),
                            entity_type: target.1,
                            name: target.2,
                            description: target.3,
                            depth: current_depth + 1,
                        });
                        
                        // Mark as visited
                        visited.insert(target.0);
                        
                        // Add to queue for next iteration
                        queue.push_back((target_id.clone(), current_depth + 1));
                    }
                }
                
                // Add edge
                edges.push(TraversalEdge {
                    id: edge.0,
                    source_id: edge.2,
                    target_id: target_id.clone(),
                    relation_type: edge.3,
                    fact: edge.4,
                    strength: edge.5,
                });
            }
        }
        
        let max_depth = nodes.iter().map(|n| n.depth).max().unwrap_or(0);
        
        // Compute depth breakdown
        let mut depth_breakdown = Vec::new();
        for depth in 0..=max_depth {
            let nodes_at_depth = nodes.iter().filter(|n| n.depth == depth).count();
            let edges_from_depth = edges.iter()
                .filter(|e| {
                    let source_depth = nodes.iter()
                        .find(|n| n.id == e.source_id)
                        .map(|n| n.depth)
                        .unwrap_or(0);
                    source_depth == depth
                })
                .count();
            
            depth_breakdown.push(DepthBreakdown {
                depth,
                node_count: nodes_at_depth,
                edge_count: edges_from_depth,
            });
        }
        
        let edge_count = edges.len();
        Ok(GraphData {
            nodes,
            edges,
            root_id: root_id.to_string(),
            requested_depth: config.max_depth,
            max_depth_reached: max_depth,
            node_count: visited.len(),
            edge_count,
            depth_breakdown,
            traversal_time_ms: start_time.elapsed().as_millis() as u64,
        })
    }
    
    /// Load single entity from DB
    /// Returns: (id, entity_type, name, description)
    async fn load_entity(&self, id: &str) -> Result<Option<(String, String, String, Option<String>)>, String> {
        let query = r#"
            SELECT id, entity_type, name, description
            FROM memory_entity
            WHERE id = $1 AND deleted_at IS NULL
            LIMIT 1;
        "#;
        
        let row = sqlx::query(query)
            .bind(id)
            .fetch_optional(&self.pool)
            .await
            .map_err(|e| format!("Entity query failed: {}", e))?;
        
        Ok(row.map(|r| (
            r.get::<String, _>("id"),
            r.get::<String, _>("entity_type"),
            r.get::<String, _>("name"),
            r.get::<Option<String>, _>("description"),
        )))
    }
    
    /// Load outgoing edges from entity (sampled)
    /// Returns: (edge_id, target_id, source_id, relation_type, fact, strength)
    async fn load_edges_from(&self, source_id: &str, limit: usize) -> Result<Vec<(String, String, String, String, String, f32)>, String> {
        let query = r#"
            SELECT id, target_id, source_id, relation_type, fact, strength
            FROM memory_edge
            WHERE source_id = $1 AND t_expired IS NULL AND t_invalid IS NULL
            ORDER BY strength DESC
            LIMIT $2;
        "#;
        
        let rows = sqlx::query(query)
            .bind(source_id)
            .bind(limit as i64)
            .fetch_all(&self.pool)
            .await
            .map_err(|e| format!("Edge query failed: {}", e))?;
        
        Ok(rows.iter().map(|r| (
            r.get::<String, _>("id"),
            r.get::<String, _>("target_id"),
            r.get::<String, _>("source_id"),
            r.get::<String, _>("relation_type"),
            r.get::<String, _>("fact"),
            r.get::<f32, _>("strength"),
        )).collect())
    }
    
    /// Get nodes at a specific depth from traversal result
    pub fn nodes_at_depth(graph: &GraphData, depth: i32) -> Vec<&TraversalNode> {
        graph.nodes.iter()
            .filter(|n| n.depth == depth)
            .collect()
    }
    
    /// Get edges from nodes at a specific depth
    pub fn edges_from_depth(graph: &GraphData, depth: i32) -> Vec<&TraversalEdge> {
        let nodes_at_depth: std::collections::HashSet<_> = graph.nodes.iter()
            .filter(|n| n.depth == depth)
            .map(|n| n.id.as_str())
            .collect();
        
        graph.edges.iter()
            .filter(|e| nodes_at_depth.contains(e.source_id.as_str()))
            .collect()
    }
    
    /// Traverse to a specific depth only (filter out deeper results)
    pub fn truncate_to_depth(graph: &mut GraphData, max_depth: i32) {
        graph.nodes.retain(|n| n.depth <= max_depth);
        graph.edges.retain(|e| {
            let source_depth = graph.nodes.iter()
                .find(|n| n.id == e.source_id)
                .map(|n| n.depth)
                .unwrap_or(i32::MAX);
            source_depth <= max_depth
        });
        
        graph.max_depth_reached = graph.max_depth_reached.min(max_depth);
        
        // Recalculate breakdown
        let mut depth_breakdown = Vec::new();
        for depth in 0..=graph.max_depth_reached {
            let nodes_at_depth = graph.nodes.iter().filter(|n| n.depth == depth).count();
            let edges_from_depth = graph.edges.iter()
                .filter(|e| {
                    let source_depth = graph.nodes.iter()
                        .find(|n| n.id == e.source_id)
                        .map(|n| n.depth)
                        .unwrap_or(0);
                    source_depth == depth
                })
                .count();
            
            depth_breakdown.push(DepthBreakdown {
                depth,
                node_count: nodes_at_depth,
                edge_count: edges_from_depth,
            });
        }
        graph.depth_breakdown = depth_breakdown;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_bfs_config_defaults() {
        let config = BfsConfig::default();
        assert_eq!(config.max_depth, 2);
        assert_eq!(config.max_nodes, 50);
        assert_eq!(config.max_edges_per_node, 5);
    }
    
    #[test]
    fn test_traversal_node_creation() {
        let node = TraversalNode {
            id: "entity-1".to_string(),
            entity_type: "person".to_string(),
            name: "Alice".to_string(),
            description: Some("A person".to_string()),
            depth: 0,
        };
        
        assert_eq!(node.depth, 0);
        assert_eq!(node.entity_type, "person");
    }
    
    #[test]
    fn test_traversal_edge_creation() {
        let edge = TraversalEdge {
            id: "edge-1".to_string(),
            source_id: "entity-1".to_string(),
            target_id: "entity-2".to_string(),
            relation_type: "knows".to_string(),
            fact: "Alice knows Bob".to_string(),
            strength: 0.95,
        };
        
        assert_eq!(edge.source_id, "entity-1");
        assert_eq!(edge.strength, 0.95);
    }
    
    #[test]
    fn test_graph_data_creation() {
        let graph = GraphData {
            nodes: vec![],
            edges: vec![],
            root_id: "entity-1".to_string(),
            requested_depth: 2,
            max_depth_reached: 0,
            node_count: 0,
            edge_count: 0,
            depth_breakdown: vec![],
            traversal_time_ms: 100,
        };
        
        assert_eq!(graph.traversal_time_ms, 100);
    }
    
    #[test]
    fn test_nodes_at_depth() {
        let nodes = vec![
            TraversalNode {
                id: "n1".to_string(),
                entity_type: "person".to_string(),
                name: "Alice".to_string(),
                description: None,
                depth: 0,
            },
            TraversalNode {
                id: "n2".to_string(),
                entity_type: "person".to_string(),
                name: "Bob".to_string(),
                description: None,
                depth: 1,
            },
            TraversalNode {
                id: "n3".to_string(),
                entity_type: "person".to_string(),
                name: "Charlie".to_string(),
                description: None,
                depth: 1,
            },
        ];
        
        let graph = GraphData {
            nodes,
            edges: vec![],
            root_id: "n1".to_string(),
            requested_depth: 2,
            max_depth_reached: 1,
            node_count: 3,
            edge_count: 0,
            depth_breakdown: vec![],
            traversal_time_ms: 100,
        };
        
        let depth_1_nodes = BfsGraphTraversal::nodes_at_depth(&graph, 1);
        assert_eq!(depth_1_nodes.len(), 2);
        
        let depth_0_nodes = BfsGraphTraversal::nodes_at_depth(&graph, 0);
        assert_eq!(depth_0_nodes.len(), 1);
    }
    
    #[test]
    fn test_truncate_to_depth() {
        let nodes = vec![
            TraversalNode { id: "n1".to_string(), entity_type: "person".to_string(), name: "A".to_string(), description: None, depth: 0 },
            TraversalNode { id: "n2".to_string(), entity_type: "person".to_string(), name: "B".to_string(), description: None, depth: 1 },
            TraversalNode { id: "n3".to_string(), entity_type: "person".to_string(), name: "C".to_string(), description: None, depth: 2 },
        ];
        
        let edges = vec![
            TraversalEdge { id: "e1".to_string(), source_id: "n1".to_string(), target_id: "n2".to_string(), relation_type: "knows".to_string(), fact: "A knows B".to_string(), strength: 0.9 },
            TraversalEdge { id: "e2".to_string(), source_id: "n2".to_string(), target_id: "n3".to_string(), relation_type: "knows".to_string(), fact: "B knows C".to_string(), strength: 0.8 },
        ];
        
        let mut graph = GraphData {
            nodes,
            edges,
            root_id: "n1".to_string(),
            requested_depth: 2,
            max_depth_reached: 2,
            node_count: 3,
            edge_count: 2,
            depth_breakdown: vec![],
            traversal_time_ms: 100,
        };
        
        BfsGraphTraversal::truncate_to_depth(&mut graph, 1);
        
        assert_eq!(graph.nodes.len(), 2);  // Only n1 and n2
        assert_eq!(graph.edges.len(), 1);  // Only e1
        assert_eq!(graph.max_depth_reached, 1);
    }
    
    #[test]
    fn test_depth_breakdown() {
        let breakdown = DepthBreakdown {
            depth: 1,
            node_count: 5,
            edge_count: 8,
        };
        
        assert_eq!(breakdown.depth, 1);
        assert_eq!(breakdown.node_count, 5);
        assert_eq!(breakdown.edge_count, 8);
    }
    
    #[test]
    fn test_edges_from_depth() {
        let nodes = vec![
            TraversalNode { id: "n1".to_string(), entity_type: "person".to_string(), name: "A".to_string(), description: None, depth: 0 },
            TraversalNode { id: "n2".to_string(), entity_type: "person".to_string(), name: "B".to_string(), description: None, depth: 1 },
        ];
        
        let edges = vec![
            TraversalEdge { id: "e1".to_string(), source_id: "n1".to_string(), target_id: "n2".to_string(), relation_type: "knows".to_string(), fact: "knows".to_string(), strength: 0.9 },
            TraversalEdge { id: "e2".to_string(), source_id: "n2".to_string(), target_id: "n1".to_string(), relation_type: "knows".to_string(), fact: "knows".to_string(), strength: 0.8 },
        ];
        
        let graph = GraphData {
            nodes,
            edges,
            root_id: "n1".to_string(),
            requested_depth: 2,
            max_depth_reached: 1,
            node_count: 2,
            edge_count: 2,
            depth_breakdown: vec![],
            traversal_time_ms: 100,
        };
        
        let depth_0_edges = BfsGraphTraversal::edges_from_depth(&graph, 0);
        assert_eq!(depth_0_edges.len(), 1);  // Only e1 from n1 (depth 0)
    }
}