//! Hybrid query engine combining vector, text, and metadata search
//!
//! Provides a unified interface for multi-modal queries with result fusion.

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

use serde::{Deserialize, Serialize};
use ordered_float::OrderedFloat;

use crate::types::MemoryNode;
use crate::vector::HNSWIndex;
use crate::storage::{DocumentStore, GraphStore};
use crate::storage::backend::{StorageBackend, SemanticGateFilter};
use super::filter::{Filter, FilterCondition};
use super::sifu::{SiFuQueryEngine, SiFuConfig, EdgeWeightStore};

/// A hybrid query combining multiple search modes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HybridQuery {
    /// Vector query (embedding to search for)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub vector: Option<Vec<f32>>,
    
    /// Weight for vector search results (0.0 - 1.0)
    #[serde(default = "default_vector_weight")]
    pub vector_weight: f64,
    
    /// Text query for keyword matching
    #[serde(skip_serializing_if = "Option::is_none")]
    pub text: Option<String>,
    
    /// Weight for text search results (0.0 - 1.0)
    #[serde(default = "default_text_weight")]
    pub text_weight: f64,
    
    /// Metadata filters
    #[serde(default)]
    pub filters: Vec<FilterCondition>,
    
    /// Time range filter (Unix milliseconds)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub time_range: Option<TimeRange>,
    
    /// Minimum importance threshold
    #[serde(skip_serializing_if = "Option::is_none")]
    pub min_importance: Option<f64>,
    
    /// Include graph relationships in results
    #[serde(default)]
    pub include_relations: bool,
    
    /// Maximum number of results
    #[serde(default = "default_limit")]
    pub limit: usize,
    
    /// Result offset for pagination
    #[serde(default)]
    pub offset: usize,
    
    /// Collections to search in (None = all collections)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub collections: Option<Vec<String>>,

    /// 语义门控过滤（BriLLM SiFu 启发）
    ///
    /// 在向量搜索和文本搜索之后，按语义轴范围进一步缩小结果集。
    /// 类似大脑皮层先定位语义分区再精确搜索。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub semantic_gate: Option<SemanticGateFilter>,

    /// SiFu 信号传播配置（可选，启用第三路搜索）
    ///
    /// 启用后，向量搜索结果作为种子节点注入能量，
    /// 通过图的边进行信号传播扩散，发现语义关联节点。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sifu_config: Option<SiFuConfig>,

    /// SiFu 搜索权重（0.0~1.0）
    #[serde(default)]
    pub sifu_weight: f64,
}

fn default_vector_weight() -> f64 { 0.7 }
fn default_text_weight() -> f64 { 0.3 }
fn default_limit() -> usize { 10 }

/// Time range for filtering
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeRange {
    /// Start timestamp (inclusive, Unix milliseconds)
    pub start: Option<i64>,
    /// End timestamp (inclusive, Unix milliseconds)
    pub end: Option<i64>,
}

impl TimeRange {
    /// Check if a timestamp is within this range
    pub fn contains(&self, timestamp: i64) -> bool {
        let after_start = self.start.map_or(true, |s| timestamp >= s);
        let before_end = self.end.map_or(true, |e| timestamp <= e);
        after_start && before_end
    }
}

/// A search result with score and metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResult {
    /// The memory node
    pub node: MemoryNode,
    
    /// Combined relevance score
    pub score: f64,
    
    /// Score breakdown by source
    pub score_breakdown: ScoreBreakdown,
}

/// Score breakdown for debugging and transparency
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ScoreBreakdown {
    /// Vector similarity score
    pub vector_score: Option<f64>,
    
    /// Text match score (legacy keyword matching or BM25)
    pub text_score: Option<f64>,

    /// BM25 full-text search score (from Tantivy, when available)
    pub bm25_score: Option<f64>,
    
    /// Importance score
    pub importance_score: f64,
    
    /// Recency score
    pub recency_score: f64,

    /// SiFu 信号传播能量分数（越高越相关）
    pub sifu_score: Option<f64>,
}

impl HybridQuery {
    /// Create a new query with vector search
    pub fn vector(embedding: Vec<f32>) -> Self {
        Self {
            vector: Some(embedding),
            vector_weight: 0.7,
            text: None,
            text_weight: 0.3,
            filters: Vec::new(),
            time_range: None,
            min_importance: None,
            include_relations: false,
            limit: 10,
            offset: 0,
            collections: None,
            semantic_gate: None,
            sifu_config: None,
            sifu_weight: 0.0,
        }
    }
    
