//! Multi-collection query engine
//!
//! Provides query capabilities across multiple collections with result merging.

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

use ordered_float::OrderedFloat;

use crate::storage::Collection;
use crate::types::MemoryNode;
use super::hybrid::{HybridQuery, SearchResult, ScoreBreakdown};
use super::filter::Filter;

/// Multi-collection query engine
/// 
/// Enables searching across multiple collections and merging results.
pub struct MultiCollectionQueryEngine {
    /// Collections to search
    collections: Vec<Arc<Collection>>,
}

impl MultiCollectionQueryEngine {
    /// Create a new multi-collection query engine
    pub fn new(collections: Vec<Arc<Collection>>) -> Self {
        Self { collections }
    }
    
    /// Execute a query across all collections
    pub fn query(&self, query: &HybridQuery) -> Vec<SearchResult> {
        let mut all_candidates: HashMap<String, (f64, ScoreBreakdown, MemoryNode, String)> = HashMap::new();
        
        // Determine which collections to search
        let collections_to_search: Vec<&Arc<Collection>> = match &query.collections {
            Some(names) => {
                self.collections.iter()
                    .filter(|c| names.contains(&c.name))
                    .collect()
            }
            None => self.collections.iter().collect(),
        };
        
        // Search each collection
        for collection in collections_to_search {
            let candidates = self.search_collection(collection, query);
            
            for (id, (score, breakdown, node)) in candidates {
                // Use collection-prefixed ID to avoid conflicts
                let unique_id = format!("{}:{}", collection.name, id);
                
                // If same document found in multiple collections, keep the highest score
                let entry = all_candidates.entry(unique_id).or_insert_with(|| {
                    (0.0, ScoreBreakdown::default(), node.clone(), collection.name.clone())
                });
                
                if score > entry.0 {
                    *entry = (score, breakdown, node, collection.name.clone());
                }
            }
        }
        
        // Build results
        let mut results: Vec<SearchResult> = all_candidates
            .into_iter()
            .map(|(_, (score, breakdown, mut node, collection_name))| {
                // Add collection name to node metadata
                node.metadata.insert(
                    "_collection".to_string(), 
                    serde_json::json!(collection_name)
                );
                
                SearchResult {
                    node,
                    score,
                    score_breakdown: breakdown,
                }
            })
            .collect();
        
        // Sort by score descending
        results.sort_by(|a, b| {
            OrderedFloat(b.score).cmp(&OrderedFloat(a.score))
        });
        
        // Apply pagination
        results.into_iter()
            .skip(query.offset)
            .take(query.limit)
            .collect()
    }
    
