//! Hybrid search combining semantic vector search with BM25 keyword ranking
//!
//! Provides hybrid search capabilities for better coverage and relevance:
//! - Vector/semantic search for semantic similarity
//! - BM25 keyword ranking for exact term matching
//! - Result fusion with configurable weight tuning
//! - Normalized scoring (0.0-1.0) for fair comparison

use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use tracing::info;

use crate::error::Result;
use crate::retrieval::RetrieverEngine;

/// BM25 parameters
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct BM25Parameters {
    /// k1 parameter (term saturation, typical 1.5)
    pub k1: f32,
    /// b parameter (length normalization, typical 0.75)
    pub b: f32,
}

impl Default for BM25Parameters {
    fn default() -> Self {
        Self { k1: 1.5, b: 0.75 }
    }
}

/// Keyword index for BM25 ranking
#[derive(Debug, Clone)]
pub struct KeywordIndex {
    /// Document ID to content mapping
    documents: HashMap<String, String>,
    /// Inverted index: term -> (doc_id, frequency)
    inverted_index: HashMap<String, Vec<(String, usize)>>,
    /// Document lengths for BM25 calculation
    doc_lengths: HashMap<String, usize>,
    /// Average document length
    avg_doc_length: f32,
    /// BM25 parameters
    params: BM25Parameters,
}

impl KeywordIndex {
    /// Create a new keyword index
    pub fn new() -> Self {
        Self {
            documents: HashMap::new(),
            inverted_index: HashMap::new(),
            doc_lengths: HashMap::new(),
            avg_doc_length: 0.0,
            params: BM25Parameters::default(),
        }
    }

    /// Set BM25 parameters
    pub fn with_params(mut self, params: BM25Parameters) -> Self {
        self.params = params;
        self
    }

    /// Index a document
    pub fn index_document(&mut self, doc_id: String, content: &str) {
        self.documents.insert(doc_id.clone(), content.to_string());

        // Tokenize content
        let tokens = self.tokenize(content);
        let doc_length = tokens.len();

        // Store document length
        self.doc_lengths.insert(doc_id.clone(), doc_length);

        // Build inverted index
        let mut term_freqs: HashMap<String, usize> = HashMap::new();
        for token in tokens {
            *term_freqs.entry(token).or_insert(0) += 1;
        }

        for (term, freq) in term_freqs {
            self.inverted_index
                .entry(term)
                .or_default()
                .push((doc_id.clone(), freq));
        }

        // Update average document length
        self.update_avg_length();
    }

    /// Search for keyword matches and return BM25 scores
    pub fn search(&self, query: &str) -> Vec<(String, f32)> {
        let query_terms = self.tokenize(query);

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

        // Calculate BM25 score for each document
        let mut doc_scores: HashMap<String, f32> = HashMap::new();

        for term in &query_terms {
            if let Some(postings) = self.inverted_index.get(term) {
                let idf = self.calculate_idf(postings.len());

                for (doc_id, freq) in postings {
                    let doc_length = self.doc_lengths.get(doc_id).copied().unwrap_or(0);
                    let bm25_component =
                        self.calculate_bm25_component(idf, *freq as f32, doc_length);

                    doc_scores
                        .entry(doc_id.clone())
                        .and_modify(|score| *score += bm25_component)
                        .or_insert(bm25_component);
                }
            }
        }

        // Sort by score descending
        let mut results: Vec<_> = doc_scores
            .into_iter()
            .map(|(doc_id, score)| {
                // Normalize to 0.0-1.0 range
                let normalized = (score / (1.0 + score)).min(1.0);
                (doc_id, normalized)
            })
            .collect();

        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        results
    }

    /// Tokenize text into terms
    fn tokenize(&self, text: &str) -> Vec<String> {
        text.to_lowercase()
            .split(|c: char| !c.is_alphanumeric())
            .filter(|s| !s.is_empty() && s.len() > 1)
            .map(|s| s.to_string())
            .collect()
    }

