//! Cross-Database Index - Global indexing for heterogeneous stores
//!
//! Provides unified entity location tracking and cross-type relationship indexing.
//! Supports persist()/restore() for durable storage to JSONL files.

use std::collections::HashMap;
use std::io::{BufRead, Write};
use std::path::{Path, PathBuf};
use std::sync::RwLock;

use serde::{Serialize, Deserialize};

use super::error::{FederationError, Result};
use super::traits::Store;
use super::types::{StoreType, EntityLocation, TypeCrossRelation, Record};

/// Statistics for the cross-database index
#[derive(Debug, Clone, Default)]
pub struct IndexStats {
    /// Number of indexed entities
    pub entity_count: usize,
    /// Number of cross-store relations
    pub cross_relation_count: usize,
    /// Number of cross-type relations
    pub type_cross_relation_count: usize,
    /// Index size in bytes (estimated)
    pub size_bytes: usize,
    /// Last rebuild timestamp
    pub last_rebuild: Option<u64>,
}

/// Cross-database index for entity location and relationships
pub struct CrossDbIndex {
    /// Entity locations: entity_id -> (store_alias, store_type)
    entity_locations: RwLock<HashMap<String, EntityLocation>>,
    /// Cross-store relations (same type): (source_id, target_id) -> relation info
    cross_relations: RwLock<HashMap<(String, String), CrossRelationInfo>>,
    /// Cross-type relations: source_id -> [(target_id, type_cross_relation)]
    type_cross_relations: RwLock<HashMap<String, Vec<TypeCrossRelation>>>,
    /// Reverse index: target_id -> source_ids (for refs_to queries)
    reverse_index: RwLock<HashMap<String, Vec<String>>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CrossRelationInfo {
    relation_type: String,
    source_store: String,
    target_store: String,
    attributes: HashMap<String, String>,
}

/// Serializable snapshot of the entire index (used by persist/restore)
#[derive(Debug, Serialize, Deserialize)]
struct IndexSnapshot {
    /// Version for forward-compatible deserialization
    version: u32,
    /// Entity locations
    entity_locations: HashMap<String, EntityLocation>,
    /// Cross-store relations keyed as "source_id\0target_id"
    cross_relations: Vec<CrossRelationEntry>,
    /// Cross-type relations
    type_cross_relations: HashMap<String, Vec<TypeCrossRelation>>,
    /// Reverse index
    reverse_index: HashMap<String, Vec<String>>,
}

#[derive(Debug, Serialize, Deserialize)]
struct CrossRelationEntry {
    source_id: String,
    target_id: String,
    info: CrossRelationInfo,
}

impl CrossDbIndex {
    /// Create a new cross-database index
    pub fn new() -> Self {
        Self {
            entity_locations: RwLock::new(HashMap::new()),
            cross_relations: RwLock::new(HashMap::new()),
            type_cross_relations: RwLock::new(HashMap::new()),
            reverse_index: RwLock::new(HashMap::new()),
        }
    }
    
    /// Clear all index data
    pub fn clear(&self) -> Result<()> {
        let mut entities = self.entity_locations.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        entities.clear();
        
        let mut relations = self.cross_relations.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        relations.clear();
        
        let mut type_rels = self.type_cross_relations.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        type_rels.clear();
        
        let mut reverse = self.reverse_index.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        reverse.clear();
        
        Ok(())
    }
    
    // ===== Entity Location =====
    
    /// Index an entity's location
    pub fn index_entity(&self, entity_id: &str, store_alias: &str, store_type: StoreType) -> Result<()> {
        let mut locations = self.entity_locations.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        locations.insert(entity_id.to_string(), EntityLocation {
            entity_id: entity_id.to_string(),
            store_alias: store_alias.to_string(),
            store_type,
        });
        
        Ok(())
    }
    
    /// Remove an entity from the index
    pub fn remove_entity(&self, entity_id: &str) -> Result<()> {
        let mut locations = self.entity_locations.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        locations.remove(entity_id);
        Ok(())
    }
    
