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

//! Path Finding Engine
//!
//! Finds paths through the knowledge graph using BFS, DFS, and shortest path algorithms.
//! Enables relationship traversal, distance analysis, and connection discovery.

use serde::{Deserialize, Serialize};
use sqlx::{Pool, Postgres};
use std::collections::{HashMap, HashSet, VecDeque};
use tracing::{debug, info};

/// A single path through the graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Path {
    pub source_id: String,
    pub target_id: String,
    pub entity_ids: Vec<String>,           // All entities in path
    pub entity_names: Vec<String>,         // Human-readable names
    pub relation_types: Vec<String>,       // Relations along path
    pub distance: usize,                   // Number of hops
    pub total_confidence: f32,             // Product of edge confidences
}

/// K-hop neighborhood around an entity
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KHopNeighborhood {
    pub center_id: String,
    pub center_name: String,
    pub k: usize,                          // Hop distance
    pub entities: Vec<(String, String, usize)>,  // (id, name, hops_away)
    pub entity_count: usize,
    pub edge_count: usize,
}

/// Results from path finding
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PathFindingResult {
    pub source_id: String,
    pub target_id: String,
    pub paths_found: Vec<Path>,
    pub path_count: usize,
    pub shortest_distance: Option<usize>,
    pub average_distance: f32,
}

/// Edge representation for path finding
#[derive(Debug, Clone)]
struct GraphEdge {
    from_id: String,
    to_id: String,
    relation_type: String,
    confidence: f32,
}

/// Path Finder for graph traversal
pub struct PathFinder {
    pub pool: Pool<Postgres>,
}

