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

/// HTTP handler for POST /memory/visualize endpoint.
///
/// Receives request with root entity ID and optional depth parameter.
/// Returns React Flow JSON with nodes, edges, and performance metrics.

use actix_web::{web, HttpRequest, HttpResponse};
use serde_json::json;
use crate::query::visualize_types::{VisualizeRequest, VisualizeResponse, ReactFlowNode, ReactFlowEdge, NodeData, EdgeData, NodeStyle, PerformanceMetrics, SummaryMetrics, TypeCount};
use crate::query::bfs_graph_traversal::BfsConfig;
use crate::query::force_directed_layout::ForceDirectedLayout;
use crate::http_server::AppState;
use crate::jwt_validator::JwtValidator;
use std::time::Instant;
use std::collections::HashMap;

/// POST /memory/visualize - Graph visualization with BFS + layout
///
/// Query params (in JSON body):
/// - root_id (required): Starting entity ID
/// - depth (optional, default 2): Max traversal depth (1-3)
/// - max_nodes (optional, default 50): Max nodes to return (1-500)
/// - max_edges_per_node (optional, default 5): Max edges per node (1-100)
///
/// Response:
/// - nodes: React Flow node objects with positions
/// - edges: React Flow edge objects
/// - depth_breakdown: Nodes/edges per depth level
/// - performance: Traversal + layout timing
/// - summary: Entity/relation type counts
pub async fn visualize_handler(
    req: HttpRequest,
    body: web::Json<VisualizeRequest>,
    state: web::Data<AppState>,
) -> HttpResponse {
    // 1. Validate JWT + rate limiting (centralized middleware)
    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(&req, &state, "visualize", 100) {
        return response;
    }
    
    // 2. Call handler
    match execute_visualize(&state, body.into_inner()).await {
        Ok(response) => {
            HttpResponse::Ok().json(response)
        }
        Err(e) => {
            eprintln!("Visualization error: {}", e);
            HttpResponse::InternalServerError().json(json!({
                "error": format!("Visualization failed: {}", e)
            }))
        }
    }
}

/// Execute visualization: BFS traversal + force-directed layout
async fn execute_visualize(
    state: &AppState,
    req: VisualizeRequest,
) -> Result<VisualizeResponse, String> {
    // Validate request
    req.validate()?;
    
    let start_time = Instant::now();
    
    // BFS traversal
    let bfs_config = BfsConfig {
        max_depth: req.depth.unwrap_or(2).min(3),
        max_nodes: req.max_nodes.unwrap_or(50),
        max_edges_per_node: req.max_edges_per_node.unwrap_or(5),
    };
    
    // Use pool from AppState
    let bfs = crate::query::bfs_graph_traversal::BfsGraphTraversal::new(state.pool.clone());
    let graph = bfs.traverse(&req.root_id, &bfs_config).await?;
    
    let traversal_time_ms = Instant::now().elapsed().as_millis() as u64;
    
    // Force-directed layout
    let layout_start = Instant::now();
    let layout = ForceDirectedLayout::layout(&graph, &crate::query::force_directed_layout::LayoutConfig::default());
    let layout_time_ms = layout_start.elapsed().as_millis() as u64;
    
    // Build React Flow nodes
    let nodes: Vec<ReactFlowNode> = graph.nodes.iter().map(|n| {
        let pos = layout.positions.get(&n.id)
            .copied()
            .unwrap_or_default();
        
        let background = NodeStyle::for_entity_type(&n.entity_type);
        
        ReactFlowNode {
            id: n.id.clone(),
            label: n.name.clone(),
            position: pos,
            data: NodeData {
                entity_type: n.entity_type.clone(),
                depth: n.depth,
                description: n.description.clone(),
            },
            style: Some(NodeStyle {
                background,
                border: "#333333".to_string(),
                width: 100.0,
                height: 60.0,
            }),
        }
    }).collect();
    
    // Build React Flow edges
    let edges: Vec<ReactFlowEdge> = graph.edges.iter().map(|e| {
        ReactFlowEdge {
            id: e.id.clone(),
            source: e.source_id.clone(),
            target: e.target_id.clone(),
            label: e.relation_type.clone(),
            data: EdgeData {
                relation_type: e.relation_type.clone(),
                strength: e.strength,
            },
        }
    }).collect();
    
    // Compute summary metrics
    let mut entity_types: HashMap<String, usize> = HashMap::new();
    for node in &graph.nodes {
        *entity_types.entry(node.entity_type.clone()).or_insert(0) += 1;
    }
    
    let mut relation_types: HashMap<String, usize> = HashMap::new();
    for edge in &graph.edges {
        *relation_types.entry(edge.relation_type.clone()).or_insert(0) += 1;
    }
    
    let entity_type_counts: Vec<TypeCount> = entity_types
        .into_iter()
        .map(|(name, count)| TypeCount { name, count })
        .collect();
    
    let relation_type_counts: Vec<TypeCount> = relation_types
        .into_iter()
        .map(|(name, count)| TypeCount { name, count })
        .collect();
    
    let total_time_ms = start_time.elapsed().as_millis() as u64;
    
    Ok(VisualizeResponse {
        nodes,
        edges,
        root_id: req.root_id,
        depth_breakdown: graph.depth_breakdown,
        performance: PerformanceMetrics {
            traversal_time_ms,
            layout_time_ms,
            total_time_ms,
        },
        summary: SummaryMetrics {
            total_nodes: graph.node_count,
            total_edges: graph.edge_count,
            max_depth_reached: graph.max_depth_reached,
            entity_types: entity_type_counts,
            relation_types: relation_type_counts,
        },
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_visualize_request_serialization() {
        let json_str = r#"{
            "root_id": "entity-1",
            "depth": 2,
            "max_nodes": 50,
            "max_edges_per_node": 5
        }"#;
        
        let req: VisualizeRequest = serde_json::from_str(json_str).unwrap();
        assert_eq!(req.root_id, "entity-1");
        assert_eq!(req.depth, Some(2));
    }
    
    #[test]
    fn test_visualize_request_minimal() {
        let json_str = r#"{"root_id": "entity-1"}"#;
        
        let req: VisualizeRequest = serde_json::from_str(json_str).unwrap();
        assert_eq!(req.root_id, "entity-1");
        assert_eq!(req.depth, None);
        assert_eq!(req.max_nodes, None);
    }
}