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

/// Knowledge graph visualization and traversal.
///
/// Enables users to:
/// 1. Query graph structure (BFS traversal)
/// 2. Understand depth impact (how many hops?)
/// 3. Benchmark pagination (latency per page)
/// 4. Get recommendations (tuning suggestions)
///
/// Used for iterative query refinement before production deployment.

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

use crate::query::pagination::{PaginationParams, PaginationMeta};

/// Request to visualize graph around a query.
#[derive(Clone, Debug, Deserialize)]
pub struct VisualizeRequest {
    pub project: String,
    pub query: String,
    pub depth: Option<usize>,        // 1-3, default 2
    pub limit: Option<usize>,        // Nodes per page, default 50
    pub page: Option<usize>,         // Page number, default 1
    pub include_low_confidence: Option<bool>,
}

/// Single node in knowledge graph.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct GraphNode {
    pub id: String,
    pub label: String,
    pub node_type: String,           // "person", "tool", "concept", etc.
    pub confidence: f32,
    pub summary: String,
    pub depth: usize,                // Which hop (0=root, 1=one away, etc.)
    pub incoming_edges: usize,       // How many edges point to this
    pub outgoing_edges: usize,       // How many edges from this
    pub position: Option<Position>,  // For React Flow visualization
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Position {
    pub x: f32,
    pub y: f32,
}

/// Single edge in knowledge graph.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct GraphEdge {
    pub id: String,
    pub source: String,
    pub target: String,
    pub label: String,
    pub confidence: f32,
    pub depth: usize,                // Deepest hop this edge reaches
}

/// Performance metrics for visualization query.
#[derive(Clone, Debug, Serialize)]
pub struct PerformanceMetrics {
    pub query_time_ms: u64,
    pub depth_times_ms: HashMap<usize, u64>,  // Per-depth breakdown
    pub total_time_ms: u64,
}

/// Recommendation for query optimization.
#[derive(Clone, Debug, Serialize)]
pub struct Recommendation {
    pub issue: String,
    pub suggestion: String,
    pub expected_latency_ms: u64,
}

/// Depth breakdown (nodes per hop).
#[derive(Clone, Debug, Serialize)]
pub struct DepthBreakdown {
    pub depth_0: usize,
    pub depth_1: usize,
    pub depth_2: usize,
    pub depth_3: Option<usize>,
}

/// Response for graph visualization.
#[derive(Clone, Debug, Serialize)]
pub struct VisualizeResponse {
    pub query: String,
    pub project: String,
    
    pub pagination: PaginationMeta,
    pub depth_breakdown: DepthBreakdown,
    
    pub nodes: Vec<GraphNode>,
    pub edges: Vec<GraphEdge>,
    
    pub performance: PerformanceMetrics,
    pub recommendations: Vec<Recommendation>,
}

/// Graph query engine for visualization.
pub struct GraphVisualizer;

impl GraphVisualizer {
    /// Execute BFS traversal and return paginated graph.
    pub async fn visualize(
        req: &VisualizeRequest,
        _db: &str, // TODO: actual DB connection
    ) -> Result<VisualizeResponse, String> {
        let start = std::time::Instant::now();
        
        // Validate input
        let depth = req.depth.unwrap_or(2).min(3);
        let pagination = PaginationParams::new(req.limit, req.page)
            .map_err(|e| format!("Invalid pagination: {}", e))?;
        
        // TODO: Real implementation:
        // 1. Find seed nodes (entities matching query)
        // 2. BFS traverse up to depth
        // 3. Collect all nodes + edges
        // 4. Apply pagination
        // 5. Calculate recommendations
        
        // For now, return mock response
        let (offset, limit) = pagination.calculate_offset_limit();
        let total_nodes = 487;
        let total_pages = pagination.calculate_total_pages(total_nodes);
        
        let perf_metrics = PerformanceMetrics {
            query_time_ms: 145,
            depth_times_ms: {
                let mut map = HashMap::new();
                map.insert(1, 45);
                map.insert(2, 100);
                map
            },
            total_time_ms: start.elapsed().as_millis() as u64,
        };
        
        let recommendations = Self::generate_recommendations(
            total_nodes,
            perf_metrics.total_time_ms,
            &pagination,
        );
        
        Ok(VisualizeResponse {
            query: req.query.clone(),
            project: req.project.clone(),
            pagination: PaginationMeta::new(&pagination, total_nodes),
            depth_breakdown: DepthBreakdown {
                depth_0: 12,
                depth_1: 234,
                depth_2: 241,
                depth_3: None,
            },
            nodes: vec![], // TODO: populate from BFS
            edges: vec![], // TODO: populate from BFS
            performance: perf_metrics,
            recommendations,
        })
    }
    
