//! HybridRetriever — RAG 检索器实现
//!
//! 将现有的 `HybridQueryEngine` 包装为 `Retriever` trait 实现，
//! 支持向量 + BM25 全文混合检索 + 可选 Reranker 精排。
//!
//! # Architecture
//!
//! ```text
//! query text + embedding
//!       │
//!       ▼
//! ┌─────────────────────┐
//! │  HybridQueryEngine  │  ← vector (HNSW) + BM25 (Tantivy) + metadata filters
//! └────────┬────────────┘
//!          │ Vec<SearchResult>
//!          ▼
//! ┌─────────────────────┐
//! │  Reranker (optional) │  ← RRF / Linear / CrossEncoder
//! └────────┬────────────┘
//!          │ Vec<RankedResult>
//!          ▼
//!   Vec<RetrievalResult>  → returned to RAG pipeline
//! ```

use std::sync::Arc;

use crate::query::hybrid::{HybridQuery, HybridQueryEngine, SearchResult};
use crate::query::reranker::{Reranker, RerankCandidate, RRFReranker, LinearReranker, CrossEncoderReranker};

use super::error::RagResult;
use super::traits::Retriever;
use super::types::{RagConfig, RetrievalResult};

// ============================================================================
// HybridRetriever
// ============================================================================

/// A Retriever that wraps `HybridQueryEngine` + optional `Reranker`.
///
/// Performs hybrid search (vector + BM25 + metadata) and optionally
/// re-ranks results using a configurable reranker strategy.
pub struct HybridRetriever {
    /// The underlying hybrid query engine
    engine: Arc<HybridQueryEngine>,
    /// Optional reranker (created lazily based on config)
    reranker: Option<Box<dyn Reranker>>,
}

impl HybridRetriever {
    /// Create a new HybridRetriever with a query engine.
    pub fn new(engine: Arc<HybridQueryEngine>) -> Self {
        Self {
            engine,
            reranker: None,
        }
    }

    /// Set a custom reranker.
    pub fn with_reranker(mut self, reranker: Box<dyn Reranker>) -> Self {
        self.reranker = Some(reranker);
        self
    }

    /// Build a `HybridQuery` from RAG config and optional embedding.
    fn build_query(
        query_text: &str,
        embedding: Option<&[f32]>,
        config: &RagConfig,
    ) -> HybridQuery {
        let mut hq = if let Some(emb) = embedding {
            HybridQuery::vector(emb.to_vec())
                .with_text(query_text, config.bm25_weight)
        } else {
            HybridQuery::text(query_text)
        };

        // Apply weights
        hq.vector_weight = config.vector_weight;
        hq.text_weight = config.bm25_weight;

        // Set limit (over-fetch for reranking)
        let fetch_k = if config.enable_reranker {
            config.top_k * 3 // Fetch more for reranking
        } else {
            config.top_k
        };
        hq.limit = fetch_k;

        // Set collections if specified
        if let Some(ref cols) = config.collections {
            hq.collections = Some(cols.clone());
        }

        hq
    }

    /// Convert a `SearchResult` to a `RetrievalResult`.
    fn to_retrieval_result(sr: &SearchResult) -> RetrievalResult {
        let mut metadata = std::collections::HashMap::new();

        // Copy importance and recency as metadata
        metadata.insert(
            "importance".to_string(),
            serde_json::json!(sr.node.importance),
        );

        // Copy node metadata if available
        for (k, v) in &sr.node.metadata {
            metadata.insert(k.clone(), v.clone());
        }

        let mut score_details = std::collections::HashMap::new();
        if let Some(vs) = sr.score_breakdown.vector_score {
            score_details.insert("vector".to_string(), vs);
        }
        if let Some(bs) = sr.score_breakdown.bm25_score {
            score_details.insert("bm25".to_string(), bs);
        }
        if let Some(ts) = sr.score_breakdown.text_score {
            score_details.insert("text".to_string(), ts);
        }
        if let Some(ss) = sr.score_breakdown.sifu_score {
            score_details.insert("sifu".to_string(), ss);
        }

        RetrievalResult {
            id: sr.node.id.clone(),
            content: sr.node.content.clone(),
            score: sr.score,
            collection: None, // set by caller if needed
            score_details,
            metadata,
        }
    }

