//! 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.
//!
//! The index is **incremental**: [`BM25Index::add_document`] and
//! [`BM25Index::remove_document`] keep it exactly equivalent to one built from
//! scratch over the same live documents, so a store can maintain one index for
//! its lifetime instead of re-tokenising the whole corpus per query. To make
//! that possible IDF is computed at query time (it depends on the live
//! document count) rather than cached at build time.
//!
//! - Posting lists sorted by doc id
//! - Bounded-heap top-k; results ordered by score, then doc id (deterministic)

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)>>,
    /// Number of tokens in each document (0 for tombstoned docs).
    doc_lengths: Vec<u32>,
    /// Sum of `doc_lengths` over live documents (keeps `avg_dl` exact under
    /// incremental updates).
    total_length: u64,
    /// 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(),
            doc_lengths: vec![0; documents.len()],
            total_length: 0,
            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 k == 0 {
            return Vec::new();
        }
        // Top-k with a bounded min-heap: O(matches * log k) instead of sorting
        // every match. Ties break towards the lower doc id so results are
        // deterministic.
        let mut heap: BinaryHeap<Reverse<(HeapScore, Reverse<usize>)>> =
            BinaryHeap::with_capacity(k.min(1024) + 1);
        for (doc_id, score) in self.scores(query) {
            heap.push(Reverse((HeapScore(score), Reverse(doc_id))));
            if heap.len() > k {
                heap.pop();
            }
        }
        let mut results: Vec<(usize, f32)> = heap
            .into_iter()
            .map(|Reverse((HeapScore(score), Reverse(doc_id)))| (doc_id, score))
            .collect();
        results.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
        results
    }

    /// The BM25 score of **every** matching document, in doc-id order, unsorted
    /// by score. Score fusion normalises over the whole matching set, so it
    /// needs all of these but not their ranking; producing a ranked list of
    /// every match (`search(query, corpus_len)`) spent most of its time sorting.
    pub fn scores(&self, query: &str) -> Vec<(usize, f32)> {
        if self.num_docs == 0 {
            return Vec::new();
        }
        // Term-at-a-time accumulation into a dense array: a common term has a
        // posting per document, and hashing each one dominated query time.
        // IDF is computed here rather than cached at build time: it depends on
        // the live document count, which changes with every incremental
        // add/remove, and costs one `ln` per query term.
        let mut acc = vec![0.0f32; self.doc_lengths.len()];
        let mut matched = false;
        for token in tokenize(query) {
            let Some(postings) = self.inverted.get(token.as_str()) else {
                continue;
            };
            matched = true;
            let df = postings.len() as f32;
            let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln();
            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));
                acc[doc_id] += idf * tf;
            }
        }
        if !matched {
            return Vec::new();
        }
        // Every contribution is strictly positive (idf = ln(1 + x), x > 0), so
        // a zero entry is a document no query term touched.
        acc.into_iter()
            .enumerate()
            .filter(|&(_, score)| score > 0.0)
            .collect()
    }

    /// Number of document slots (live or not) the index covers. Ids are
    /// positions in the document list it mirrors.
    pub fn len(&self) -> usize {
        self.doc_lengths.len()
    }

    /// `true` when the index covers no document slots.
    pub fn is_empty(&self) -> bool {
        self.doc_lengths.is_empty()
    }

    /// Index `text` as document `doc_id`, which must be the next free id
    /// (`self.len()`) or an existing slot that is currently empty (removed or
    /// tombstoned). After any sequence of `add_document` / `remove_document`
    /// calls the index scores exactly as one freshly built from the same live
    /// documents.
    pub fn add_document(&mut self, doc_id: usize, text: &str) {
        if doc_id >= self.doc_lengths.len() {
            self.doc_lengths.resize(doc_id + 1, 0);
        }
        debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied");

        let tokens = tokenize(text);
        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 {
            let postings = self.inverted.entry(token.to_string()).or_default();
            // Posting lists stay sorted by doc id; appends are the common case.
            match postings.last() {
                Some(&(last, _)) if last >= doc_id => {
                    let at = postings.partition_point(|&(id, _)| id < doc_id);
                    postings.insert(at, (doc_id, freq));
                }
                _ => postings.push((doc_id, freq)),
            }
        }
        self.doc_lengths[doc_id] = tokens.len() as u32;
        self.total_length += tokens.len() as u64;
        self.num_docs += 1;
        self.refresh_avg_dl();
    }

    /// Extend the index to cover `len` document slots, leaving new ones empty.
    /// Used for slots that hold no live document (tombstoned records).
    pub fn pad_to(&mut self, len: usize) {
        if len > self.doc_lengths.len() {
            self.doc_lengths.resize(len, 0);
        }
    }

    /// Remove document `doc_id`, whose indexed text was `text`. The text is
    /// needed to find its postings; pass exactly what was added.
    pub fn remove_document(&mut self, doc_id: usize, text: &str) {
        let tokens = tokenize(text);
        let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
        for token in &tokens {
            if !seen.insert(token) {
                continue;
            }
            if let Some(postings) = self.inverted.get_mut(token.as_str()) {
                if let Ok(at) = postings.binary_search_by_key(&doc_id, |&(id, _)| id) {
                    postings.remove(at);
                }
                if postings.is_empty() {
                    self.inverted.remove(token.as_str());
                }
            }
        }
        if let Some(len) = self.doc_lengths.get_mut(doc_id) {
            self.total_length = self.total_length.saturating_sub(u64::from(*len));
            *len = 0;
        }
        self.num_docs = self.num_docs.saturating_sub(1);
        self.refresh_avg_dl();
    }

    fn refresh_avg_dl(&mut self) {
        self.avg_dl = if self.num_docs > 0 {
            self.total_length as f32 / self.num_docs as f32
        } else {
            0.0
        };
    }

    /// Rebuild the index from scratch (e.g., after compaction).
    pub fn rebuild(&mut self, documents: &[String], tombstones: &[u8]) {
        self.inverted.clear();
        self.doc_lengths = vec![0; documents.len()];
        self.total_length = 0;
        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.total_length = total_length;
        self.refresh_avg_dl();

        // 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);
        }
    }
}