    /// Find an entity's location
    pub fn find_entity(&self, entity_id: &str) -> Result<Option<EntityLocation>> {
        let locations = self.entity_locations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        Ok(locations.get(entity_id).cloned())
    }
    
    /// Find entities by store type
    pub fn find_by_type(&self, store_type: StoreType) -> Result<Vec<EntityLocation>> {
        let locations = self.entity_locations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        Ok(locations.values()
            .filter(|loc| loc.store_type == store_type)
            .cloned()
            .collect())
    }
    
    // ===== Cross-Store Relations =====
    
    /// Index a cross-store relation
    pub fn index_cross_relation(
        &self,
        source_id: &str,
        target_id: &str,
        relation_type: &str,
        source_store: &str,
        target_store: &str,
        attributes: HashMap<String, String>,
    ) -> Result<()> {
        let mut relations = self.cross_relations.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        relations.insert(
            (source_id.to_string(), target_id.to_string()),
            CrossRelationInfo {
                relation_type: relation_type.to_string(),
                source_store: source_store.to_string(),
                target_store: target_store.to_string(),
                attributes,
            },
        );
        
        // Update reverse index
        let mut reverse = self.reverse_index.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        reverse.entry(target_id.to_string())
            .or_insert_with(Vec::new)
            .push(source_id.to_string());
        
        Ok(())
    }
    
    // ===== Cross-Type Relations =====
    
    /// Index a cross-type relation (e.g., graph entity -> document)
    pub fn index_type_cross_relation(&self, source_id: &str, relation: TypeCrossRelation) -> Result<()> {
        let mut relations = self.type_cross_relations.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        relations.entry(source_id.to_string())
            .or_insert_with(Vec::new)
            .push(relation.clone());
        
        // Update reverse index
        let mut reverse = self.reverse_index.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        reverse.entry(relation.target_id.clone())
            .or_insert_with(Vec::new)
            .push(source_id.to_string());
        
        Ok(())
    }
    
    /// Find cross-type relations from an entity
    pub fn find_type_cross_from(&self, source_id: &str) -> Result<Vec<TypeCrossRelation>> {
        let relations = self.type_cross_relations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        Ok(relations.get(source_id)
            .cloned()
            .unwrap_or_default())
    }
    
    // ===== Reference Queries =====
    
    /// Find all entities that reference the given entity
    pub fn find_refs_to(&self, entity_id: &str) -> Result<Vec<String>> {
        let reverse = self.reverse_index.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        Ok(reverse.get(entity_id)
            .cloned()
            .unwrap_or_default())
    }
    
    /// Find all entities referenced by the given entity
    pub fn find_refs_from(&self, entity_id: &str) -> Result<Vec<String>> {
        let relations = self.cross_relations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        let mut refs: Vec<String> = relations.keys()
            .filter(|(src, _)| src == entity_id)
            .map(|(_, tgt)| tgt.clone())
            .collect();
        
        // Also check type_cross_relations
        let type_rels = self.type_cross_relations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        if let Some(rels) = type_rels.get(entity_id) {
            refs.extend(rels.iter().map(|r| r.target_id.clone()));
        }
        
        refs.sort();
        refs.dedup();
        Ok(refs)
    }
    
    // ===== Store Indexing =====
    
    /// Index all entities from a store
    pub fn index_store(&self, store_alias: &str, store: &dyn Store) -> Result<()> {
        let store_type = store.store_type();
        
        // Query all records from the store
        // This is a simplified implementation - in production, we'd use pagination
        let query = super::types::UniversalQuery {
            target: "*".to_string(), // All types
            filters: vec![],
            limit: Some(10000),
            offset: None,
            order_by: None,
            traverse: None,
            vector_search: None,
        };
        
        // Note: This won't work with "*" target in most stores
        // In practice, we'd need to iterate over known types
        if let Ok(records) = store.query(&query) {
            for record in records {
                let entity_id = match &record {
                    Record::GraphEntity { id, .. } => id,
                    Record::GraphRelation { id, .. } => &id.to_string(),
                    Record::Document { id, .. } => id,
                    Record::Vector { id, .. } => id,
                };
                
                self.index_entity(entity_id, store_alias, store_type)?;
            }
        }
        
        Ok(())
    }
    