    /// Create a new query with text search
    pub fn text(query: impl Into<String>) -> Self {
        Self {
            vector: None,
            vector_weight: 0.7,
            text: Some(query.into()),
            text_weight: 0.3,
            filters: Vec::new(),
            time_range: None,
            min_importance: None,
            include_relations: false,
            limit: 10,
            offset: 0,
            collections: None,
            semantic_gate: None,
            sifu_config: None,
            sifu_weight: 0.0,
        }
    }
    
    /// Create an empty query (for filter-only queries)
    pub fn empty() -> Self {
        Self {
            vector: None,
            vector_weight: 0.7,
            text: None,
            text_weight: 0.3,
            filters: Vec::new(),
            time_range: None,
            min_importance: None,
            include_relations: false,
            limit: 10,
            offset: 0,
            collections: None,
            semantic_gate: None,
            sifu_config: None,
            sifu_weight: 0.0,
        }
    }
    
    /// Set vector and weight
    pub fn with_vector(mut self, embedding: Vec<f32>, weight: f64) -> Self {
        self.vector = Some(embedding);
        self.vector_weight = weight;
        self
    }
    
    /// Set text and weight
    pub fn with_text(mut self, text: impl Into<String>, weight: f64) -> Self {
        self.text = Some(text.into());
        self.text_weight = weight;
        self
    }
    
    /// Add a filter condition
    pub fn filter(mut self, condition: FilterCondition) -> Self {
        self.filters.push(condition);
        self
    }
    
    /// Set time range
    pub fn in_time_range(mut self, start: Option<i64>, end: Option<i64>) -> Self {
        self.time_range = Some(TimeRange { start, end });
        self
    }
    
    /// Set minimum importance
    pub fn min_importance(mut self, min: f64) -> Self {
        self.min_importance = Some(min);
        self
    }
    
    /// Set result limit
    pub fn limit(mut self, limit: usize) -> Self {
        self.limit = limit;
        self
    }
    
    /// Set offset
    pub fn offset(mut self, offset: usize) -> Self {
        self.offset = offset;
        self
    }
    
    /// Include graph relations
    pub fn with_relations(mut self) -> Self {
        self.include_relations = true;
        self
    }
    
    /// Set collections to search in
    pub fn in_collections(mut self, collections: Vec<String>) -> Self {
        self.collections = Some(collections);
        self
    }
    
    /// Set a single collection to search in
    pub fn in_collection(mut self, collection: impl Into<String>) -> Self {
        self.collections = Some(vec![collection.into()]);
        self
    }

    /// 设置语义门控过滤（BriLLM SiFu 启发）
    ///
    /// 按语义轴范围缩小候选集，类似大脑皮层的"语义分区定位"。
    pub fn with_semantic_gate(mut self, gate: SemanticGateFilter) -> Self {
        self.semantic_gate = Some(gate);
        self
    }

    /// 快捷方法：按单个语义轴范围过滤
    pub fn gate_axis_0(mut self, min: f32, max: f32) -> Self {
        self.semantic_gate = Some(SemanticGateFilter::axis_0(min, max));
        self
    }

    /// 快捷方法：按三个语义轴范围过滤
    pub fn gate_all_axes(mut self, a0: (f32, f32), a1: (f32, f32), a2: (f32, f32)) -> Self {
        self.semantic_gate = Some(SemanticGateFilter::all_axes(a0, a1, a2));
        self
    }

    /// 启用 SiFu 信号传播作为第三路搜索
    ///
    /// 向量搜索结果将作为种子节点，通过图边进行信号传播扩散。
    /// `weight` 控制 SiFu 分数在最终融合中的权重。
    pub fn with_sifu(mut self, config: SiFuConfig, weight: f64) -> Self {
        self.sifu_config = Some(config);
        self.sifu_weight = weight.clamp(0.0, 1.0);
        self
    }
}

/// Hybrid query engine
pub struct HybridQueryEngine {
    /// HNSW index for vector search
    hnsw_index: Arc<HNSWIndex>,
    
    /// Document store (StorageBackend trait object)
    doc_store: Arc<dyn StorageBackend>,
    
    /// Graph store (optional)
    graph_store: Option<Arc<GraphStore>>,

