//! Reference Resolver - Resolves cross-store and cross-type references
//!
//! Provides utilities for resolving entity references across different
//! stores and storage types in the federation.

use std::collections::HashMap;
use std::sync::Arc;

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

/// Reference resolution mode
#[derive(Debug, Clone, Copy)]
pub enum ResolutionMode {
    /// Resolve lazily (on demand)
    Lazy,
    /// Resolve eagerly (prefetch)
    Eager,
    /// Resolve in batch
    Batch,
}

/// Reference resolution options
#[derive(Debug, Clone)]
pub struct ResolutionOptions {
    /// Resolution mode
    pub mode: ResolutionMode,
    /// Maximum depth for recursive resolution
    pub max_depth: u32,
    /// Include reverse references
    pub include_reverse: bool,
    /// Filter by store types
    pub store_type_filter: Option<Vec<StoreType>>,
}

impl Default for ResolutionOptions {
    fn default() -> Self {
        Self {
            mode: ResolutionMode::Lazy,
            max_depth: 1,
            include_reverse: false,
            store_type_filter: None,
        }
    }
}

/// Resolved reference with full entity data
#[derive(Debug, Clone)]
pub struct ResolvedRef {
    /// Source entity ID
    pub source_id: String,
    /// Target entity ID
    pub target_id: String,
    /// Relation type
    pub relation_type: String,
    /// Source entity (if resolved)
    pub source: Option<Record>,
    /// Target entity (if resolved)
    pub target: Option<Record>,
    /// Source location
    pub source_location: Option<EntityLocation>,
    /// Target location
    pub target_location: Option<EntityLocation>,
}

/// Reference resolver for cross-store entity references
pub struct RefResolver {
    /// Cross-database index
    index: Arc<CrossDbIndex>,
    /// Registered stores
    stores: HashMap<String, Arc<dyn Store>>,
}

impl RefResolver {
    /// Create a new reference resolver
    pub fn new(index: Arc<CrossDbIndex>) -> Self {
        Self {
            index,
            stores: HashMap::new(),
        }
    }
    
    /// Register a store
    pub fn register_store(&mut self, alias: &str, store: Arc<dyn Store>) {
        self.stores.insert(alias.to_string(), store);
    }
    
    /// Resolve an entity ID to its full record
    pub fn resolve_entity(&self, entity_id: &str) -> Result<Option<Record>> {
        // First, find the entity location
        let location = self.index.find_entity(entity_id)?;
        
        if let Some(loc) = location {
            // Get the store
            let store = self.stores.get(&loc.store_alias)
                .ok_or_else(|| FederationError::StoreNotFound(loc.store_alias.clone()))?;
            
            // Fetch the entity
            store.get(entity_id)
        } else {
            // Try all stores
            for store in self.stores.values() {
                if let Ok(Some(record)) = store.get(entity_id) {
                    return Ok(Some(record));
                }
            }
            Ok(None)
        }
    }
    
    /// Resolve multiple entities in batch
    pub fn resolve_batch(&self, entity_ids: &[String]) -> Result<HashMap<String, Record>> {
        let mut results = HashMap::new();
        
        // Group by store for efficient batch fetching
        let mut by_store: HashMap<String, Vec<String>> = HashMap::new();
        
        for id in entity_ids {
            if let Ok(Some(loc)) = self.index.find_entity(id) {
                by_store.entry(loc.store_alias)
                    .or_default()
                    .push(id.clone());
            } else {
                // Unknown location - will try all stores
                by_store.entry("__unknown__".to_string())
                    .or_default()
                    .push(id.clone());
            }
        }
        
        // Fetch from known stores
        for (store_alias, ids) in &by_store {
            if store_alias == "__unknown__" {
                continue;
            }
            
            if let Some(store) = self.stores.get(store_alias) {
                for id in ids {
                    if let Ok(Some(record)) = store.get(id) {
                        results.insert(id.clone(), record);
                    }
                }
            }
        }
        
        // Try unknown entities in all stores
        if let Some(unknown_ids) = by_store.get("__unknown__") {
            for id in unknown_ids {
                if results.contains_key(id) {
                    continue;
                }
                
                for store in self.stores.values() {
                    if let Ok(Some(record)) = store.get(id) {
                        results.insert(id.clone(), record);
                        break;
                    }
                }
            }
        }
        
        Ok(results)
    }
    