    // ===== Persistence =====
    
    /// Persist the entire index to a JSON file.
    ///
    /// The file is written atomically (write to `.tmp` then rename) to avoid
    /// corruption on crash.
    pub fn persist(&self, path: &Path) -> Result<()> {
        // Acquire all read locks
        let entities = self.entity_locations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock for persist".to_string()))?;
        let relations = self.cross_relations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock for persist".to_string()))?;
        let type_rels = self.type_cross_relations.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock for persist".to_string()))?;
        let reverse = self.reverse_index.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock for persist".to_string()))?;
        
        // Build snapshot
        let cross_relation_entries: Vec<CrossRelationEntry> = relations.iter()
            .map(|((src, tgt), info)| CrossRelationEntry {
                source_id: src.clone(),
                target_id: tgt.clone(),
                info: info.clone(),
            })
            .collect();
        
        let snapshot = IndexSnapshot {
            version: 1,
            entity_locations: entities.clone(),
            cross_relations: cross_relation_entries,
            type_cross_relations: type_rels.clone(),
            reverse_index: reverse.clone(),
        };
        
        // Drop locks before IO
        drop(entities);
        drop(relations);
        drop(type_rels);
        drop(reverse);
        
        // Write atomically: tmp file → rename
        let tmp_path = path.with_extension("tmp");
        
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)
                .map_err(|e| FederationError::StorageError(format!("Failed to create index dir: {}", e)))?;
        }
        
        let file = std::fs::File::create(&tmp_path)
            .map_err(|e| FederationError::StorageError(format!("Failed to create index file: {}", e)))?;
        let mut writer = std::io::BufWriter::new(file);
        
        serde_json::to_writer(&mut writer, &snapshot)
            .map_err(|e| FederationError::StorageError(format!("Failed to serialize index: {}", e)))?;
        writer.flush()
            .map_err(|e| FederationError::StorageError(format!("Failed to flush index file: {}", e)))?;
        
        // Atomic rename
        std::fs::rename(&tmp_path, path)
            .map_err(|e| FederationError::StorageError(format!("Failed to rename index file: {}", e)))?;
        
        Ok(())
    }
    
    /// Restore the index from a previously persisted JSON file.
    ///
    /// Replaces all current in-memory data with the file contents.
    pub fn restore(path: &Path) -> Result<Self> {
        let file = std::fs::File::open(path)
            .map_err(|e| FederationError::StorageError(format!("Failed to open index file: {}", e)))?;
        let reader = std::io::BufReader::new(file);
        
        let snapshot: IndexSnapshot = serde_json::from_reader(reader)
            .map_err(|e| FederationError::StorageError(format!("Failed to deserialize index: {}", e)))?;
        
        if snapshot.version != 1 {
            return Err(FederationError::StorageError(
                format!("Unsupported index snapshot version: {}", snapshot.version)
            ));
        }
        
        // Rebuild cross_relations HashMap from entries
        let mut cross_relations = HashMap::new();
        for entry in snapshot.cross_relations {
            cross_relations.insert(
                (entry.source_id, entry.target_id),
                entry.info,
            );
        }
        
        Ok(Self {
            entity_locations: RwLock::new(snapshot.entity_locations),
            cross_relations: RwLock::new(cross_relations),
            type_cross_relations: RwLock::new(snapshot.type_cross_relations),
            reverse_index: RwLock::new(snapshot.reverse_index),
        })
    }
    
    /// Restore from file if it exists, otherwise create a new empty index.
    pub fn restore_or_new(path: &Path) -> Result<Self> {
        if path.exists() {
            Self::restore(path)
        } else {
            Ok(Self::new())
        }
    }
    
    // ===== Statistics =====
    
    /// Get index statistics
    pub fn stats(&self) -> IndexStats {
        let entity_count = self.entity_locations.read()
            .map(|l| l.len())
            .unwrap_or(0);
        
        let cross_relation_count = self.cross_relations.read()
            .map(|r| r.len())
            .unwrap_or(0);
        
        let type_cross_relation_count = self.type_cross_relations.read()
            .map(|r| r.values().map(|v| v.len()).sum())
            .unwrap_or(0);
        
        // Rough size estimation
        let size_bytes = entity_count * 100 + cross_relation_count * 200 + type_cross_relation_count * 150;
        
        IndexStats {
            entity_count,
            cross_relation_count,
            type_cross_relation_count,
            size_bytes,
            last_rebuild: None,
        }
    }
}