    /// Calculate IDF (inverse document frequency)
    fn calculate_idf(&self, posting_count: usize) -> f32 {
        let total_docs = self.documents.len() as f32 + 1.0;
        ((total_docs - posting_count as f32 + 0.5) / (posting_count as f32 + 0.5) + 1.0).ln()
    }

    /// Calculate BM25 component for a term
    fn calculate_bm25_component(&self, idf: f32, term_freq: f32, doc_length: usize) -> f32 {
        let doc_length_norm = doc_length as f32 / (self.avg_doc_length + 0.0001);
        let numerator = term_freq * (self.params.k1 + 1.0);
        let denominator =
            term_freq + self.params.k1 * (1.0 - self.params.b + self.params.b * doc_length_norm);

        idf * (numerator / denominator)
    }

    /// Update average document length
    fn update_avg_length(&mut self) {
        let total_length: usize = self.doc_lengths.values().sum();
        let count = self.doc_lengths.len();
        self.avg_doc_length = if count > 0 {
            total_length as f32 / count as f32
        } else {
            0.0
        };
    }
}

impl Default for KeywordIndex {
    fn default() -> Self {
        Self::new()
    }
}

/// Hybrid search result combining vector and keyword scores
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HybridSearchResult {
    /// Document ID
    pub doc_id: String,
    /// Document content preview
    pub preview: String,
    /// Vector search score (0.0-1.0)
    pub vector_score: f32,
    /// Keyword search score (0.0-1.0)
    pub keyword_score: f32,
    /// Fused score (weighted combination)
    pub fused_score: f32,
    /// Search method that contributed most
    pub primary_method: String,
}

/// Hybrid search engine combining vector and keyword search
pub struct HybridSearchEngine {
    retriever: RetrieverEngine,
    keyword_index: KeywordIndex,
    vector_weight: f32,
    keyword_weight: f32,
}

impl HybridSearchEngine {
    /// Create a new hybrid search engine
    pub fn new(retriever: RetrieverEngine, keyword_index: KeywordIndex) -> Self {
        Self {
            retriever,
            keyword_index,
            vector_weight: 0.5, // Default 50/50 split
            keyword_weight: 0.5,
        }
    }

    /// Set search weights (must sum to <= 1.0)
    pub fn with_weights(mut self, vector_weight: f32, keyword_weight: f32) -> Self {
        self.vector_weight = vector_weight;
        self.keyword_weight = keyword_weight;
        self
    }