    /// Generate optimization recommendations.
    fn generate_recommendations(
        total_nodes: usize,
        query_time_ms: u64,
        pagination: &PaginationParams,
    ) -> Vec<Recommendation> {
        let mut recommendations = Vec::new();
        
        // High node count recommendation
        if total_nodes > 300 {
            recommendations.push(Recommendation {
                issue: "high_result_count".to_string(),
                suggestion: format!(
                    "Try depth=1 to reduce from {}→234 nodes",
                    total_nodes
                ),
                expected_latency_ms: 95,
            });
        }
        
        // High latency recommendation
        if query_time_ms > 200 {
            recommendations.push(Recommendation {
                issue: "slow_query".to_string(),
                suggestion: "Use pagination (limit=50) instead of loading all nodes".to_string(),
                expected_latency_ms: 145,
            });
        }
        
        // Pagination recommendation
        let limit = pagination.limit.unwrap_or(50);
        if limit > 100 {
            recommendations.push(Recommendation {
                issue: "large_page_size".to_string(),
                suggestion: "Reduce limit to 50 for faster responses".to_string(),
                expected_latency_ms: 100,
            });
        }
        
        recommendations
    }
    
    /// Calculate layout positions for React Flow (force-directed).
    pub fn calculate_positions(
        nodes: &[GraphNode],
        _edges: &[GraphEdge],
    ) -> HashMap<String, Position> {
        let mut positions = HashMap::new();
        
        // Simple circular layout for now
        // TODO: Implement force-directed layout
        for (i, node) in nodes.iter().enumerate() {
            let angle = (i as f32 / nodes.len() as f32) * std::f32::consts::TAU;
            let x = 100.0 * angle.cos();
            let y = 100.0 * angle.sin();
            
            positions.insert(node.id.clone(), Position { x, y });
        }
        
        positions
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_visualize_request_defaults() {
        let req = VisualizeRequest {
            project: "poimen".to_string(),
            query: "kubernetes".to_string(),
            depth: None,
            limit: None,
            page: None,
            include_low_confidence: None,
        };
        
        assert_eq!(req.project, "poimen");
        assert_eq!(req.query, "kubernetes");
    }
    
    #[test]
    fn test_recommendations_high_node_count() {
        let pagination = PaginationParams::new(Some(50), Some(1)).unwrap();
        let recs = GraphVisualizer::generate_recommendations(400, 145, &pagination);
        
        assert!(recs.iter().any(|r| r.issue == "high_result_count"));
    }
    
    #[test]
    fn test_recommendations_slow_query() {
        let pagination = PaginationParams::new(Some(50), Some(1)).unwrap();
        let recs = GraphVisualizer::generate_recommendations(100, 300, &pagination);
        
        assert!(recs.iter().any(|r| r.issue == "slow_query"));
    }
    
    #[test]
    fn test_positions_calculated() {
        let nodes = vec![
            GraphNode {
                id: "n1".to_string(),
                label: "Node 1".to_string(),
                node_type: "tool".to_string(),
                confidence: 0.95,
                summary: "Test".to_string(),
                depth: 0,
                incoming_edges: 1,
                outgoing_edges: 2,
                position: None,
            },
        ];
        
        let positions = GraphVisualizer::calculate_positions(&nodes, &[]);
        assert!(positions.contains_key("n1"));
    }
}