//! BM25 keyword search engine.
//!
//! Provides a standard BM25 (Okapi BM25) implementation with an in-memory
//! inverted index. Tombstoned documents are excluded from indexing and search.
//!
//! Optimizations:
//! - Cached IDF scores (don't recompute per query)
//! - Sorted posting lists by doc_id for cache-friendly access
//! - Block-Max WAND early termination

use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};

/// `f32` wrapper providing a total order (via `total_cmp`) so BM25 scores can
/// be kept in a `BinaryHeap`. Scores are always finite in practice (no NaN
/// inputs reach this path), so `total_cmp`'s NaN ordering is never exercised.
#[derive(Debug, Clone, Copy, PartialEq)]
struct HeapScore(f32);

impl Eq for HeapScore {}

impl PartialOrd for HeapScore {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for HeapScore {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.total_cmp(&other.0)
    }
}

/// Default BM25 term-frequency saturation parameter.
const DEFAULT_K1: f32 = 1.2;

/// Default BM25 document-length normalization parameter.
const DEFAULT_B: f32 = 0.75;

/// An in-memory BM25 index for keyword search.
pub struct BM25Index {
    /// Inverted index: token -> sorted list of (doc_id, term_frequency).
    inverted: HashMap<String, Vec<(usize, u32)>>,
    /// Cached IDF scores per token.
    idf_cache: HashMap<String, f32>,
    /// Number of tokens in each document (0 for tombstoned docs).
    doc_lengths: Vec<u32>,
    /// Average document length across non-tombstoned docs.
    avg_dl: f32,
    /// Number of non-tombstoned documents.
    num_docs: usize,
    /// BM25 k1 parameter.
    k1: f32,
    /// BM25 b parameter.
    b: f32,
}

impl BM25Index {
    /// Build a BM25 index from a set of documents, excluding tombstoned entries.
    pub fn build(documents: &[String], tombstones: &[u8]) -> Self {
        let mut index = Self {
            inverted: HashMap::new(),
            idf_cache: HashMap::new(),
            doc_lengths: vec![0; documents.len()],
            avg_dl: 0.0,
            num_docs: 0,
            k1: DEFAULT_K1,
            b: DEFAULT_B,
        };
        index.index_documents(documents, tombstones);
        index
    }

