//! Graph Traversal and Relationship Inference
//!
//! Provides algorithms for traversing the knowledge graph and inferring
//! new relationships based on existing connections.

use crate::error::{ErrorWrapper, Result};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};

#[cfg(all(test, feature = "ai_backend"))]
use super::entities::Relationship;
#[cfg(feature = "ai_backend")]
use super::entities::{Entity, RelationType};
#[cfg(feature = "ai_backend")]
use super::graph::KnowledgeGraph;

/// Path in the knowledge graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Path {
    /// Sequence of entity IDs
    pub entities: Vec<String>,
    /// Relationships along the path
    pub relationships: Vec<RelationType>,
    /// Total path weight (sum of relationship weights)
    pub total_weight: f32,
}

impl Path {
    /// Get path length
    pub fn length(&self) -> usize {
        self.entities.len()
    }

    /// Get hop count (number of relationships)
    pub fn hops(&self) -> usize {
        self.relationships.len()
    }
}

/// Result of relationship inference
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferredRelationship {
    /// Source entity ID
    pub source: String,
    /// Target entity ID
    pub target: String,
    /// Inferred relationship type
    pub rel_type: RelationType,
    /// Confidence score (0.0 - 1.0)
    pub confidence: f32,
    /// The path that led to this inference
    pub path: Path,
}

/// Graph traversal methods
#[cfg(feature = "ai_backend")]
pub struct GraphTraversal;

#[cfg(feature = "ai_backend")]
impl GraphTraversal {
    /// Breadth-first search for entities
    pub fn bfs<'a>(
        kg: &'a KnowledgeGraph,
        start_id: &str,
        max_depth: usize,
    ) -> Result<Vec<&'a Entity>> {
        if !kg.contains_entity(start_id) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Entity '{}' not found",
                start_id
            )));
        }

        let mut visited = std::collections::HashSet::new();
        let mut queue = VecDeque::new();
        let mut results = Vec::new();

        queue.push_back((start_id.to_string(), 0));
        visited.insert(start_id.to_string());

        while let Some((id, depth)) = queue.pop_front() {
            if let Some(entity) = kg.get_entity(&id) {
                results.push(entity);

                if depth < max_depth {
                    if let Ok(neighbors) = kg.get_neighbors(&id) {
                        for neighbor in neighbors {
                            if !visited.contains(&neighbor.id) {
                                visited.insert(neighbor.id.clone());
                                queue.push_back((neighbor.id.clone(), depth + 1));
                            }
                        }
                    }
                }
            }
        }

        Ok(results)
    }

    /// Depth-first search for entities
    pub fn dfs<'a>(
        kg: &'a KnowledgeGraph,
        start_id: &str,
        max_depth: usize,
    ) -> Result<Vec<&'a Entity>> {
        if !kg.contains_entity(start_id) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Entity '{}' not found",
                start_id
            )));
        }

        let mut visited = std::collections::HashSet::new();
        let mut results = Vec::new();

        fn dfs_recursive<'a>(
            kg: &'a KnowledgeGraph,
            id: &str,
            depth: usize,
            max_depth: usize,
            visited: &mut std::collections::HashSet<String>,
            results: &mut Vec<&'a Entity>,
        ) {
            if let Some(entity) = kg.get_entity(id) {
                results.push(entity);

                if depth < max_depth {
                    if let Ok(neighbors) = kg.get_neighbors(id) {
                        for neighbor in neighbors {
                            if !visited.contains(&neighbor.id) {
                                visited.insert(neighbor.id.clone());
                                dfs_recursive(
                                    kg,
                                    &neighbor.id,
                                    depth + 1,
                                    max_depth,
                                    visited,
                                    results,
                                );
                            }
                        }
                    }
                }
            }
        }

        visited.insert(start_id.to_string());
        dfs_recursive(kg, start_id, 0, max_depth, &mut visited, &mut results);

        Ok(results)
    }

    /// Find all entities within n hops
    pub fn find_within_distance(
        kg: &KnowledgeGraph,
        start_id: &str,
        max_hops: usize,
    ) -> Result<HashMap<String, usize>> {
        if !kg.contains_entity(start_id) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Entity '{}' not found",
                start_id
            )));
        }

        let mut distances = HashMap::new();
        let mut queue = VecDeque::new();

        queue.push_back((start_id.to_string(), 0));
        distances.insert(start_id.to_string(), 0);

        while let Some((id, dist)) = queue.pop_front() {
            if dist < max_hops {
                if let Ok(neighbors) = kg.get_neighbors(&id) {
                    for neighbor in neighbors {
                        if !distances.contains_key(&neighbor.id) {
                            distances.insert(neighbor.id.clone(), dist + 1);
                            queue.push_back((neighbor.id.clone(), dist + 1));
                        }
                    }
                }
            }
        }

        Ok(distances)
    }

    /// Infer transitive relationships
    pub fn infer_transitive(
        kg: &KnowledgeGraph,
        start_id: &str,
        rel_type: &RelationType,
        max_hops: usize,
    ) -> Result<Vec<InferredRelationship>> {
        if !kg.contains_entity(start_id) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Entity '{}' not found",
                start_id
            )));
        }

        let mut inferred = Vec::new();
        let mut visited = std::collections::HashSet::new();
        let mut queue = VecDeque::new();

        queue.push_back((start_id.to_string(), Vec::new(), 1.0, Vec::new()));
        visited.insert(start_id.to_string());

        while let Some((current_id, mut entity_path, confidence, mut rel_path)) = queue.pop_front()
        {
            entity_path.push(current_id.clone());

            if let Ok(relationships) = kg.get_outgoing_relationships(&current_id) {
                for rel in relationships {
                    if rel.rel_type == *rel_type && rel_path.len() < max_hops {
                        rel_path.push(rel.rel_type.clone());
                        let new_confidence = confidence * rel.weight;

                        if !visited.contains(&rel.target) {
                            visited.insert(rel.target.clone());

                            if entity_path.len() > 1 {
                                inferred.push(InferredRelationship {
                                    source: start_id.to_string(),
                                    target: rel.target.clone(),
                                    rel_type: rel_type.clone(),
                                    confidence: new_confidence,
                                    path: Path {
                                        entities: entity_path.clone(),
                                        relationships: rel_path.clone(),
                                        total_weight: new_confidence,
                                    },
                                });
                            }

                            queue.push_back((
                                rel.target.clone(),
                                entity_path.clone(),
                                new_confidence,
                                rel_path.clone(),
                            ));
                        }
                    }
                }
            }
        }

        Ok(inferred)
    }

    /// Find common neighbors
    pub fn find_common_neighbors<'a>(
        kg: &'a KnowledgeGraph,
        id1: &str,
        id2: &str,
    ) -> Result<Vec<&'a Entity>> {
        let neighbors1 = kg.get_neighbors(id1)?;
        let neighbors2 = kg.get_neighbors(id2)?;

        let ids1: std::collections::HashSet<_> = neighbors1.iter().map(|e| e.id.as_str()).collect();

        Ok(neighbors2
            .into_iter()
            .filter(|e| ids1.contains(e.id.as_str()))
            .collect())
    }
}