    /// SiFu 信号传播引擎（可选，用于三路融合查询）
    sifu_engine: Option<Arc<SiFuQueryEngine>>,
}

impl HybridQueryEngine {
    /// Create a new hybrid query engine
    pub fn new(
        hnsw_index: Arc<HNSWIndex>,
        doc_store: Arc<DocumentStore>,
    ) -> Self {
        Self {
            hnsw_index,
            doc_store,
            graph_store: None,
            sifu_engine: None,
        }
    }

    /// Create a new hybrid query engine with a generic StorageBackend
    pub fn with_backend(
        hnsw_index: Arc<HNSWIndex>,
        doc_store: Arc<dyn StorageBackend>,
    ) -> Self {
        Self {
            hnsw_index,
            doc_store,
            graph_store: None,
            sifu_engine: None,
        }
    }
    
    /// Set the graph store
    pub fn with_graph_store(mut self, graph_store: Arc<GraphStore>) -> Self {
        self.graph_store = Some(graph_store);
        self
    }

    /// 设置 SiFu 信号传播引擎（启用三路融合查询）
    pub fn with_sifu_engine(mut self, sifu: Arc<SiFuQueryEngine>) -> Self {
        self.sifu_engine = Some(sifu);
        self
    }
    
    /// Execute a hybrid query
    pub fn query(&self, query: &HybridQuery) -> Vec<SearchResult> {
        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) = self.hnsw_index.search_with_ids(vector, k) {
                let max_dist = results.iter()
                    .map(|(_, d)| *d)
                    .fold(f32::MIN, f32::max);
                
                for (id, distance) in results {
                    // Convert distance to similarity score (0-1, higher is better)
                    let similarity = if max_dist > 0.0 {
                        1.0 - (distance / max_dist) as f64
                    } else {
                        1.0
                    };
                    
                    let weighted_score = similarity * query.vector_weight;
                    
                    let entry = candidates.entry(id.clone()).or_insert_with(|| {
                        (0.0, ScoreBreakdown::default())
                    });
                    entry.0 += weighted_score;
                    entry.1.vector_score = Some(similarity);
                }
            }
        }
        
        // Text search — BM25 via StorageBackend::full_text_search (Tantivy) or keyword fallback
        if let Some(ref text) = query.text {
            let fetch_limit = (query.limit + query.offset) * 3;
            let fts_results = self.doc_store.full_text_search(text, fetch_limit)
                .unwrap_or_default();

            let uses_bm25 = self.doc_store.supports_fts();

            for (id, raw_score) in fts_results {
                let weighted_score = raw_score * query.text_weight;

                let entry = candidates.entry(id).or_insert_with(|| {
                    (0.0, ScoreBreakdown::default())
                });
                entry.0 += weighted_score;
                if uses_bm25 {
                    entry.1.bm25_score = Some(raw_score);
                } else {
                    entry.1.text_score = Some(raw_score);
                }
            }
        }
        
        // If no vector or text search, get all documents
        if query.vector.is_none() && query.text.is_none() {
            for doc in self.doc_store.all().unwrap_or_default() {
                candidates.entry(doc.id.clone()).or_insert_with(|| {
                    (0.0, ScoreBreakdown::default())
                });
            }
        }
        
        // Apply filters and build results
        let mut results: Vec<SearchResult> = Vec::new();
        let filter = Filter {
            and: query.filters.clone(),
            or: Vec::new(),
        };
        
        for (id, (score, mut breakdown)) in candidates {
            // Get the document
            let node = match self.doc_store.get(&id) {
                Ok(Some(n)) => n,
                _ => 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;
                }
            }

            // Apply semantic gate filter (BriLLM SiFu — cortical region localization)
            if let Some(ref gate) = query.semantic_gate {
                if gate.is_active() {
                    match &node.semantic_axes {
                        Some(axes) if gate.matches(axes) => { /* pass */ }
                        _ => continue, // no axes or out of range → skip
                    }
                }
            }
            
            // Calculate recency score (higher for more recent)
            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(); // Decay over ~1 week
            breakdown.recency_score = recency_score;
            
            // Combine scores
            let final_score = score + (importance * 0.1) + (recency_score * 0.05);
            
            results.push(SearchResult {
                node,
                score: final_score,
                score_breakdown: breakdown,
            });
        }

        // === SiFu 第三路：信号传播扩散 ===
        if let (Some(ref sifu_config), Some(ref sifu_engine)) = (&query.sifu_config, &self.sifu_engine) {
            if query.sifu_weight > 0.0 && !results.is_empty() {
                // 用当前候选结果作为种子节点
                let seeds: Vec<(String, f64)> = results.iter()
                    .take(5) // top-5 作为种子
                    .map(|r| (r.node.id.clone(), r.score.clamp(0.0, 1.0)))
                    .collect();

                let sifu_results = sifu_engine.propagate(&seeds, sifu_config);

                // 将 SiFu 能量归一化后融合到结果中
                let max_energy = sifu_results.iter()
                    .map(|r| r.energy)
                    .fold(0.0f64, f64::max);

                if max_energy > 0.0 {
                    let sifu_map: HashMap<String, f64> = sifu_results.into_iter()
                        .map(|r| (r.node_id, r.energy / max_energy))
                        .collect();

                    for result in &mut results {
                        if let Some(&norm_energy) = sifu_map.get(&result.node.id) {
                            result.score_breakdown.sifu_score = Some(norm_energy);
                            result.score += norm_energy * query.sifu_weight;
                        }
                    }

                    // SiFu 可能发现新节点（不在原始候选中）
                    for (node_id, norm_energy) in &sifu_map {
                        if results.iter().any(|r| r.node.id == *node_id) {
                            continue; // 已经在结果中
                        }
                        if let Ok(Some(node)) = self.doc_store.get(node_id) {
                            let mut breakdown = ScoreBreakdown::default();
                            breakdown.sifu_score = Some(*norm_energy);
                            results.push(SearchResult {
                                node,
                                score: norm_energy * query.sifu_weight,
                                score_breakdown: breakdown,
                            });
                        }
                    }
                }
            }
        }
        
        // 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 with RRF (Reciprocal Rank Fusion)
    pub fn query_rrf(&self, query: &HybridQuery, k: f64) -> Vec<SearchResult> {
        let mut rrf_scores: HashMap<String, f64> = HashMap::new();
        
        // Vector search ranking
        if let Some(ref vector) = query.vector {
            let num_results = (query.limit + query.offset) * 3;
            
            if let Ok(results) = self.hnsw_index.search_with_ids(vector, num_results) {
                for (rank, (id, _)) in results.iter().enumerate() {
                    let rrf_score = 1.0 / (k + rank as f64 + 1.0);
                    *rrf_scores.entry(id.clone()).or_default() += rrf_score * query.vector_weight;
                }
            }
        }
        
        // Text search ranking — BM25 via StorageBackend::full_text_search or keyword fallback
        if let Some(ref text) = query.text {
            let fetch_limit = (query.limit + query.offset) * 3;
            let mut text_results = self.doc_store.full_text_search(text, fetch_limit)
                .unwrap_or_default();

            // Sort by score for RRF ranking
            text_results.sort_by(|a, b| OrderedFloat(b.1).cmp(&OrderedFloat(a.1)));
            
            for (rank, (id, _)) in text_results.iter().enumerate() {
                let rrf_score = 1.0 / (k + rank as f64 + 1.0);
                *rrf_scores.entry(id.clone()).or_default() += rrf_score * query.text_weight;
            }
        }
        
        // Build results with RRF scores
        let filter = Filter {
            and: query.filters.clone(),
            or: Vec::new(),
        };
        
        let mut results: Vec<SearchResult> = Vec::new();
        
        for (id, rrf_score) in rrf_scores {
            let node = match self.doc_store.get(&id) {
                Ok(Some(n)) => n,
                _ => continue,
            };
            
            // Apply filters (same as regular query)
            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;
                }
            }

            // Semantic gate filter (BriLLM SiFu)
            if let Some(ref gate) = query.semantic_gate {
                if gate.is_active() {
                    match &node.semantic_axes {
                        Some(axes) if gate.matches(axes) => { /* pass */ }
                        _ => continue,
                    }
                }
            }
            
            results.push(SearchResult {
                node,
                score: rrf_score,
                score_breakdown: ScoreBreakdown::default(),
            });
        }

        // === SiFu 第三路 RRF 融合 ===
        if let (Some(ref sifu_config), Some(ref sifu_engine)) = (&query.sifu_config, &self.sifu_engine) {
            if query.sifu_weight > 0.0 && !results.is_empty() {
                let seeds: Vec<(String, f64)> = results.iter()
                    .take(5)
                    .map(|r| (r.node.id.clone(), r.score.clamp(0.0, 1.0)))
                    .collect();

                let sifu_results = sifu_engine.propagate(&seeds, sifu_config);

                // SiFu 结果也用 RRF 排名融合
                for (rank, sr) in sifu_results.iter().enumerate() {
                    let sifu_rrf = 1.0 / (k + rank as f64 + 1.0);
                    let weighted = sifu_rrf * query.sifu_weight;

                    if let Some(existing) = results.iter_mut().find(|r| r.node.id == sr.node_id) {
                        existing.score += weighted;
                        existing.score_breakdown.sifu_score = Some(sr.energy);
                    } else if let Ok(Some(node)) = self.doc_store.get(&sr.node_id) {
                        let mut breakdown = ScoreBreakdown::default();
                        breakdown.sifu_score = Some(sr.energy);
                        results.push(SearchResult {
                            node,
                            score: weighted,
                            score_breakdown: breakdown,
                        });
                    }
                }
            }
        }
        
        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 tempfile::tempdir;
    use crate::types::HNSWConfig;

    fn setup_test_engine() -> (HybridQueryEngine, Arc<HNSWIndex>, Arc<DocumentStore>) {
        let dir = tempdir().unwrap();
        
        let hnsw_config = HNSWConfig::new(4)
            .with_m(4)
            .with_ef_construction(16)
            .with_ef_search(16)
            .with_max_elements(100);
        
        let hnsw = Arc::new(HNSWIndex::new(hnsw_config));
        let doc_store = Arc::new(
            DocumentStore::open_or_create(dir.path().join("docs.jsonl")).unwrap()
        );
        
        let engine = HybridQueryEngine::new(hnsw.clone(), doc_store.clone());
        
        (engine, hnsw, doc_store)
    }

    #[test]
    fn test_hybrid_query_builder() {
        let query = HybridQuery::vector(vec![1.0, 2.0, 3.0, 4.0])
            .with_text("search term", 0.3)
            .filter(FilterCondition::eq("status", "active"))
            .min_importance(0.5)
            .limit(20);
        
        assert!(query.vector.is_some());
        assert_eq!(query.text, Some("search term".to_string()));
        assert_eq!(query.filters.len(), 1);
        assert_eq!(query.min_importance, Some(0.5));
        assert_eq!(query.limit, 20);
    }

    #[test]
    fn test_text_search() {
        let (engine, _, doc_store) = setup_test_engine();
        
        // Add some documents
        doc_store.put(MemoryNode::new("The quick brown fox jumps over the lazy dog")).unwrap();
        doc_store.put(MemoryNode::new("A quick test document")).unwrap();
        doc_store.put(MemoryNode::new("Something completely different")).unwrap();
        
        let query = HybridQuery::text("quick");
        let results = engine.query(&query);
        
        assert_eq!(results.len(), 2);
        assert!(results[0].node.content.contains("quick"));
    }

    #[test]
    fn test_empty_query_with_filter() {
        let (engine, _, doc_store) = setup_test_engine();
        
        let mut node1 = MemoryNode::new("Important document");
        node1.metadata.insert("status".to_string(), serde_json::json!("active"));
        
        let mut node2 = MemoryNode::new("Archived document");
        node2.metadata.insert("status".to_string(), serde_json::json!("archived"));
        
        doc_store.put(node1).unwrap();
        doc_store.put(node2).unwrap();
        
        let query = HybridQuery::empty()
            .filter(FilterCondition::eq("status", "active"));
        
        let results = engine.query(&query);
        
        assert_eq!(results.len(), 1);
        assert!(results[0].node.content.contains("Important"));
    }

    #[test]
    fn test_time_range_filter() {
        let (engine, _, doc_store) = setup_test_engine();
        
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as i64;
        
        let node = MemoryNode::new("Recent document");
        doc_store.put(node).unwrap();
        
        // Query with time range that includes now
        let query = HybridQuery::empty()
            .in_time_range(Some(now - 1000), Some(now + 1000));
        
        let results = engine.query(&query);
        assert_eq!(results.len(), 1);
        
        // Query with time range that excludes now
        let query = HybridQuery::empty()
            .in_time_range(Some(now + 1000), None);
        
        let results = engine.query(&query);
        assert_eq!(results.len(), 0);
    }
}