/// 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 score_matches_the_bm25_formula() {
        let docs = vec![
            "rust programming".to_string(),
            "rust systems".to_string(),
            "python scripting".to_string(),
        ];
        let index = BM25Index::build(&docs, &[0, 0, 0]);

        // "python": df = 1 of N = 3. Every doc has the average length (2) and
        // tf = 1, so the tf factor is exactly 1 and the score is the IDF.
        let results = index.search("python", 3);
        let expected_idf = ((3.0f32 - 1.0 + 0.5) / (1.0 + 0.5) + 1.0).ln();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 2);
        assert!((results[0].1 - expected_idf).abs() < 1e-6, "{results:?}");
    }

    #[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"
            );
        }
    }
    /// Documents drawn from a small vocabulary so terms collide heavily.
    fn random_doc(state: &mut u64) -> String {
        const VOCAB: &[&str] = &[
            "alpha", "beta", "gamma", "delta", "eps", "zeta", "eta", "x1",
        ];
        let mut next = || {
            *state = state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            (*state >> 33) as usize
        };
        let len = 1 + next() % 9;
        (0..len)
            .map(|_| VOCAB[next() % VOCAB.len()])
            .collect::<Vec<_>>()
            .join(" ")
    }

    #[test]
    fn incremental_updates_match_a_fresh_build_exactly() {
        for seed in 0..60u64 {
            let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
            let mut docs: Vec<String> = Vec::new();
            let mut tombstones: Vec<u8> = Vec::new();
            let mut index = BM25Index::build(&docs, &tombstones);

            for step in 0..80 {
                state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
                let live: Vec<usize> = (0..docs.len()).filter(|&i| tombstones[i] == 0).collect();
                match (state >> 40) % 4 {
                    0 if !live.is_empty() => {
                        // delete
                        let id = live[(state >> 20) as usize % live.len()];
                        index.remove_document(id, &docs[id]);
                        tombstones[id] = 1;
                    }
                    1 if !live.is_empty() => {
                        // update in place
                        let id = live[(state >> 20) as usize % live.len()];
                        let new_text = random_doc(&mut state);
                        index.remove_document(id, &docs[id]);
                        index.add_document(id, &new_text);
                        docs[id] = new_text;
                    }
                    _ => {
                        let text = random_doc(&mut state);
                        index.add_document(docs.len(), &text);
                        docs.push(text);
                        tombstones.push(0);
                    }
                }

                let fresh = BM25Index::build(&docs, &tombstones);
                for query in ["alpha", "beta gamma", "x1 zeta alpha delta", "missing"] {
                    let got = index.search(query, 5);
                    let want = fresh.search(query, 5);
                    assert_eq!(got.len(), want.len(), "seed {seed} step {step} {query:?}");
                    for (g, w) in got.iter().zip(&want) {
                        assert_eq!(
                            g.0, w.0,
                            "seed {seed} step {step} {query:?}: {got:?} vs {want:?}"
                        );
                        assert!(
                            (g.1 - w.1).abs() < 1e-5,
                            "seed {seed} step {step} {query:?}"
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn scores_is_the_unranked_form_of_a_full_search() {
        let mut state = 99u64;
        let docs: Vec<String> = (0..200).map(|_| random_doc(&mut state)).collect();
        let tombstones: Vec<u8> = (0..200).map(|i| u8::from(i % 7 == 0)).collect();
        let index = BM25Index::build(&docs, &tombstones);
        for query in ["alpha", "beta gamma x1", "missing", ""] {
            let mut all = index.scores(query);
            all.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
            assert_eq!(all, index.search(query, docs.len()), "{query:?}");
            assert!(all.iter().all(|(id, _)| tombstones[*id] == 0));
        }
    }

    #[test]
    fn ties_break_towards_the_lower_doc_id() {
        let docs: Vec<String> = (0..6).map(|_| "same text".to_string()).collect();
        let index = BM25Index::build(&docs, &[0; 6]);
        let ids: Vec<usize> = index.search("same", 3).into_iter().map(|r| r.0).collect();
        assert_eq!(ids, [0, 1, 2]);
    }
}