//! Knowledge Graph Implementation using petgraph
//!
//! Provides graph construction, traversal, and relationship management
//! using the petgraph library for efficient graph operations.

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

#[cfg(feature = "ai_backend")]
use petgraph::graph::{DiGraph, NodeIndex};
#[cfg(feature = "ai_backend")]
use petgraph::Direction;

#[cfg(test)]
use super::entities::RelationType;
use super::entities::{Entity, EntityType, Relationship};

/// Knowledge Graph using petgraph DiGraph
#[cfg(feature = "ai_backend")]
pub struct KnowledgeGraph {
    graph: DiGraph<Entity, Relationship>,
    entity_index: HashMap<String, NodeIndex>,
}

/// Serializable representation of a knowledge graph
#[cfg(feature = "ai_backend")]
#[derive(Serialize, Deserialize)]
struct KnowledgeGraphSnapshot {
    /// All entities in the graph
    entities: Vec<(String, Entity)>,
    /// All relationships in the graph
    relationships: Vec<(String, String, Relationship)>,
}

#[cfg(feature = "ai_backend")]
impl KnowledgeGraph {
    /// Create a new knowledge graph
    pub fn new() -> Self {
        KnowledgeGraph {
            graph: DiGraph::new(),
            entity_index: HashMap::new(),
        }
    }