impl PathFinder {
    /// Create a new path finder
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }

    /// Find shortest path between two entities using BFS
    ///
    /// # Arguments
    /// * `source_id` - Starting entity ID
    /// * `target_id` - Ending entity ID
    /// * `max_depth` - Maximum hops to explore (default 5, max 10)
    ///
    /// # Returns
    /// Path with shortest distance, or error if no path found
    pub async fn shortest_path(
        &self,
        source_id: &str,
        target_id: &str,
        max_depth: usize,
    ) -> Result<Option<Path>, String> {
        let max_depth = max_depth.max(1).min(10);

        debug!("Finding shortest path: {} → {}, max_depth={}", 
               source_id, target_id, max_depth);

        if source_id == target_id {
            return Ok(Some(Path {
                source_id: source_id.to_string(),
                target_id: target_id.to_string(),
                entity_ids: vec![source_id.to_string()],
                entity_names: vec![],
                relation_types: vec![],
                distance: 0,
                total_confidence: 1.0,
            }));
        }

        // BFS: level by level traversal
        let mut queue: VecDeque<(String, Vec<String>, Vec<String>, f32)> = VecDeque::new();
        let mut visited: HashSet<String> = HashSet::new();

        queue.push_back((source_id.to_string(), vec![source_id.to_string()], vec![], 1.0));
        visited.insert(source_id.to_string());

        while let Some((current_id, path_entities, path_relations, confidence)) = queue.pop_front() {
            if path_entities.len() - 1 >= max_depth {
                continue;  // Depth limit reached
            }

            // Fetch neighbors of current entity
            let neighbors = self.fetch_neighbors(&current_id).await?;

            for edge in neighbors {
                if edge.to_id == target_id {
                    // Found target!
                    let mut final_entities = path_entities.clone();
                    final_entities.push(target_id.to_string());

                    let mut final_relations = path_relations.clone();
                    final_relations.push(edge.relation_type.clone());

                    let final_confidence = confidence * edge.confidence;

                    info!("Found shortest path: {} → {} (distance: {})", 
                          source_id, target_id, final_entities.len() - 1);

                    return Ok(Some(Path {
                        source_id: source_id.to_string(),
                        target_id: target_id.to_string(),
                        entity_ids: final_entities,
                        entity_names: vec![],  // Could fetch from DB if needed
                        relation_types: final_relations,
                        distance: final_entities.len() - 1,
                        total_confidence: final_confidence.max(0.0).min(1.0),
                    }));
                }

                if !visited.contains(&edge.to_id) {
                    visited.insert(edge.to_id.clone());
                    let mut next_entities = path_entities.clone();
                    next_entities.push(edge.to_id.clone());

                    let mut next_relations = path_relations.clone();
                    next_relations.push(edge.relation_type.clone());

                    let next_confidence = confidence * edge.confidence;

                    queue.push_back((
                        edge.to_id.clone(),
                        next_entities,
                        next_relations,
                        next_confidence,
                    ));
                }
            }
        }

        debug!("No path found between {} and {}", source_id, target_id);
        Ok(None)
    }

    /// Find all entities within K hops of a source entity
    ///
    /// # Arguments
    /// * `source_id` - Starting entity ID
    /// * `k` - Number of hops (default 2, max 5)
    ///
    /// # Returns
    /// KHopNeighborhood with all entities within k hops
    pub async fn k_hop_neighbors(
        &self,
        source_id: &str,
        k: usize,
    ) -> Result<KHopNeighborhood, String> {
        let k = k.max(1).min(5);

        debug!("Finding {}-hop neighbors of {}", k, source_id);

        let mut current_level = vec![source_id.to_string()];
        let mut all_neighbors: HashMap<String, (String, usize)> = HashMap::new();  // id → (name, hops)
        let mut edge_count = 0;

        for hop in 1..=k {
            let mut next_level = Vec::new();

            for entity_id in &current_level {
                let neighbors = self.fetch_neighbors(entity_id).await?;

                for edge in neighbors {
                    if !all_neighbors.contains_key(&edge.to_id) && edge.to_id != source_id {
                        all_neighbors.insert(edge.to_id.clone(), ("".to_string(), hop));
                        next_level.push(edge.to_id.clone());
                    }
                    edge_count += 1;
                }
            }

            current_level = next_level;
            if current_level.is_empty() {
                break;  // No more neighbors to explore
            }
        }

        let entity_count = all_neighbors.len();
        let entities: Vec<_> = all_neighbors
            .into_iter()
            .map(|(id, (name, hops))| (id, name, hops))
            .collect();

        info!("Found {}-hop neighborhood: {} entities", k, entity_count);

        Ok(KHopNeighborhood {
            center_id: source_id.to_string(),
            center_name: "".to_string(),
            k,
            entities,
            entity_count,
            edge_count: edge_count.min(1000),  // Cap to prevent explosion
        })
    }

    /// Find all paths (up to max_paths) between two entities using DFS
    ///
    /// # Arguments
    /// * `source_id` - Starting entity ID
    /// * `target_id` - Ending entity ID
    /// * `max_depth` - Maximum hops per path (default 4)
    /// * `max_paths` - Maximum paths to find (default 10, max 50)
    ///
    /// # Returns
    /// PathFindingResult with all paths found (sorted by distance)
    pub async fn all_paths(
        &self,
        source_id: &str,
        target_id: &str,
        max_depth: usize,
        max_paths: usize,
    ) -> Result<PathFindingResult, String> {
        let max_depth = max_depth.max(1).min(6);
        let max_paths = max_paths.max(1).min(50);

        debug!("Finding all paths: {} → {}, max_depth={}, max_paths={}", 
               source_id, target_id, max_depth, max_paths);

        if source_id == target_id {
            return Ok(PathFindingResult {
                source_id: source_id.to_string(),
                target_id: target_id.to_string(),
                paths_found: vec![],
                path_count: 0,
                shortest_distance: Some(0),
                average_distance: 0.0,
            });
        }

        let mut paths_found = Vec::new();
        let mut visited = HashSet::new();

        self.dfs_paths(
            source_id,
            target_id,
            vec![source_id.to_string()],
            vec![],
            1.0,
            0,
            max_depth,
            &mut paths_found,
            &mut visited,
            max_paths,
        ).await?;

        // Sort by distance
        paths_found.sort_by_key(|p| p.distance);

        let shortest_distance = paths_found.first().map(|p| p.distance);
        let average_distance = if !paths_found.is_empty() {
            paths_found.iter().map(|p| p.distance as f32).sum::<f32>() / paths_found.len() as f32
        } else {
            0.0
        };

        let path_count = paths_found.len();
        info!("Found {} paths between {} and {} (avg distance: {:.2})", 
              path_count, source_id, target_id, average_distance);

        Ok(PathFindingResult {
            source_id: source_id.to_string(),
            target_id: target_id.to_string(),
            paths_found,
            path_count,
            shortest_distance,
            average_distance,
        })
    }

    /// DFS helper for finding all paths
    async fn dfs_paths(
        &self,
        source_id: &str,
        target_id: &str,
        current_path: Vec<String>,
        relations_path: Vec<String>,
        confidence: f32,
        depth: usize,
        max_depth: usize,
        paths_found: &mut Vec<Path>,
        visited: &mut HashSet<String>,
        max_paths: usize,
    ) -> Result<(), String> {
        if paths_found.len() >= max_paths {
            return Ok(());  // Found enough paths
        }

        if depth >= max_depth {
            return Ok(());  // Depth limit reached
        }

        let current_id = current_path.last().unwrap();
        let neighbors = self.fetch_neighbors(current_id).await?;

        for edge in neighbors {
            if edge.to_id == target_id {
                // Found a path!
                let mut final_path = current_path.clone();
                final_path.push(target_id.to_string());

                let mut final_relations = relations_path.clone();
                final_relations.push(edge.relation_type.clone());

                let final_confidence = confidence * edge.confidence;

                paths_found.push(Path {
                    source_id: source_id.to_string(),
                    target_id: target_id.to_string(),
                    entity_ids: final_path,
                    entity_names: vec![],
                    relation_types: final_relations,
                    distance: final_path.len() - 1,
                    total_confidence: final_confidence.max(0.0).min(1.0),
                });

                if paths_found.len() >= max_paths {
                    return Ok(());
                }
            } else if !current_path.contains(&edge.to_id) && !visited.contains(&edge.to_id) {
                // Continue DFS
                visited.insert(edge.to_id.clone());
                let mut next_path = current_path.clone();
                next_path.push(edge.to_id.clone());

                let mut next_relations = relations_path.clone();
                next_relations.push(edge.relation_type.clone());

                let next_confidence = confidence * edge.confidence;

                self.dfs_paths(
                    source_id,
                    target_id,
                    next_path,
                    next_relations,
                    next_confidence,
                    depth + 1,
                    max_depth,
                    paths_found,
                    visited,
                    max_paths,
                ).await?;

                visited.remove(&edge.to_id);  // Backtrack for DFS
            }
        }

        Ok(())
    }

    /// Fetch direct neighbors of an entity
    async fn fetch_neighbors(&self, entity_id: &str) -> Result<Vec<GraphEdge>, String> {
        let edges = sqlx::query_as::<_, (String, String, String, f32)>(
            "SELECT source_entity_id, target_entity_id, relation_type, confidence 
             FROM memory_edge 
             WHERE (source_entity_id = $1 OR target_entity_id = $1)
             AND fact_invalid_at IS NULL 
             AND deleted_at IS NULL"
        )
            .bind(entity_id)
            .fetch_all(&self.pool)
            .await
            .map_err(|e| format!("Failed to fetch neighbors: {}", e))?
            .into_iter()
            .map(|(source, target, rel_type, conf)| {
                // Normalize direction: always point forward from input entity
                if source == entity_id {
                    GraphEdge {
                        from_id: source,
                        to_id: target,
                        relation_type: rel_type,
                        confidence: conf.max(0.0).min(1.0),
                    }
                } else {
                    GraphEdge {
                        from_id: target,
                        to_id: source,
                        relation_type: format!("{}(reverse)", rel_type),
                        confidence: conf.max(0.0).min(1.0),
                    }
                }
            })
            .collect();

        Ok(edges)
    }
}

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

    #[test]
    fn test_path_creation() {
        let path = Path {
            source_id: "e1".to_string(),
            target_id: "e3".to_string(),
            entity_ids: vec!["e1".to_string(), "e2".to_string(), "e3".to_string()],
            entity_names: vec!["Entity1".to_string(), "Entity2".to_string(), "Entity3".to_string()],
            relation_types: vec!["related".to_string(), "connected".to_string()],
            distance: 2,
            total_confidence: 0.9,
        };

        assert_eq!(path.distance, 2);
        assert_eq!(path.entity_ids.len(), 3);
    }

    #[test]
    fn test_k_hop_neighborhood() {
        let neighborhood = KHopNeighborhood {
            center_id: "e1".to_string(),
            center_name: "Entity1".to_string(),
            k: 2,
            entities: vec![
                ("e2".to_string(), "Entity2".to_string(), 1),
                ("e3".to_string(), "Entity3".to_string(), 2),
            ],
            entity_count: 2,
            edge_count: 3,
        };

        assert_eq!(neighborhood.k, 2);
        assert_eq!(neighborhood.entity_count, 2);
    }

    #[test]
    fn test_path_distance_zero() {
        let path = Path {
            source_id: "e1".to_string(),
            target_id: "e1".to_string(),
            entity_ids: vec!["e1".to_string()],
            entity_names: vec![],
            relation_types: vec![],
            distance: 0,
            total_confidence: 1.0,
        };

        assert_eq!(path.distance, 0);
    }

    #[test]
    fn test_path_distance_one() {
        let path = Path {
            source_id: "e1".to_string(),
            target_id: "e2".to_string(),
            entity_ids: vec!["e1".to_string(), "e2".to_string()],
            entity_names: vec![],
            relation_types: vec!["related".to_string()],
            distance: 1,
            total_confidence: 0.95,
        };

        assert_eq!(path.distance, 1);
    }

    #[test]
    fn test_confidence_normalization() {
        let confidence = 0.7 * 0.8 * 0.9;  // 0.504
        let normalized = (confidence as f32).max(0.0).min(1.0);
        
        assert!(normalized >= 0.0 && normalized <= 1.0);
    }

    #[test]
    fn test_max_depth_clamping() {
        let max_depth = 0;
        let clamped = max_depth.max(1).min(10);
        assert_eq!(clamped, 1);

        let max_depth = 15;
        let clamped = max_depth.max(1).min(10);
        assert_eq!(clamped, 10);
    }

    #[test]
    fn test_k_hop_clamping() {
        let k = 0;
        let clamped = k.max(1).min(5);
        assert_eq!(clamped, 1);

        let k = 10;
        let clamped = k.max(1).min(5);
        assert_eq!(clamped, 5);
    }

    #[test]
    fn test_max_paths_clamping() {
        let max_paths = 0;
        let clamped = max_paths.max(1).min(50);
        assert_eq!(clamped, 1);

        let max_paths = 100;
        let clamped = max_paths.max(1).min(50);
        assert_eq!(clamped, 50);
    }

    #[test]
    fn test_path_finding_result() {
        let result = PathFindingResult {
            source_id: "e1".to_string(),
            target_id: "e5".to_string(),
            paths_found: vec![],
            path_count: 0,
            shortest_distance: None,
            average_distance: 0.0,
        };

        assert_eq!(result.path_count, 0);
        assert!(result.shortest_distance.is_none());
    }

    #[test]
    fn test_path_ordering_by_distance() {
        let mut paths = vec![
            Path {
                source_id: "e1".to_string(),
                target_id: "e4".to_string(),
                entity_ids: vec!["e1".to_string(), "e2".to_string(), "e3".to_string(), "e4".to_string()],
                entity_names: vec![],
                relation_types: vec![],
                distance: 3,
                total_confidence: 0.7,
            },
            Path {
                source_id: "e1".to_string(),
                target_id: "e4".to_string(),
                entity_ids: vec!["e1".to_string(), "e4".to_string()],
                entity_names: vec![],
                relation_types: vec![],
                distance: 1,
                total_confidence: 0.9,
            },
        ];

        paths.sort_by_key(|p| p.distance);

        assert_eq!(paths[0].distance, 1);
        assert_eq!(paths[1].distance, 3);
    }

    #[test]
    fn test_average_distance_calculation() {
        let distances = vec![1, 2, 3, 4, 5];
        let avg = distances.iter().map(|&d| d as f32).sum::<f32>() / distances.len() as f32;
        
        assert!((avg - 3.0).abs() < 0.01);
    }

    #[test]
    fn test_edge_representation() {
        let edge = GraphEdge {
            from_id: "e1".to_string(),
            to_id: "e2".to_string(),
            relation_type: "related".to_string(),
            confidence: 0.85,
        };

        assert_eq!(edge.from_id, "e1");
        assert_eq!(edge.to_id, "e2");
        assert!(edge.confidence >= 0.0 && edge.confidence <= 1.0);
    }

    #[test]
    fn test_reverse_edge_naming() {
        let relation = "depends_on".to_string();
        let reverse = format!("{}(reverse)", relation);

        assert_eq!(reverse, "depends_on(reverse)");
    }
}