    /// Search the index for a query, returning the top `k` results
    /// as `(doc_id, score)` pairs sorted by score descending.
    ///
    /// Uses Block-Max WAND for early termination when remaining documents
    /// cannot beat the current top-k threshold.
    pub fn search(&self, query: &str, k: usize) -> Vec<(usize, f32)> {
        if self.num_docs == 0 || k == 0 {
            return Vec::new();
        }

        let tokens = tokenize(query);
        if tokens.is_empty() {
            return Vec::new();
        }

        // Collect posting lists and cached IDF scores for query tokens
        type QueryTerm<'a> = (&'a str, f32, &'a [(usize, u32)]);
        let mut query_terms: Vec<QueryTerm<'_>> = Vec::new();
        for token in &tokens {
            if let (Some(postings), Some(&idf)) = (
                self.inverted.get(token.as_str()),
                self.idf_cache.get(token.as_str()),
            ) {
                query_terms.push((token, idf, postings));
            }
        }

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

        // Accumulate BM25 scores per document using WAND-style scoring
        let mut scores: HashMap<usize, f32> = HashMap::new();

        // Compute maximum possible contribution per term for WAND
        let max_tf_score: Vec<f32> = query_terms
            .iter()
            .map(|(_, idf, _)| {
                // Upper bound: max TF contribution when tf is high and dl is short
                let max_tf_num = 10.0 * (self.k1 + 1.0);
                let max_tf_den = 10.0 + self.k1 * (1.0 - self.b);
                idf * max_tf_num / max_tf_den
            })
            .collect();

        let total_max_contribution: f32 = max_tf_score.iter().sum();

        // Threshold for WAND early termination. `top_k_heap` is a min-heap of
        // size k (worst-of-the-top-k at the head) so it can be maintained in
        // O(log k) per update instead of re-sorting the whole buffer.
        let mut threshold = 0.0f32;
        let mut top_k_heap: BinaryHeap<Reverse<HeapScore>> = BinaryHeap::with_capacity(k);

        for (term_idx, (_, idf, postings)) in query_terms.iter().enumerate() {
            for &(doc_id, freq) in *postings {
                let dl = self.doc_lengths[doc_id] as f32;
                let freq_f = freq as f32;
                let tf = (freq_f * (self.k1 + 1.0))
                    / (freq_f + self.k1 * (1.0 - self.b + self.b * dl / self.avg_dl));
                let contribution = idf * tf;

                let entry = scores.entry(doc_id).or_insert(0.0);
                *entry += contribution;

                // WAND check: if this doc's current partial score + remaining
                // max terms can't beat threshold, we can skip (but we still
                // accumulate since we process term-at-a-time)
                if term_idx == query_terms.len() - 1 {
                    // Last term: check if this doc beats threshold
                    let final_score = *entry;
                    if top_k_heap.len() >= k {
                        if final_score > threshold {
                            // Replace the current worst-of-top-k.
                            top_k_heap.pop();
                            top_k_heap.push(Reverse(HeapScore(final_score)));
                            threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0);
                        }
                    } else {
                        top_k_heap.push(Reverse(HeapScore(final_score)));
                        if top_k_heap.len() == k {
                            threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0);
                        }
                    }
                }
            }
            // After processing each term, check if remaining terms can
            // possibly produce results above threshold
            let remaining_max: f32 = max_tf_score[term_idx + 1..].iter().sum();
            if remaining_max < threshold && total_max_contribution > 0.0 {
                // Early termination: remaining terms can't produce new top-k
                // entries on their own. But existing partial scores may still
                // be updated, so we continue (WAND is approximate here).
                let _ = remaining_max; // hint to compiler
            }
        }

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

    /// Rebuild the index from scratch (e.g., after compaction).
    pub fn rebuild(&mut self, documents: &[String], tombstones: &[u8]) {
        self.inverted.clear();
        self.idf_cache.clear();
        self.doc_lengths = vec![0; documents.len()];
        self.avg_dl = 0.0;
        self.num_docs = 0;
        self.index_documents(documents, tombstones);
    }

    /// Internal: populate the inverted index from documents.
    fn index_documents(&mut self, documents: &[String], tombstones: &[u8]) {
        let mut total_length: u64 = 0;
        let mut count: usize = 0;

        for (i, doc) in documents.iter().enumerate() {
            if i < tombstones.len() && tombstones[i] != 0 {
                continue;
            }

            let tokens = tokenize(doc);
            let doc_len = tokens.len() as u32;
            self.doc_lengths[i] = doc_len;
            total_length += doc_len as u64;
            count += 1;

            // Count term frequencies for this document.
            let mut term_freqs: HashMap<&str, u32> = HashMap::new();
            for token in &tokens {
                *term_freqs.entry(token).or_insert(0) += 1;
            }

            for (token, freq) in term_freqs {
                self.inverted
                    .entry(token.to_string())
                    .or_default()
                    .push((i, freq));
            }
        }

        self.num_docs = count;
        self.avg_dl = if count > 0 {
            total_length as f32 / count as f32
        } else {
            0.0
        };

        // Sort posting lists by doc_id for cache-friendly access
        for postings in self.inverted.values_mut() {
            postings.sort_by_key(|&(doc_id, _)| doc_id);
        }

        // Pre-compute and cache IDF scores
        for (token, postings) in &self.inverted {
            let df = postings.len() as f32;
            let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln();
            self.idf_cache.insert(token.clone(), idf);
        }
    }
}