    /// Add an entity to the graph
    pub fn add_entity(&mut self, entity: Entity) -> Result<()> {
        if entity.id.is_empty() {
            return Err(ErrorWrapper::validation_failed("Entity ID cannot be empty"));
        }

        if self.entity_index.contains_key(&entity.id) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Entity with ID '{}' already exists",
                entity.id
            )));
        }

        let node_idx = self.graph.add_node(entity.clone());
        self.entity_index.insert(entity.id.clone(), node_idx);
        Ok(())
    }

    /// Add multiple entities
    pub fn add_entities(&mut self, entities: Vec<Entity>) -> Result<()> {
        for entity in entities {
            self.add_entity(entity)?;
        }
        Ok(())
    }

    /// Add a relationship between entities
    pub fn add_relationship(&mut self, relationship: Relationship) -> Result<()> {
        let source_idx = self.entity_index.get(&relationship.source).ok_or_else(|| {
            ErrorWrapper::validation_failed(format!(
                "Source entity '{}' not found",
                relationship.source
            ))
        })?;

        let target_idx = self.entity_index.get(&relationship.target).ok_or_else(|| {
            ErrorWrapper::validation_failed(format!(
                "Target entity '{}' not found",
                relationship.target
            ))
        })?;

        self.graph.add_edge(*source_idx, *target_idx, relationship);
        Ok(())
    }

    /// Get entity by ID
    pub fn get_entity(&self, id: &str) -> Option<&Entity> {
        self.entity_index
            .get(id)
            .and_then(|idx| self.graph.node_weight(*idx))
    }

    /// Get mutable entity by ID
    pub fn get_entity_mut(&mut self, id: &str) -> Option<&mut Entity> {
        if let Some(idx) = self.entity_index.get(id).copied() {
            self.graph.node_weight_mut(idx)
        } else {
            None
        }
    }

    /// Remove entity and its relationships
    pub fn remove_entity(&mut self, id: &str) -> Option<Entity> {
        if let Some(idx) = self.entity_index.remove(id) {
            self.graph.remove_node(idx)
        } else {
            None
        }
    }

    /// Find entities by type
    pub fn find_by_type(&self, entity_type: EntityType) -> Vec<&Entity> {
        self.graph
            .node_weights()
            .filter(|entity| entity.entity_type == entity_type)
            .collect()
    }

    /// Find entities by name (substring matching)
    pub fn find_by_name(&self, name: &str) -> Vec<&Entity> {
        let name_lower = name.to_lowercase();
        self.graph
            .node_weights()
            .filter(|entity| entity.name.to_lowercase().contains(&name_lower))
            .collect()
    }

    /// Get entities connected to a given entity
    pub fn get_neighbors(&self, id: &str) -> Result<Vec<&Entity>> {
        let idx = self
            .entity_index
            .get(id)
            .ok_or_else(|| ErrorWrapper::validation_failed(format!("Entity '{}' not found", id)))?;

        Ok(self
            .graph
            .neighbors(*idx)
            .filter_map(|neighbor_idx| self.graph.node_weight(neighbor_idx))
            .collect())
    }

    /// Get incoming neighbors (entities pointing to this entity)
    pub fn get_incoming_neighbors(&self, id: &str) -> Result<Vec<&Entity>> {
        let idx = self
            .entity_index
            .get(id)
            .ok_or_else(|| ErrorWrapper::validation_failed(format!("Entity '{}' not found", id)))?;

        Ok(self
            .graph
            .neighbors_directed(*idx, Direction::Incoming)
            .filter_map(|neighbor_idx| self.graph.node_weight(neighbor_idx))
            .collect())
    }

    /// Get relationships from an entity
    pub fn get_outgoing_relationships(&self, id: &str) -> Result<Vec<&Relationship>> {
        let idx = self
            .entity_index
            .get(id)
            .ok_or_else(|| ErrorWrapper::validation_failed(format!("Entity '{}' not found", id)))?;

        Ok(self
            .graph
            .edges_directed(*idx, Direction::Outgoing)
            .map(|edge| edge.weight())
            .collect())
    }

    /// Get all relationships to an entity
    pub fn get_incoming_relationships(&self, id: &str) -> Result<Vec<&Relationship>> {
        let idx = self
            .entity_index
            .get(id)
            .ok_or_else(|| ErrorWrapper::validation_failed(format!("Entity '{}' not found", id)))?;

        Ok(self
            .graph
            .edges_directed(*idx, Direction::Incoming)
            .map(|edge| edge.weight())
            .collect())
    }

    /// Find shortest path between two entities
    pub fn find_path(&self, from_id: &str, to_id: &str) -> Result<Option<Vec<String>>> {
        use petgraph::algo::astar;

        let from_idx = self.entity_index.get(from_id).ok_or_else(|| {
            ErrorWrapper::validation_failed(format!("Source '{}' not found", from_id))
        })?;

        let to_idx = self.entity_index.get(to_id).ok_or_else(|| {
            ErrorWrapper::validation_failed(format!("Target '{}' not found", to_id))
        })?;

        let result = astar(
            &self.graph,
            *from_idx,
            |finish| finish == *to_idx,
            |_| 1,
            |_| 0,
        );

        Ok(result.map(|(_, path)| {
            path.into_iter()
                .filter_map(|idx| self.graph.node_weight(idx).map(|entity| entity.id.clone()))
                .collect()
        }))
    }

    /// Get entity count
    pub fn entity_count(&self) -> usize {
        self.graph.node_count()
    }

    /// Get relationship count
    pub fn relationship_count(&self) -> usize {
        self.graph.edge_count()
    }

    /// Check if entity exists
    pub fn contains_entity(&self, id: &str) -> bool {
        self.entity_index.contains_key(id)
    }

    /// Clear the graph
    pub fn clear(&mut self) {
        self.graph.clear();
        self.entity_index.clear();
    }

    /// Get all entities
    pub fn all_entities(&self) -> Vec<&Entity> {
        self.graph.node_weights().collect()
    }

    /// Get all relationships
    pub fn all_relationships(&self) -> Vec<&Relationship> {
        self.graph.edge_weights().collect()
    }

    /// Serialize knowledge graph for persistence
    pub fn serialize(&self) -> Result<Vec<u8>> {
        // Convert to serializable snapshot format
        let entities: Vec<_> = self
            .entity_index
            .iter()
            .filter_map(|(id, &idx)| self.graph.node_weight(idx).map(|e| (id.clone(), e.clone())))
            .collect();

        let relationships: Vec<_> = self
            .graph
            .edge_indices()
            .filter_map(|edge_idx| {
                let (source_idx, target_idx) = self.graph.edge_endpoints(edge_idx)?;
                let source_id = self
                    .entity_index
                    .iter()
                    .find(|(_, &idx)| idx == source_idx)
                    .map(|(id, _)| id.clone())?;
                let target_id = self
                    .entity_index
                    .iter()
                    .find(|(_, &idx)| idx == target_idx)
                    .map(|(id, _)| id.clone())?;
                let rel = self.graph.edge_weight(edge_idx)?.clone();
                Some((source_id, target_id, rel))
            })
            .collect();

        let snapshot = KnowledgeGraphSnapshot {
            entities,
            relationships,
        };

        bincode::serialize(&snapshot).map_err(|e| {
            ErrorWrapper::validation_failed(format!("Failed to serialize knowledge graph: {}", e))
        })
    }

    /// Deserialize knowledge graph from persistence
    pub fn deserialize(data: &[u8]) -> Result<Self> {
        let snapshot: KnowledgeGraphSnapshot = bincode::deserialize(data).map_err(|e| {
            ErrorWrapper::validation_failed(format!("Failed to deserialize knowledge graph: {}", e))
        })?;

        let mut kg = KnowledgeGraph::new();

        // Add all entities
        for (_id, entity) in snapshot.entities {
            kg.add_entity(entity)?;
        }

        // Add all relationships
        for (source_id, target_id, rel) in snapshot.relationships {
            // Ensure relationship's source and target match the provided IDs
            let mut rel = rel;
            rel.source = source_id;
            rel.target = target_id;
            kg.add_relationship(rel)?;
        }

        Ok(kg)
    }

    /// Save knowledge graph to a file
    pub fn save_to_file(&self, path: &str) -> Result<()> {
        let graph_data = self.serialize()?;
        let snapshot = super::super::persistence::KnowledgeGraphSnapshot {
            version: super::super::persistence::Persistence::current_version(),
            graph_data,
        };
        super::super::persistence::Persistence::save_binary(&snapshot, path)
    }

    /// Load knowledge graph from a file
    pub fn load_from_file(path: &str) -> Result<Self> {
        let snapshot: super::super::persistence::KnowledgeGraphSnapshot =
            super::super::persistence::Persistence::load_binary(path)?;

        if !super::super::persistence::Persistence::is_version_compatible(snapshot.version) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Incompatible snapshot version: {}",
                snapshot.version
            )));
        }

        Self::deserialize(&snapshot.graph_data)
    }
}

