//! RAG + Knowledge Graph Integration
//!
//! Combines retrieval-augmented generation with knowledge graph
//! for contextual entity extraction and relationship building.

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

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

/// Extracted entities from a document
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DocumentEntities {
    /// Document ID
    pub doc_id: String,
    /// Extracted entities
    pub entities: Vec<Entity>,
    /// Relationships between entities
    pub relationships: Vec<Relationship>,
}

impl DocumentEntities {
    /// Create new document entities
    pub fn new(doc_id: String) -> Self {
        DocumentEntities {
            doc_id,
            entities: Vec::new(),
            relationships: Vec::new(),
        }
    }

    /// Add extracted entity
    pub fn add_entity(&mut self, entity: Entity) {
        self.entities.push(entity);
    }

    /// Add relationship
    pub fn add_relationship(&mut self, rel: Relationship) {
        self.relationships.push(rel);
    }
}

/// KG-augmented retrieval result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KgAugmentedResult {
    /// Original retrieval result
    pub retrieval: RetrievalResult,
    /// Entities mentioned in the document
    pub entities: Vec<Entity>,
    /// Related entities from knowledge graph
    pub related_entities: Vec<Entity>,
    /// Relationships to explore
    pub relationships: Vec<Relationship>,
}

/// RAG + KG Integration Manager
#[cfg(feature = "ai_backend")]
pub struct RagKgIntegration {
    kg: KnowledgeGraph,
    doc_entities: HashMap<String, DocumentEntities>,
}

#[cfg(feature = "ai_backend")]
impl RagKgIntegration {
    /// Create new integration manager
    pub fn new() -> Self {
        RagKgIntegration {
            kg: KnowledgeGraph::new(),
            doc_entities: HashMap::new(),
        }
    }

    /// Add document with extracted entities and relationships
    pub fn add_document(&mut self, doc_entities: DocumentEntities) -> Result<()> {
        if doc_entities.doc_id.is_empty() {
            return Err(ErrorWrapper::validation_failed(
                "Document ID cannot be empty",
            ));
        }

        // Add entities to knowledge graph
        for entity in doc_entities.entities.clone() {
            let entity_with_source = entity.with_source(doc_entities.doc_id.clone());
            self.kg.add_entity(entity_with_source)?;
        }

        // Add relationships to knowledge graph
        for rel in doc_entities.relationships.clone() {
            self.kg.add_relationship(rel)?;
        }

        // Store document entities
        self.doc_entities
            .insert(doc_entities.doc_id.clone(), doc_entities);

        Ok(())
    }

    /// Augment retrieval results with KG context
    pub fn augment_result(&self, result: RetrievalResult) -> Result<KgAugmentedResult> {
        let doc_entities = self
            .doc_entities
            .get(&result.doc_id)
            .map(|de| de.entities.clone())
            .unwrap_or_default();

        // Find related entities in the graph
        let mut related_entities = Vec::new();
        let mut relationships = Vec::new();

        for entity in &doc_entities {
            // Get neighbors
            if let Ok(neighbors) = self.kg.get_neighbors(&entity.id) {
                related_entities.extend(neighbors.iter().map(|e| (*e).clone()));
            }

            // Get outgoing relationships
            if let Ok(rels) = self.kg.get_outgoing_relationships(&entity.id) {
                relationships.extend(rels.iter().map(|r| (*r).clone()));
            }
        }

        Ok(KgAugmentedResult {
            retrieval: result,
            entities: doc_entities,
            related_entities,
            relationships,
        })
    }

    /// Get entity context
    pub fn get_entity_context(&self, entity_id: &str) -> Result<EntityContext> {
        if !self.kg.contains_entity(entity_id) {
            return Err(ErrorWrapper::validation_failed(format!(
                "Entity '{}' not found",
                entity_id
            )));
        }

        let entity = self.kg.get_entity(entity_id).unwrap().clone();

        let incoming = self
            .kg
            .get_incoming_neighbors(entity_id)
            .unwrap_or_default()
            .iter()
            .map(|e| (*e).clone())
            .collect();

        let outgoing = self
            .kg
            .get_neighbors(entity_id)
            .unwrap_or_default()
            .iter()
            .map(|e| (*e).clone())
            .collect();

        Ok(EntityContext {
            entity,
            incoming_entities: incoming,
            outgoing_entities: outgoing,
        })
    }