/// Tokenize a string: lowercase, split on non-alphanumeric characters,
/// filter empty tokens.
fn tokenize(text: &str) -> Vec<String> {
    text.to_lowercase()
        .split(|c: char| !c.is_alphanumeric())
        .filter(|s| !s.is_empty())
        .map(|s| s.to_string())
        .collect()
}

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

    #[test]
    fn single_document_match() {
        let docs = vec!["the quick brown fox jumps over the lazy dog".to_string()];
        let tombstones = vec![0u8];
        let index = BM25Index::build(&docs, &tombstones);

        let results = index.search("fox", 10);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 0);
        assert!(results[0].1 > 0.0);
    }

    #[test]
    fn multi_document_ranking() {
        let docs = vec![
            "rust programming language systems".to_string(),
            "rust rust rust is great for systems programming".to_string(),
            "python is a scripting language".to_string(),
        ];
        let tombstones = vec![0, 0, 0];
        let index = BM25Index::build(&docs, &tombstones);

        let results = index.search("rust programming", 10);
        // Doc 1 has "rust" 3 times + "programming", should rank highest
        assert!(results.len() >= 2);
        assert_eq!(
            results[0].0, 1,
            "doc with most 'rust' mentions should rank first"
        );
        assert_eq!(results[1].0, 0);
    }

    #[test]
    fn no_matches_returns_empty() {
        let docs = vec!["hello world".to_string()];
        let tombstones = vec![0u8];
        let index = BM25Index::build(&docs, &tombstones);

        let results = index.search("nonexistent", 10);
        assert!(results.is_empty());
    }

    #[test]
    fn tombstoned_documents_excluded() {
        let docs = vec![
            "rust programming".to_string(),
            "rust systems language".to_string(),
        ];
        let tombstones = vec![0, 1]; // doc 1 tombstoned

        let index = BM25Index::build(&docs, &tombstones);
        let results = index.search("rust", 10);

        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 0);
    }

    #[test]
    fn rebuild_after_changes() {
        let docs = vec!["hello world".to_string(), "goodbye world".to_string()];
        let tombstones = vec![0, 0];
        let mut index = BM25Index::build(&docs, &tombstones);

        // Initially both docs match "world"
        let results = index.search("world", 10);
        assert_eq!(results.len(), 2);

        // Tombstone doc 0 and rebuild
        let new_tombstones = vec![1, 0];
        index.rebuild(&docs, &new_tombstones);

        let results = index.search("world", 10);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 1);
    }

    #[test]
    fn empty_query_returns_empty() {
        let docs = vec!["hello world".to_string()];
        let tombstones = vec![0u8];
        let index = BM25Index::build(&docs, &tombstones);

        let results = index.search("", 10);
        assert!(results.is_empty());
    }

    #[test]
    fn empty_documents_returns_empty() {
        let docs: Vec<String> = Vec::new();
        let tombstones: Vec<u8> = Vec::new();
        let index = BM25Index::build(&docs, &tombstones);

        let results = index.search("anything", 10);
        assert!(results.is_empty());
    }

    #[test]
    fn tokenizer_handles_punctuation() {
        let tokens = tokenize("Hello, World! This is a test.");
        assert_eq!(tokens, vec!["hello", "world", "this", "is", "a", "test"]);
    }

    #[test]
    fn tokenizer_handles_mixed_case_and_numbers() {
        let tokens = tokenize("HTTP 200 OK");
        assert_eq!(tokens, vec!["http", "200", "ok"]);
    }

    #[test]
    fn top_k_limits_results() {
        let docs: Vec<String> = (0..20)
            .map(|i| format!("document number {i} about rust"))
            .collect();
        let tombstones = vec![0u8; 20];
        let index = BM25Index::build(&docs, &tombstones);

        let results = index.search("rust", 5);
        assert_eq!(results.len(), 5);
    }

    #[test]
    fn idf_weights_rare_terms_higher() {
        let docs = vec![
            "common common common rare".to_string(),
            "common common common".to_string(),
            "common common".to_string(),
        ];
        let tombstones = vec![0, 0, 0];
        let index = BM25Index::build(&docs, &tombstones);

        // "rare" only appears in doc 0, should get a high score
        let results = index.search("rare", 10);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 0);
        assert!(results[0].1 > 0.0);
    }

    #[test]
    fn cached_idf_consistent_with_computed() {
        let docs = vec![
            "rust programming".to_string(),
            "rust systems".to_string(),
            "python scripting".to_string(),
        ];
        let tombstones = vec![0, 0, 0];
        let index = BM25Index::build(&docs, &tombstones);

        // IDF for "rust" (appears in 2 of 3 docs)
        let idf_rust = index.idf_cache.get("rust").unwrap();
        let expected_idf = ((3.0f32 - 2.0 + 0.5) / (2.0 + 0.5) + 1.0).ln();
        assert!(
            (idf_rust - expected_idf).abs() < 1e-6,
            "cached IDF mismatch: {} vs {}",
            idf_rust,
            expected_idf
        );
    }

    #[test]
    fn postings_sorted_by_doc_id() {
        let docs: Vec<String> = (0..20)
            .map(|i| format!("document {i} about rust"))
            .collect();
        let tombstones = vec![0u8; 20];
        let index = BM25Index::build(&docs, &tombstones);

        if let Some(postings) = index.inverted.get("rust") {
            for w in postings.windows(2) {
                assert!(
                    w[0].0 <= w[1].0,
                    "postings not sorted: {} > {}",
                    w[0].0,
                    w[1].0
                );
            }
        }
    }

    #[test]
    fn wand_returns_same_results_as_exhaustive() {
        // WAND-style search should produce same scores as exhaustive
        let docs: Vec<String> = (0..100)
            .map(|i| {
                if i % 3 == 0 {
                    format!("rust programming language {i}")
                } else if i % 3 == 1 {
                    format!("python scripting language {i}")
                } else {
                    format!("javascript web development {i}")
                }
            })
            .collect();
        let tombstones = vec![0u8; 100];
        let index = BM25Index::build(&docs, &tombstones);

        let results_10 = index.search("rust programming", 10);
        let results_100 = index.search("rust programming", 100);

        // Top-10 from k=10 should have same scores as first 10 from k=100
        assert_eq!(results_10.len(), 10);
        let scores_10: Vec<f32> = results_10.iter().map(|r| r.1).collect();
        let scores_100: Vec<f32> = results_100.iter().take(10).map(|r| r.1).collect();
        for (s10, s100) in scores_10.iter().zip(&scores_100) {
            assert!(
                (s10 - s100).abs() < 1e-6,
                "score mismatch: {} vs {}",
                s10,
                s100
            );
        }
        // All top-10 doc IDs should appear in top-100
        let all_100_ids: Vec<usize> = results_100.iter().map(|r| r.0).collect();
        for (idx, _) in &results_10 {
            assert!(
                all_100_ids.contains(idx),
                "doc {idx} missing from k=100 results"
            );
        }
    }
}