    /// Create a reranker from config strategy name.
    fn create_reranker(strategy: &str) -> Box<dyn Reranker> {
        match strategy {
            "linear" => Box::new(LinearReranker::default()),
            "cross-encoder" => Box::new(CrossEncoderReranker::default()),
            _ => Box::new(RRFReranker::default()), // "rrf" is default
        }
    }

    /// Convert search results to reranker candidates.
    fn to_rerank_candidates(results: &[SearchResult]) -> Vec<RerankCandidate> {
        results
            .iter()
            .map(|sr| {
                let mut c = RerankCandidate::new(&sr.node.id, &sr.node.content)
                    .with_importance(sr.score_breakdown.importance_score)
                    .with_recency(sr.score_breakdown.recency_score);

                if let Some(vs) = sr.score_breakdown.vector_score {
                    c = c.with_vector_score(vs);
                }
                if let Some(bs) = sr.score_breakdown.bm25_score {
                    c = c.with_bm25_score(bs);
                }
                if let Some(ts) = sr.score_breakdown.text_score {
                    c = c.with_text_score(ts);
                }
                if let Some(ss) = sr.score_breakdown.sifu_score {
                    c = c.with_sifu_score(ss);
                }
                c
            })
            .collect()
    }
}

impl Retriever for HybridRetriever {
    fn retrieve(
        &self,
        query: &str,
        embedding: Option<&[f32]>,
        config: &RagConfig,
    ) -> RagResult<Vec<RetrievalResult>> {
        // 1. Build hybrid query
        let hq = Self::build_query(query, embedding, config);

        // 2. Execute search
        let search_results = self.engine.query(&hq);

        if search_results.is_empty() {
            return Ok(Vec::new());
        }

        // 3. Optionally rerank
        let final_results: Vec<RetrievalResult> = if config.enable_reranker {
            let reranker = self
                .reranker
                .as_ref()
                .map(|r| r.as_ref());

            let candidates = Self::to_rerank_candidates(&search_results);

            let ranked = if let Some(rr) = reranker {
                rr.rerank(query, candidates)?
            } else {
                // Use config-specified strategy
                let rr = Self::create_reranker(&config.reranker_strategy);
                rr.rerank(query, candidates)?
            };

            // Map ranked results back to RetrievalResult
            let result_map: std::collections::HashMap<String, &SearchResult> = search_results
                .iter()
                .map(|sr| (sr.node.id.clone(), sr))
                .collect();

            ranked
                .iter()
                .take(config.top_k)
                .filter_map(|rr| {
                    result_map.get(&rr.id).map(|sr| {
                        let mut ret = Self::to_retrieval_result(sr);
                        ret.score = rr.score; // Use reranked score
                        // Merge reranker score details
                        for (k, v) in &rr.score_details {
                            ret.score_details.insert(format!("rerank_{}", k), *v);
                        }
                        ret
                    })
                })
                .collect()
        } else {
            // No reranking — convert directly
            search_results
                .iter()
                .take(config.top_k)
                .map(Self::to_retrieval_result)
                .collect()
        };

        // 4. Apply min_score filter
        let filtered = if let Some(min_score) = config.min_score {
            final_results
                .into_iter()
                .filter(|r| r.score >= min_score)
                .collect()
        } else {
            final_results
        };

        Ok(filtered)
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::query::hybrid::{ScoreBreakdown, SearchResult};
    use crate::types::MemoryNode;

    fn make_node(id: &str, content: &str, importance: f64) -> MemoryNode {
        let mut node = MemoryNode::new(content);
        node.id = id.to_string();
        node.importance = importance;
        node
    }

    fn make_search_results() -> Vec<SearchResult> {
        vec![
            SearchResult {
                node: make_node("doc-1", "Rust is a systems programming language", 0.8),
                score: 0.9,
                score_breakdown: ScoreBreakdown {
                    vector_score: Some(0.85),
                    bm25_score: Some(0.7),
                    ..Default::default()
                },
            },
            SearchResult {
                node: make_node("doc-2", "Python is great for data science", 0.6),
                score: 0.7,
                score_breakdown: ScoreBreakdown {
                    vector_score: Some(0.5),
                    bm25_score: Some(0.9),
                    ..Default::default()
                },
            },
        ]
    }

    #[test]
    fn test_to_retrieval_result() {
        let sr = &make_search_results()[0];
        let rr = HybridRetriever::to_retrieval_result(sr);

        assert_eq!(rr.id, "doc-1");
        assert_eq!(rr.score, 0.9);
        assert_eq!(rr.content, "Rust is a systems programming language");
        assert_eq!(rr.score_details.get("vector"), Some(&0.85));
        assert_eq!(rr.score_details.get("bm25"), Some(&0.7));
    }

    #[test]
    fn test_to_rerank_candidates() {
        let srs = make_search_results();
        let candidates = HybridRetriever::to_rerank_candidates(&srs);

        assert_eq!(candidates.len(), 2);
        assert_eq!(candidates[0].id, "doc-1");
        assert_eq!(candidates[0].vector_score, Some(0.85));
        assert_eq!(candidates[0].bm25_score, Some(0.7));
        assert_eq!(candidates[1].id, "doc-2");
    }

    #[test]
    fn test_create_reranker() {
        let rrf = HybridRetriever::create_reranker("rrf");
        assert_eq!(rrf.name(), "rrf");

        let linear = HybridRetriever::create_reranker("linear");
        assert_eq!(linear.name(), "linear");

        let ce = HybridRetriever::create_reranker("cross-encoder");
        assert_eq!(ce.name(), "cross-encoder");

        // Unknown defaults to RRF
        let unknown = HybridRetriever::create_reranker("unknown");
        assert_eq!(unknown.name(), "rrf");
    }

    #[test]
    fn test_build_query_with_embedding() {
        let config = RagConfig {
            top_k: 5,
            vector_weight: 0.8,
            bm25_weight: 0.2,
            enable_reranker: true,
            ..Default::default()
        };
        let emb = vec![0.1, 0.2, 0.3];
        let hq = HybridRetriever::build_query("test query", Some(&emb), &config);

        assert!(hq.vector.is_some());
        assert_eq!(hq.vector.unwrap(), vec![0.1, 0.2, 0.3]);
        assert!(hq.text.is_some());
        assert_eq!(hq.vector_weight, 0.8);
        assert_eq!(hq.text_weight, 0.2);
        assert_eq!(hq.limit, 15); // top_k * 3 because reranker is enabled
    }

    #[test]
    fn test_build_query_text_only() {
        let config = RagConfig {
            top_k: 10,
            enable_reranker: false,
            ..Default::default()
        };
        let hq = HybridRetriever::build_query("just text", None, &config);

        assert!(hq.vector.is_none());
        assert!(hq.text.is_some());
        assert_eq!(hq.limit, 10); // No over-fetch when reranker disabled
    }

    #[test]
    fn test_build_query_with_collections() {
        let config = RagConfig {
            collections: Some(vec!["notes".to_string(), "articles".to_string()]),
            ..Default::default()
        };
        let hq = HybridRetriever::build_query("query", None, &config);

        assert_eq!(
            hq.collections,
            Some(vec!["notes".to_string(), "articles".to_string()])
        );
    }

    #[test]
    fn test_rerank_candidates_roundtrip() {
        let srs = make_search_results();
        let candidates = HybridRetriever::to_rerank_candidates(&srs);

        // Rerank with RRF
        let reranker = RRFReranker::default();
        let ranked = reranker.rerank("rust programming", candidates).unwrap();

        assert_eq!(ranked.len(), 2);
        // Both should have positive scores
        for r in &ranked {
            assert!(r.score > 0.0);
        }
    }
}