#[cfg(all(test, feature = "ai_backend"))]
mod tests {
    use super::*;

    fn create_test_kg() -> KnowledgeGraph {
        let mut kg = KnowledgeGraph::new();
        let alice = Entity::new(
            "e1".into(),
            "Alice".into(),
            crate::ai::kg::entities::EntityType::Person,
        );
        let bob = Entity::new(
            "e2".into(),
            "Bob".into(),
            crate::ai::kg::entities::EntityType::Person,
        );
        let charlie = Entity::new(
            "e3".into(),
            "Charlie".into(),
            crate::ai::kg::entities::EntityType::Person,
        );
        let acme = Entity::new(
            "e4".into(),
            "ACME".into(),
            crate::ai::kg::entities::EntityType::Organization,
        );

        kg.add_entity(alice).unwrap();
        kg.add_entity(bob).unwrap();
        kg.add_entity(charlie).unwrap();
        kg.add_entity(acme).unwrap();

        kg.add_relationship(Relationship::new(
            "e1".into(),
            "e4".into(),
            RelationType::WorksFor,
        ))
        .unwrap();
        kg.add_relationship(Relationship::new(
            "e2".into(),
            "e4".into(),
            RelationType::WorksFor,
        ))
        .unwrap();
        kg.add_relationship(Relationship::new(
            "e3".into(),
            "e4".into(),
            RelationType::WorksFor,
        ))
        .unwrap();
        kg.add_relationship(Relationship::new(
            "e1".into(),
            "e2".into(),
            RelationType::RelatedTo,
        ))
        .unwrap();

        kg
    }

    #[test]
    fn test_bfs() {
        let kg = create_test_kg();
        let results = GraphTraversal::bfs(&kg, "e1", 2).unwrap();
        assert!(!results.is_empty());
    }

    #[test]
    fn test_dfs() {
        let kg = create_test_kg();
        let results = GraphTraversal::dfs(&kg, "e1", 2).unwrap();
        assert!(!results.is_empty());
    }

    #[test]
    fn test_find_within_distance() {
        let kg = create_test_kg();
        let distances = GraphTraversal::find_within_distance(&kg, "e1", 2).unwrap();
        assert!(distances.contains_key("e1"));
        assert_eq!(distances.get("e1"), Some(&0));
    }

    #[test]
    fn test_common_neighbors() {
        let kg = create_test_kg();
        let common = GraphTraversal::find_common_neighbors(&kg, "e1", "e2").unwrap();
        assert!(!common.is_empty());
    }
}