impl Default for CrossDbIndex {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_entity_indexing() {
        let index = CrossDbIndex::new();
        
        index.index_entity("user1", "social_graph", StoreType::Graph).unwrap();
        index.index_entity("doc1", "documents", StoreType::Document).unwrap();
        
        let loc = index.find_entity("user1").unwrap();
        assert!(loc.is_some());
        assert_eq!(loc.unwrap().store_type, StoreType::Graph);
        
        let graph_entities = index.find_by_type(StoreType::Graph).unwrap();
        assert_eq!(graph_entities.len(), 1);
    }
    
    #[test]
    fn test_cross_type_relation() {
        let index = CrossDbIndex::new();
        
        // Person (graph) -> Resume (document)
        index.index_type_cross_relation("person1", TypeCrossRelation {
            source_id: "person1".to_string(),
            source_type: StoreType::Graph,
            target_id: "resume1".to_string(),
            target_type: StoreType::Document,
            relation_type: "has_resume".to_string(),
        }).unwrap();
        
        let refs = index.find_type_cross_from("person1").unwrap();
        assert_eq!(refs.len(), 1);
        assert_eq!(refs[0].target_id, "resume1");
        
        let refs_to = index.find_refs_to("resume1").unwrap();
        assert_eq!(refs_to.len(), 1);
        assert_eq!(refs_to[0], "person1");
    }
    
    #[test]
    fn test_persist_and_restore() {
        let temp_dir = std::env::temp_dir().join("aifastdb_test_idx_persist");
        let _ = std::fs::create_dir_all(&temp_dir);
        let idx_path = temp_dir.join("cross_db_index.json");
        let _ = std::fs::remove_file(&idx_path);
        
        // Build index
        let index = CrossDbIndex::new();
        index.index_entity("user1", "social", StoreType::Graph).unwrap();
        index.index_entity("doc1", "docs", StoreType::Document).unwrap();
        index.index_type_cross_relation("user1", TypeCrossRelation {
            source_id: "user1".to_string(),
            source_type: StoreType::Graph,
            target_id: "doc1".to_string(),
            target_type: StoreType::Document,
            relation_type: "authored".to_string(),
        }).unwrap();
        
        // Persist
        index.persist(&idx_path).unwrap();
        assert!(idx_path.exists());
        
        // Restore into a new instance
        let restored = CrossDbIndex::restore(&idx_path).unwrap();
        
        let loc = restored.find_entity("user1").unwrap().unwrap();
        assert_eq!(loc.store_alias, "social");
        assert_eq!(loc.store_type, StoreType::Graph);
        
        let loc = restored.find_entity("doc1").unwrap().unwrap();
        assert_eq!(loc.store_alias, "docs");
        
        let refs = restored.find_type_cross_from("user1").unwrap();
        assert_eq!(refs.len(), 1);
        assert_eq!(refs[0].target_id, "doc1");
        
        let stats = restored.stats();
        assert_eq!(stats.entity_count, 2);
        
        // Cleanup
        let _ = std::fs::remove_dir_all(&temp_dir);
    }
    
    #[test]
    fn test_restore_or_new_missing_file() {
        let path = std::env::temp_dir().join("aifastdb_nonexistent_index.json");
        let _ = std::fs::remove_file(&path);
        
        let index = CrossDbIndex::restore_or_new(&path).unwrap();
        assert_eq!(index.stats().entity_count, 0);
    }
}