    /// Search a single collection
    fn search_collection(
        &self,
        collection: &Collection,
        query: &HybridQuery,
    ) -> HashMap<String, (f64, ScoreBreakdown, MemoryNode)> {
        let mut candidates: HashMap<String, (f64, ScoreBreakdown)> = HashMap::new();
        
        // Vector search
        if let Some(ref vector) = query.vector {
            let k = (query.limit + query.offset) * 3; // Over-fetch for filtering
            
            if let Ok(results) = collection.vector_engine.search(vector, k) {
                let max_dist = results.iter()
                    .map(|r| r.distance)
                    .fold(f32::MIN, f32::max);
                
                for result in results {
                    let similarity = if max_dist > 0.0 {
                        1.0 - (result.distance / max_dist) as f64
                    } else {
                        1.0
                    };
                    
                    let weighted_score = similarity * query.vector_weight;
                    
                    let entry = candidates.entry(result.id.clone()).or_insert_with(|| {
                        (0.0, ScoreBreakdown::default())
                    });
                    entry.0 += weighted_score;
                    entry.1.vector_score = Some(similarity);
                }
            }
        }
        
        // Text search
        if let Some(ref text) = query.text {
            let text_lower = text.to_lowercase();
            let keywords: Vec<&str> = text_lower.split_whitespace().collect();
            
            for doc in collection.docs.all().unwrap_or_default() {
                let content_lower = doc.content.to_lowercase();
                
                let matches: usize = keywords.iter()
                    .filter(|kw| content_lower.contains(*kw))
                    .count();
                
                if matches > 0 {
                    let text_score = matches as f64 / keywords.len() as f64;
                    let weighted_score = text_score * query.text_weight;
                    
                    let entry = candidates.entry(doc.id.clone()).or_insert_with(|| {
                        (0.0, ScoreBreakdown::default())
                    });
                    entry.0 += weighted_score;
                    entry.1.text_score = Some(text_score);
                }
            }
        }
        
        // If no vector or text search, get all documents
        if query.vector.is_none() && query.text.is_none() {
            for doc in collection.docs.all().unwrap_or_default() {
                candidates.entry(doc.id.clone()).or_insert_with(|| {
                    (0.0, ScoreBreakdown::default())
                });
            }
        }
        
        // Apply filters and build final results
        let filter = Filter {
            and: query.filters.clone(),
            or: Vec::new(),
        };
        
        let mut results: HashMap<String, (f64, ScoreBreakdown, MemoryNode)> = HashMap::new();
        
        for (id, (score, mut breakdown)) in candidates {
            let node = match collection.docs.get(&id).unwrap_or(None) {
                Some(n) => n,
                None => continue,
            };
            
            // Apply metadata filter
            if !filter.is_empty() {
                let metadata_json = serde_json::to_value(&node.metadata).unwrap_or_default();
                if !filter.evaluate(&metadata_json) {
                    continue;
                }
            }
            
            // Apply time range filter
            if let Some(ref range) = query.time_range {
                if !range.contains(node.created_at) {
                    continue;
                }
            }
            
            // Apply importance filter
            let importance = node.computed_importance();
            breakdown.importance_score = importance;
            
            if let Some(min) = query.min_importance {
                if importance < min {
                    continue;
                }
            }
            
            // Calculate recency score
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis() as i64;
            let age_hours = (now - node.last_accessed_at) as f64 / (1000.0 * 60.0 * 60.0);
            let recency_score = (-age_hours / 168.0).exp();
            breakdown.recency_score = recency_score;
            
            // Combine scores
            let final_score = score + (importance * 0.1) + (recency_score * 0.05);
            
            results.insert(id, (final_score, breakdown, node));
        }
        
        results
    }
    