    /// Find entities related to a query
    pub fn find_related_entities(&self, query: &str, limit: usize) -> Vec<Entity> {
        let query_lower = query.to_lowercase();
        self.kg
            .all_entities()
            .iter()
            .filter(|e| {
                e.name.to_lowercase().contains(&query_lower)
                    || e.description
                        .as_ref()
                        .map(|d| d.to_lowercase().contains(&query_lower))
                        .unwrap_or(false)
            })
            .take(limit)
            .map(|e| (*e).clone())
            .collect()
    }

    /// Get knowledge graph statistics
    pub fn get_stats(&self) -> KgStats {
        let entity_types = self
            .kg
            .all_entities()
            .iter()
            .fold(HashMap::new(), |mut acc, entity| {
                *acc.entry(entity.entity_type).or_insert(0) += 1;
                acc
            });

        let rel_types = self
            .kg
            .all_relationships()
            .iter()
            .fold(HashMap::new(), |mut acc, rel| {
                let type_str = rel.rel_type.as_str().to_string();
                *acc.entry(type_str).or_insert(0) += 1;
                acc
            });

        KgStats {
            total_entities: self.kg.entity_count(),
            total_relationships: self.kg.relationship_count(),
            entity_types,
            relationship_types: rel_types,
        }
    }

    /// Get knowledge graph reference
    pub fn graph(&self) -> &KnowledgeGraph {
        &self.kg
    }

    /// Get mutable knowledge graph reference
    pub fn graph_mut(&mut self) -> &mut KnowledgeGraph {
        &mut self.kg
    }
}

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

/// Entity context in the knowledge graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EntityContext {
    /// The entity itself
    pub entity: Entity,
    /// Entities pointing to this entity
    pub incoming_entities: Vec<Entity>,
    /// Entities this entity points to
    pub outgoing_entities: Vec<Entity>,
}

/// Knowledge graph statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KgStats {
    /// Total number of entities
    pub total_entities: usize,
    /// Total number of relationships
    pub total_relationships: usize,
    /// Count by entity type
    pub entity_types: HashMap<EntityType, usize>,
    /// Count by relationship type
    pub relationship_types: HashMap<String, usize>,
}

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

    fn create_doc_entities() -> DocumentEntities {
        let mut doc = DocumentEntities::new("doc1".to_string());
        let alice = Entity::new("e1".into(), "Alice".into(), EntityType::Person)
            .with_description("Software engineer".into());
        let acme = Entity::new("e2".into(), "ACME".into(), EntityType::Organization);

        doc.add_entity(alice);
        doc.add_entity(acme);
        doc.add_relationship(Relationship::new(
            "e1".into(),
            "e2".into(),
            RelationType::WorksFor,
        ));

        doc
    }

    #[test]
    fn test_add_document() {
        let mut integration = RagKgIntegration::new();
        let doc = create_doc_entities();
        assert!(integration.add_document(doc).is_ok());
        assert_eq!(integration.kg.entity_count(), 2);
    }

    #[test]
    fn test_find_related_entities() {
        let mut integration = RagKgIntegration::new();
        let doc = create_doc_entities();
        integration.add_document(doc).unwrap();

        let results = integration.find_related_entities("Alice", 10);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "e1");
    }

    #[test]
    fn test_get_stats() {
        let mut integration = RagKgIntegration::new();
        let doc = create_doc_entities();
        integration.add_document(doc).unwrap();

        let stats = integration.get_stats();
        assert_eq!(stats.total_entities, 2);
        assert_eq!(stats.total_relationships, 1);
    }

    #[test]
    fn test_get_entity_context() {
        let mut integration = RagKgIntegration::new();
        let doc = create_doc_entities();
        integration.add_document(doc).unwrap();

        let context = integration.get_entity_context("e1").unwrap();
        assert_eq!(context.entity.id, "e1");
    }
}