    /// Resolve references from an entity
    pub fn resolve_refs_from(
        &self,
        entity_id: &str,
        options: &ResolutionOptions,
    ) -> Result<Vec<ResolvedRef>> {
        let mut results = Vec::new();
        
        // Get cross-store references
        let ref_ids = self.index.find_refs_from(entity_id)?;
        
        // Get cross-type references
        let type_refs = self.index.find_type_cross_from(entity_id)?;
        
        // Resolve cross-store refs
        for target_id in ref_ids {
            let mut resolved = ResolvedRef {
                source_id: entity_id.to_string(),
                target_id: target_id.clone(),
                relation_type: "ref".to_string(), // TODO: Get actual relation type
                source: None,
                target: None,
                source_location: None,
                target_location: None,
            };
            
            // Resolve entities based on mode
            if matches!(options.mode, ResolutionMode::Eager) {
                resolved.source = self.resolve_entity(entity_id)?;
                resolved.target = self.resolve_entity(&target_id)?;
            }
            
            // Get locations
            resolved.source_location = self.index.find_entity(entity_id)?;
            resolved.target_location = self.index.find_entity(&target_id)?;
            
            // Apply filter
            if let Some(filter) = &options.store_type_filter {
                if let Some(loc) = &resolved.target_location {
                    if !filter.contains(&loc.store_type) {
                        continue;
                    }
                }
            }
            
            results.push(resolved);
        }
        
        // Resolve cross-type refs
        for type_ref in type_refs {
            let mut resolved = ResolvedRef {
                source_id: type_ref.source_id.clone(),
                target_id: type_ref.target_id.clone(),
                relation_type: type_ref.relation_type.clone(),
                source: None,
                target: None,
                source_location: Some(EntityLocation {
                    entity_id: type_ref.source_id.clone(),
                    store_alias: String::new(), // TODO: Get from index
                    store_type: type_ref.source_type,
                }),
                target_location: Some(EntityLocation {
                    entity_id: type_ref.target_id.clone(),
                    store_alias: String::new(), // TODO: Get from index
                    store_type: type_ref.target_type,
                }),
            };
            
            // Resolve entities based on mode
            if matches!(options.mode, ResolutionMode::Eager) {
                resolved.source = self.resolve_entity(&type_ref.source_id)?;
                resolved.target = self.resolve_entity(&type_ref.target_id)?;
            }
            
            // Apply filter
            if let Some(filter) = &options.store_type_filter {
                if !filter.contains(&type_ref.target_type) {
                    continue;
                }
            }
            
            results.push(resolved);
        }
        
        Ok(results)
    }
    
    /// Resolve references to an entity
    pub fn resolve_refs_to(
        &self,
        entity_id: &str,
        options: &ResolutionOptions,
    ) -> Result<Vec<ResolvedRef>> {
        let mut results = Vec::new();
        
        // Get entities that reference this entity
        let source_ids = self.index.find_refs_to(entity_id)?;
        
        for source_id in source_ids {
            let mut resolved = ResolvedRef {
                source_id: source_id.clone(),
                target_id: entity_id.to_string(),
                relation_type: "ref".to_string(),
                source: None,
                target: None,
                source_location: None,
                target_location: None,
            };
            
            // Resolve entities based on mode
            if matches!(options.mode, ResolutionMode::Eager) {
                resolved.source = self.resolve_entity(&source_id)?;
                resolved.target = self.resolve_entity(entity_id)?;
            }
            
            // Get locations
            resolved.source_location = self.index.find_entity(&source_id)?;
            resolved.target_location = self.index.find_entity(entity_id)?;
            
            // Apply filter
            if let Some(filter) = &options.store_type_filter {
                if let Some(loc) = &resolved.source_location {
                    if !filter.contains(&loc.store_type) {
                        continue;
                    }
                }
            }
            
            results.push(resolved);
        }
        
        Ok(results)
    }
    