#[cfg(feature = "ai_backend")]
impl Default for KnowledgeGraph {
    fn default() -> Self {
        Self::new()
    }
}

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

    fn create_test_graph() -> KnowledgeGraph {
        let mut kg = KnowledgeGraph::new();
        let alice = Entity::new("e1".into(), "Alice".into(), EntityType::Person);
        let bob = Entity::new("e2".into(), "Bob".into(), EntityType::Person);
        let acme = Entity::new("e3".into(), "ACME Corp".into(), EntityType::Organization);

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

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

        kg
    }

    #[test]
    fn test_add_entity() {
        let mut kg = KnowledgeGraph::new();
        let entity = Entity::new("e1".into(), "Alice".into(), EntityType::Person);
        assert!(kg.add_entity(entity).is_ok());
        assert_eq!(kg.entity_count(), 1);
    }

    #[test]
    fn test_duplicate_entity() {
        let mut kg = KnowledgeGraph::new();
        let entity = Entity::new("e1".into(), "Alice".into(), EntityType::Person);
        kg.add_entity(entity.clone()).unwrap();
        assert!(kg.add_entity(entity).is_err());
    }

    #[test]
    fn test_add_relationship() {
        let kg = create_test_graph();
        assert_eq!(kg.relationship_count(), 3);
    }

    #[test]
    fn test_get_entity() {
        let kg = create_test_graph();
        let entity = kg.get_entity("e1").unwrap();
        assert_eq!(entity.name, "Alice");
    }

    #[test]
    fn test_find_by_type() {
        let kg = create_test_graph();
        let persons = kg.find_by_type(EntityType::Person);
        assert_eq!(persons.len(), 2);
    }

    #[test]
    fn test_find_by_name() {
        let kg = create_test_graph();
        let results = kg.find_by_name("Alice");
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "e1");
    }

    #[test]
    fn test_get_neighbors() {
        let kg = create_test_graph();
        let neighbors = kg.get_neighbors("e1").unwrap();
        assert_eq!(neighbors.len(), 2); // Bob and ACME
    }

    #[test]
    fn test_get_outgoing_relationships() {
        let kg = create_test_graph();
        let rels = kg.get_outgoing_relationships("e1").unwrap();
        assert_eq!(rels.len(), 2);
    }

    #[test]
    fn test_find_path() {
        let kg = create_test_graph();
        let path = kg.find_path("e1", "e2").unwrap();
        assert!(path.is_some());
    }

    #[test]
    fn test_remove_entity() {
        let mut kg = create_test_graph();
        assert!(kg.remove_entity("e1").is_some());
        assert_eq!(kg.entity_count(), 2);
    }

    #[test]
    fn test_kg_save_and_load() {
        use std::fs;
        use tempfile::NamedTempFile;

        let temp_file = NamedTempFile::new().unwrap();
        let path = temp_file.path().to_str().unwrap().to_string();

        // Create and save graph
        let kg = create_test_graph();
        kg.save_to_file(&path).unwrap();
        assert!(fs::metadata(&path).is_ok());

        // Load
        let loaded = KnowledgeGraph::load_from_file(&path).unwrap();
        assert_eq!(loaded.entity_count(), 3);
        assert_eq!(loaded.relationship_count(), 3);

        // Verify entities
        let entity = loaded.get_entity("e1").unwrap();
        assert_eq!(entity.name, "Alice");

        fs::remove_file(&path).ok();
    }

    #[test]
    fn test_kg_save_and_load_empty() {
        use std::fs;
        use tempfile::NamedTempFile;

        let temp_file = NamedTempFile::new().unwrap();
        let path = temp_file.path().to_str().unwrap().to_string();

        let kg = KnowledgeGraph::new();
        kg.save_to_file(&path).unwrap();

        let loaded = KnowledgeGraph::load_from_file(&path).unwrap();
        assert_eq!(loaded.entity_count(), 0);
        assert_eq!(loaded.relationship_count(), 0);

        fs::remove_file(&path).ok();
    }

    #[test]
    fn test_kg_serialize_deserialize() {
        let kg = create_test_graph();

        // Serialize
        let data = kg.serialize().unwrap();
        assert!(!data.is_empty());

        // Deserialize
        let loaded = KnowledgeGraph::deserialize(&data).unwrap();
        assert_eq!(loaded.entity_count(), kg.entity_count());
        assert_eq!(loaded.relationship_count(), kg.relationship_count());
    }
}