    /// Perform hybrid search
    pub async fn search(
        &self,
        query: &str,
        limit: Option<usize>,
    ) -> Result<Vec<HybridSearchResult>> {
        info!("Performing hybrid search for query: {}", query);

        let limit = limit.unwrap_or(10);

        // 1. Vector search
        let vector_results = self.retriever.search(query, Some(limit)).await?;
        let vector_scores: HashMap<String, f32> = vector_results
            .iter()
            .map(|r| (r.doc_id.clone(), r.similarity))
            .collect();

        info!("Vector search returned {} results", vector_results.len());

        // 2. Keyword search
        let keyword_results = self.keyword_index.search(query);
        let keyword_map: HashMap<String, f32> = keyword_results.iter().cloned().collect();

        info!("Keyword search returned {} results", keyword_results.len());

        // 3. Fuse results
        let mut fused_results = Vec::new();
        let mut seen_docs = std::collections::HashSet::new();

        // Add vector results
        for vector_result in &vector_results {
            let vector_score = *vector_scores.get(&vector_result.doc_id).unwrap_or(&0.0);
            let keyword_score = *keyword_map.get(&vector_result.doc_id).unwrap_or(&0.0);

            // Weighted fusion
            let fused_score =
                (vector_score * self.vector_weight) + (keyword_score * self.keyword_weight);

            let primary_method = if vector_score >= keyword_score {
                "vector".to_string()
            } else {
                "keyword".to_string()
            };

            fused_results.push(HybridSearchResult {
                doc_id: vector_result.doc_id.clone(),
                preview: vector_result.content.clone(),
                vector_score,
                keyword_score,
                fused_score,
                primary_method,
            });

            seen_docs.insert(vector_result.doc_id.clone());
        }

        // Add keyword-only results
        for (doc_id, keyword_score) in keyword_map {
            if !seen_docs.contains(&doc_id) && fused_results.len() < limit {
                let vector_score = 0.0;
                let fused_score = keyword_score * self.keyword_weight;

                fused_results.push(HybridSearchResult {
                    doc_id: doc_id.clone(),
                    preview: format!("Document {}", doc_id),
                    vector_score,
                    keyword_score,
                    fused_score,
                    primary_method: "keyword".to_string(),
                });

                seen_docs.insert(doc_id);
            }
        }

        // 4. Sort by fused score
        fused_results.sort_by(|a, b| {
            b.fused_score
                .partial_cmp(&a.fused_score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Truncate to limit
        fused_results.truncate(limit);

        info!(
            "Hybrid search returned {} fused results",
            fused_results.len()
        );

        Ok(fused_results)
    }
}

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

    #[test]
    fn test_keyword_index_creation() {
        let index = KeywordIndex::new();
        assert_eq!(index.documents.len(), 0);
        assert_eq!(index.inverted_index.len(), 0);
    }

    #[test]
    fn test_keyword_index_document_indexing() {
        let mut index = KeywordIndex::new();
        index.index_document(
            "doc1".to_string(),
            "kubernetes is a container orchestration platform",
        );

        assert_eq!(index.documents.len(), 1);
        assert!(!index.inverted_index.is_empty());
    }

    #[test]
    fn test_keyword_index_tokenization() {
        let index = KeywordIndex::new();
        let tokens = index.tokenize("Hello, world! This is a test.");

        assert!(tokens.contains(&"hello".to_string()));
        assert!(tokens.contains(&"world".to_string()));
        assert!(!tokens.contains(&"a".to_string())); // Single char filtered
    }

    #[test]
    fn test_bm25_scoring() {
        let mut index = KeywordIndex::new();
        index.index_document(
            "doc1".to_string(),
            "kubernetes container orchestration platform",
        );
        index.index_document("doc2".to_string(), "docker containerization");
        index.index_document(
            "doc3".to_string(),
            "kubernetes and docker integration guide",
        );

        let results = index.search("kubernetes");

        assert!(!results.is_empty());
        // doc1 should have highest score (most relevant)
        assert_eq!(results[0].0, "doc1");
        // BM25 scores are normalized 0.0-1.0 but typically lower
        assert!(results[0].1 > 0.0);
        // doc1 should rank higher than doc2 (which doesn't contain kubernetes)
        assert!(results[0].1 > 0.0001);
    }

    #[test]
    fn test_bm25_empty_query() {
        let index = KeywordIndex::new();
        let results = index.search("");

        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_bm25_no_matches() {
        let mut index = KeywordIndex::new();
        index.index_document("doc1".to_string(), "kubernetes container orchestration");

        let results = index.search("xyz_nonexistent_term");

        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_hybrid_search_result_creation() {
        let result = HybridSearchResult {
            doc_id: "doc1".to_string(),
            preview: "Sample content".to_string(),
            vector_score: 0.8,
            keyword_score: 0.6,
            fused_score: 0.7,
            primary_method: "vector".to_string(),
        };

        assert_eq!(result.doc_id, "doc1");
        assert!(result.fused_score > 0.0 && result.fused_score <= 1.0);
    }

    #[test]
    fn test_bm25_parameters() {
        let params = BM25Parameters::default();
        assert_eq!(params.k1, 1.5);
        assert_eq!(params.b, 0.75);
    }

    #[test]
    fn test_bm25_custom_parameters() {
        let custom_params = BM25Parameters { k1: 2.0, b: 0.5 };
        let index = KeywordIndex::new().with_params(custom_params);

        assert_eq!(index.params.k1, 2.0);
        assert_eq!(index.params.b, 0.5);
    }

    #[test]
    fn test_idf_calculation() {
        let mut index = KeywordIndex::new();
        index.index_document("doc1".to_string(), "kubernetes");
        index.index_document("doc2".to_string(), "kubernetes");
        index.index_document("doc3".to_string(), "kubernetes docker");

        // More documents = lower IDF
        let idf_popular = index.calculate_idf(3);
        let idf_rare = index.calculate_idf(1);

        assert!(idf_rare > idf_popular);
    }

    #[test]
    fn test_average_document_length() {
        let mut index = KeywordIndex::new();
        index.index_document("doc1".to_string(), "short");
        index.index_document("doc2".to_string(), "this is a much longer document");

        assert!(index.avg_doc_length > 0.0);
    }

    #[test]
    fn test_multiple_document_indexing() {
        let mut index = KeywordIndex::new();
        for i in 0..5 {
            index.index_document(format!("doc{}", i), "some sample content for testing");
        }

        assert_eq!(index.documents.len(), 5);
    }

    #[test]
    fn test_keyword_search_ranking() {
        let mut index = KeywordIndex::new();
        index.index_document("doc1".to_string(), "kubernetes kubernetes kubernetes");
        index.index_document("doc2".to_string(), "kubernetes");
        index.index_document("doc3".to_string(), "other content");

        let results = index.search("kubernetes");

        assert_eq!(results.len(), 2);
        // doc1 should rank higher (more occurrences)
        assert!(results[0].1 > results[1].1);
    }

    #[test]
    fn test_score_normalization() {
        let mut index = KeywordIndex::new();
        index.index_document(
            "doc1".to_string(),
            "kubernetes orchestration platform containers",
        );

        let results = index.search("kubernetes");

        assert!(!results.is_empty());
        // All scores should be normalized 0.0-1.0
        assert!(results[0].1 >= 0.0 && results[0].1 <= 1.0);
    }

    #[test]
    fn test_bm25_parameters_with_engine() {
        let params = BM25Parameters { k1: 2.0, b: 0.5 };
        let index = KeywordIndex::new().with_params(params);

        assert_eq!(index.params.k1, 2.0);
        assert_eq!(index.params.b, 0.5);
    }

    #[test]
    fn test_keyword_index_get_average_length() {
        let mut index = KeywordIndex::new();
        index.index_document(
            "doc1".to_string(),
            "one two three four five six seven eight nine ten",
        );
        index.index_document("doc2".to_string(), "short");

        // Average should be between 1 and 10
        assert!(index.avg_doc_length > 1.0 && index.avg_doc_length < 10.0);
    }

    #[test]
    fn test_hybrid_search_weights() {
        let _keyword_index = KeywordIndex::new();
        // Can't fully test without RetrieverEngine, but test structure
        let weights = (0.6, 0.4);
        assert_eq!(weights.0 + weights.1, 1.0);
    }

    #[test]
    fn test_bm25_component_calculation() {
        let index = KeywordIndex::new();
        let idf = 2.0;
        let term_freq = 3.0;
        let doc_length = 100;

        let component = index.calculate_bm25_component(idf, term_freq, doc_length);
        assert!(component > 0.0);
    }

    #[test]
    fn test_multiple_term_search() {
        let mut index = KeywordIndex::new();
        index.index_document(
            "doc1".to_string(),
            "kubernetes and docker integration is important",
        );
        index.index_document("doc2".to_string(), "kubernetes basics");

        let results = index.search("kubernetes docker integration");

        assert!(!results.is_empty());
        // doc1 should rank higher (has all terms)
        assert_eq!(results[0].0, "doc1");
    }
}