    /// Resolve all references for an entity (bidirectional)
    pub fn resolve_all_refs(
        &self,
        entity_id: &str,
        options: &ResolutionOptions,
    ) -> Result<Vec<ResolvedRef>> {
        let mut results = self.resolve_refs_from(entity_id, options)?;
        
        if options.include_reverse {
            let reverse = self.resolve_refs_to(entity_id, options)?;
            results.extend(reverse);
        }
        
        Ok(results)
    }
    
    /// Build a reference graph starting from an entity
    pub fn build_ref_graph(
        &self,
        start_id: &str,
        max_depth: u32,
    ) -> Result<RefGraph> {
        let mut graph = RefGraph::new();
        let mut visited: std::collections::HashSet<String> = std::collections::HashSet::new();
        let mut queue: std::collections::VecDeque<(String, u32)> = std::collections::VecDeque::new();
        
        queue.push_back((start_id.to_string(), 0));
        
        while let Some((entity_id, depth)) = queue.pop_front() {
            if depth > max_depth || visited.contains(&entity_id) {
                continue;
            }
            
            visited.insert(entity_id.clone());
            
            // Add node
            if let Ok(Some(record)) = self.resolve_entity(&entity_id) {
                graph.add_node(entity_id.clone(), record);
            }
            
            // Get outgoing refs
            let options = ResolutionOptions {
                mode: ResolutionMode::Lazy,
                max_depth: 1,
                include_reverse: false,
                store_type_filter: None,
            };
            
            if let Ok(refs) = self.resolve_refs_from(&entity_id, &options) {
                for r in refs {
                    graph.add_edge(r.source_id.clone(), r.target_id.clone(), r.relation_type.clone());
                    
                    if !visited.contains(&r.target_id) {
                        queue.push_back((r.target_id, depth + 1));
                    }
                }
            }
        }
        
        Ok(graph)
    }
}

/// Reference graph for visualization and analysis
#[derive(Debug, Default)]
pub struct RefGraph {
    /// Nodes (entity_id -> record)
    nodes: HashMap<String, Record>,
    /// Edges (source, target, relation_type)
    edges: Vec<(String, String, String)>,
}

impl RefGraph {
    /// Create a new empty graph
    pub fn new() -> Self {
        Self::default()
    }
    
    /// Add a node to the graph
    pub fn add_node(&mut self, id: String, record: Record) {
        self.nodes.insert(id, record);
    }
    
    /// Add an edge to the graph
    pub fn add_edge(&mut self, source: String, target: String, relation_type: String) {
        self.edges.push((source, target, relation_type));
    }
    
    /// Get node count
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }
    
    /// Get edge count
    pub fn edge_count(&self) -> usize {
        self.edges.len()
    }
    
    /// Get all nodes
    pub fn nodes(&self) -> &HashMap<String, Record> {
        &self.nodes
    }
    
    /// Get all edges
    pub fn edges(&self) -> &[(String, String, String)] {
        &self.edges
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_ref_graph() {
        let mut graph = RefGraph::new();
        
        graph.add_node("node1".to_string(), Record::Document {
            id: "node1".to_string(),
            doc_type: "test".to_string(),
            content: serde_json::json!({}),
            tags: vec![],
            created_at: 0,
            updated_at: 0,
            ownership: None,
        });
        
        graph.add_edge("node1".to_string(), "node2".to_string(), "relates_to".to_string());
        
        assert_eq!(graph.node_count(), 1);
        assert_eq!(graph.edge_count(), 1);
    }
}