    /// Search with RRF (Reciprocal Rank Fusion) across collections
    pub fn query_rrf(&self, query: &HybridQuery, k: f64) -> Vec<SearchResult> {
        let mut rrf_scores: HashMap<String, (f64, MemoryNode, String)> = HashMap::new();
        
        // Determine which collections to search
        let collections_to_search: Vec<&Arc<Collection>> = match &query.collections {
            Some(names) => {
                self.collections.iter()
                    .filter(|c| names.contains(&c.name))
                    .collect()
            }
            None => self.collections.iter().collect(),
        };
        
        for collection in collections_to_search {
            // Vector search ranking
            if let Some(ref vector) = query.vector {
                let num_results = (query.limit + query.offset) * 3;
                
                if let Ok(results) = collection.vector_engine.search(vector, num_results) {
                    for (rank, result) in results.iter().enumerate() {
                        let rrf_score = 1.0 / (k + rank as f64 + 1.0);
                        let unique_id = format!("{}:{}", collection.name, result.id);
                        
                        if let Ok(Some(node)) = collection.docs.get(&result.id) {
                            let entry = rrf_scores.entry(unique_id).or_insert_with(|| {
                                (0.0, node.clone(), collection.name.clone())
                            });
                            entry.0 += rrf_score * query.vector_weight;
                        }
                    }
                }
            }
            
            // Text search ranking
            if let Some(ref text) = query.text {
                let text_lower = text.to_lowercase();
                let keywords: Vec<&str> = text_lower.split_whitespace().collect();
                
                let mut text_results: Vec<(String, f64, MemoryNode)> = Vec::new();
                
                for doc in collection.docs.all().unwrap_or_default() {
                    let content_lower = doc.content.to_lowercase();
                    let matches: usize = keywords.iter()
                        .filter(|kw| content_lower.contains(*kw))
                        .count();
                    
                    if matches > 0 {
                        let score = matches as f64 / keywords.len() as f64;
                        text_results.push((doc.id.clone(), score, doc));
                    }
                }
                
                text_results.sort_by(|a, b| OrderedFloat(b.1).cmp(&OrderedFloat(a.1)));
                
                for (rank, (id, _, node)) in text_results.iter().enumerate() {
                    let rrf_score = 1.0 / (k + rank as f64 + 1.0);
                    let unique_id = format!("{}:{}", collection.name, id);
                    
                    let entry = rrf_scores.entry(unique_id).or_insert_with(|| {
                        (0.0, node.clone(), collection.name.clone())
                    });
                    entry.0 += rrf_score * query.text_weight;
                }
            }
        }
        
        // Apply filters and build results
        let filter = Filter {
            and: query.filters.clone(),
            or: Vec::new(),
        };
        
        let mut results: Vec<SearchResult> = Vec::new();
        
        for (_, (rrf_score, mut node, collection_name)) in rrf_scores {
            // Apply filters
            if !filter.is_empty() {
                let metadata_json = serde_json::to_value(&node.metadata).unwrap_or_default();
                if !filter.evaluate(&metadata_json) {
                    continue;
                }
            }
            
            if let Some(ref range) = query.time_range {
                if !range.contains(node.created_at) {
                    continue;
                }
            }
            
            let importance = node.computed_importance();
            if let Some(min) = query.min_importance {
                if importance < min {
                    continue;
                }
            }
            
            // Add collection name to metadata
            node.metadata.insert(
                "_collection".to_string(),
                serde_json::json!(collection_name)
            );
            
            results.push(SearchResult {
                node,
                score: rrf_score,
                score_breakdown: ScoreBreakdown::default(),
            });
        }
        
        results.sort_by(|a, b| OrderedFloat(b.score).cmp(&OrderedFloat(a.score)));
        
        results.into_iter()
            .skip(query.offset)
            .take(query.limit)
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::CollectionConfig;
    use tempfile::tempdir;

    #[test]
    fn test_multi_collection_query() {
        let dir = tempdir().unwrap();
        
        // Create two collections
        let poems_path = dir.path().join("poems");
        let poems_config = CollectionConfig::new("poems", 4);
        let poems = Arc::new(Collection::open_or_create(&poems_path, poems_config, None).unwrap());
        
        let articles_path = dir.path().join("articles");
        let articles_config = CollectionConfig::new("articles", 4);
        let articles = Arc::new(Collection::open_or_create(&articles_path, articles_config, None).unwrap());
        
        // Add documents
        let poem1 = MemoryNode::new("静夜思 床前明月光");
        poems.docs.put(poem1.clone()).unwrap();
        poems.hnsw.insert(&poem1.id, vec![1.0, 0.0, 0.0, 0.0]).unwrap();
        
        let article1 = MemoryNode::new("月球探测技术发展");
        articles.docs.put(article1.clone()).unwrap();
        articles.hnsw.insert(&article1.id, vec![0.9, 0.1, 0.0, 0.0]).unwrap();
        
        // Create engine
        let engine = MultiCollectionQueryEngine::new(vec![poems.clone(), articles.clone()]);
        
        // Search all collections
        let query = HybridQuery::text("月");
        let results = engine.query(&query);
        
        assert_eq!(results.len(), 2);
        
        // Search specific collection
        let query = HybridQuery::text("月").in_collection("poems");
        let results = engine.query(&query);
        
        assert_eq!(results.len(), 1);
        assert!(results[0].node.content.contains("静夜思"));
    }
}
