//! Hybrid search combining lexical and semantic search.

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

use crate::config::{SearchConfig, SearchMode};
use crate::embeddings::Embedder;
use crate::error::Result;
use crate::search::{merge_results, LexicalIndex, SearchQuery, SemanticIndex, Searcher};
use crate::types::{Chunk, Document, SearchResult};

/// Hybrid searcher combining lexical and semantic search.
pub struct HybridSearcher<E: Embedder> {
    lexical: Arc<LexicalIndex>,
    semantic: Arc<SemanticIndex>,
    embedder: Arc<E>,
    config: SearchConfig,
}

impl<E: Embedder> HybridSearcher<E> {
    /// Create a new hybrid searcher.
    pub fn new(
        lexical: Arc<LexicalIndex>,
        semantic: Arc<SemanticIndex>,
        embedder: Arc<E>,
        config: SearchConfig,
    ) -> Self {
        Self {
            lexical,
            semantic,
            embedder,
            config,
        }
    }

    /// Perform hybrid search.
    pub fn search_hybrid(
        &self,
        query: &SearchQuery,
        chunks: &[Chunk],
        documents: &[Document],
    ) -> Result<Vec<SearchResult>> {
        // Build lookup maps
        let chunk_map: HashMap<String, &Chunk> = chunks
            .iter()
            .map(|c| (c.id.0.to_string(), c))
            .collect();

        let doc_map: HashMap<String, &Document> = documents
            .iter()
            .map(|d| (d.id.0.to_string(), d))
            .collect();

        match query.mode {
            SearchMode::Lexical => {
                self.lexical.search_index(query, &chunk_map, &doc_map)
            }
            SearchMode::Semantic => {
                self.semantic.search_with_embedder(
                    query,
                    self.embedder.as_ref(),
                    &chunk_map,
                    &doc_map,
                )
            }
            SearchMode::Hybrid => {
                // Run both searches in parallel conceptually
                // (could use rayon for actual parallelism)
                let lexical_results = self.lexical.search_index(query, &chunk_map, &doc_map)?;
                let semantic_results = self.semantic.search_with_embedder(
                    query,
                    self.embedder.as_ref(),
                    &chunk_map,
                    &doc_map,
                )?;

                // Merge results with configured weights
                let weights = vec![self.config.lexical_weight, self.config.semantic_weight];
                let merged = merge_results(
                    vec![lexical_results, semantic_results],
                    &weights,
                    query.limit,
                );

                Ok(merged)
            }
        }
    }
}

impl<E: Embedder> Searcher for HybridSearcher<E> {
    fn search(
        &self,
        query: &SearchQuery,
        chunks: &[Chunk],
        documents: &[Document],
    ) -> Result<Vec<SearchResult>> {
        self.search_hybrid(query, chunks, documents)
    }

    fn name(&self) -> &'static str {
        "hybrid"
    }
}

/// Builder for creating a hybrid searcher.
pub struct HybridSearcherBuilder<E: Embedder> {
    lexical: Option<Arc<LexicalIndex>>,
    semantic: Option<Arc<SemanticIndex>>,
    embedder: Option<Arc<E>>,
    config: SearchConfig,
}

impl<E: Embedder> HybridSearcherBuilder<E> {
    /// Create a new builder.
    pub fn new() -> Self {
        Self {
            lexical: None,
            semantic: None,
            embedder: None,
            config: SearchConfig::default(),
        }
    }

    /// Set the lexical index.
    pub fn lexical(mut self, index: Arc<LexicalIndex>) -> Self {
        self.lexical = Some(index);
        self
    }

    /// Set the semantic index.
    pub fn semantic(mut self, index: Arc<SemanticIndex>) -> Self {
        self.semantic = Some(index);
        self
    }

    /// Set the embedder.
    pub fn embedder(mut self, embedder: Arc<E>) -> Self {
        self.embedder = Some(embedder);
        self
    }

    /// Set the search configuration.
    pub fn config(mut self, config: SearchConfig) -> Self {
        self.config = config;
        self
    }

    /// Build the hybrid searcher.
    pub fn build(self) -> Result<HybridSearcher<E>> {
        let lexical = self
            .lexical
            .ok_or_else(|| crate::error::Error::config("Lexical index not set"))?;
        let semantic = self
            .semantic
            .ok_or_else(|| crate::error::Error::config("Semantic index not set"))?;
        let embedder = self
            .embedder
            .ok_or_else(|| crate::error::Error::config("Embedder not set"))?;

        Ok(HybridSearcher::new(lexical, semantic, embedder, self.config))
    }
}

impl<E: Embedder> Default for HybridSearcherBuilder<E> {
    fn default() -> Self {
        Self::new()
    }
}

/// Reciprocal Rank Fusion (RRF) for combining rankings.
///
/// RRF is an alternative to weighted combination that works well
/// when combining results from different search methods.
pub fn reciprocal_rank_fusion(
    rankings: Vec<Vec<String>>, // Vec of ranked chunk IDs
    k: f32,                      // Constant (typically 60)
) -> Vec<(String, f32)> {
    use std::collections::HashMap;

    let mut scores: HashMap<String, f32> = HashMap::new();

    for ranking in rankings {
        for (rank, chunk_id) in ranking.into_iter().enumerate() {
            let rrf_score = 1.0 / (k + rank as f32 + 1.0);
            *scores.entry(chunk_id).or_insert(0.0) += rrf_score;
        }
    }

    let mut result: Vec<_> = scores.into_iter().collect();
    result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    result
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_reciprocal_rank_fusion() {
        let rankings = vec![
            vec!["a".to_string(), "b".to_string(), "c".to_string()],
            vec!["b".to_string(), "a".to_string(), "d".to_string()],
        ];

        let fused = reciprocal_rank_fusion(rankings, 60.0);

        // "a" and "b" should be at top since they appear in both rankings
        assert!(fused.len() >= 2);
        let top_ids: Vec<_> = fused.iter().take(2).map(|(id, _)| id.as_str()).collect();
        assert!(top_ids.contains(&"a"));
        assert!(top_ids.contains(&"b"));
    }
}
