anda_db_tfs 0.8.2

//! # BM25 index implementation
//!
//! This module contains [`BM25Index`], the concurrent, bucket-sharded BM25
//! index that backs the crate. See the crate-level documentation for a
//! high-level overview.

use anda_db_utils::UniqueVec;
use dashmap::DashMap;
use parking_lot::RwLock;
use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
use serde::{Deserialize, Serialize};
use std::{
    io::{Read, Write},
    sync::atomic::{AtomicU32, AtomicU64, Ordering},
};

use crate::error::*;
use crate::query::*;
use crate::tokenizer::*;

fn cbor_serialized_size<T: ?Sized + Serialize>(value: &T) -> usize {
    cbor2::serialized_size(value)
        .expect("CBOR serialized size calculation failed")
        .try_into()
        .expect("CBOR serialized size exceeds usize")
}

/// Concurrent, bucket-sharded full-text index using BM25 scoring.
///
/// The index keeps its in-memory state in a handful of `DashMap`s so that
/// inserts, deletes and searches can run concurrently from many threads.
/// Persistence is split into two parts:
///
/// * **Metadata** — name, configuration, statistics and the current bucket/
///   document id watermarks. Serialized in CBOR by [`store_metadata`].
/// * **Buckets** — the actual postings and per-document token counts. Each
///   token is assigned to exactly one *bucket*, a self-contained CBOR blob
///   whose serialized size is bounded by [`BM25Config::bucket_overload_size`].
///   Only buckets whose `dirty_version` has advanced past their
///   `saved_version` are re-written on [`flush`], which makes repeated flushes
///   cheap even for large indices.
///
/// [`flush`]: Self::flush
/// [`store_metadata`]: Self::store_metadata
pub struct BM25Index<T: Tokenizer + Clone> {
    /// Index name
    name: String,

    /// Tokenizer used to process text
    tokenizer: T,

    /// BM25 algorithm parameters
    config: BM25Config,

    /// Maps document IDs to their token counts
    doc_tokens: DashMap<u64, usize>,

    /// Buckets store information about where posting entries are stored and their current state
    buckets: DashMap<u32, Bucket>,

    /// Inverted index mapping tokens to (bucket id, Vec<(document_id, term_frequency)>)
    postings: DashMap<String, PostingValue>,

    /// Index metadata.
    metadata: RwLock<BM25Metadata>,

    /// Maximum bucket ID currently in use
    max_bucket_id: AtomicU32,

    /// Maximum document ID currently in use
    max_document_id: AtomicU64,

    /// Average number of tokens per document
    avg_doc_tokens: RwLock<f32>,

    /// Total number of tokens indexed.
    total_tokens: AtomicU64,

    /// Number of search operations performed.
    search_count: AtomicU64,

    /// Last saved version of the index
    last_saved_version: AtomicU64,
}

#[derive(Default)]
struct Bucket {
    /// Version counter incremented on each modification
    dirty_version: u64,
    /// Version that was last successfully persisted
    saved_version: u64,
    // Current size of the bucket in bytes
    size: usize,
    // List of tokens stored in this bucket
    tokens: UniqueVec<String>,
    // Set of document IDs associated with this bucket
    doc_ids: FxHashSet<u64>,
}

impl Bucket {
    #[inline]
    fn is_dirty(&self) -> bool {
        self.dirty_version > self.saved_version
    }

    #[inline]
    fn mark_dirty(&mut self) {
        self.dirty_version += 1;
    }
}

/// Parameters controlling the BM25 scoring formula.
///
/// BM25 ranks a document `d` against a multi-term query `q` as:
///
/// ```text
/// score(d, q) = Σ_{t ∈ q} idf(t) · (tf · (k1 + 1))
///                                 / (tf + k1 · (1 − b + b · |d| / avgdl))
/// ```
///
/// - `k1` controls **term frequency saturation**. Larger values give more
///   weight to repeated occurrences of a term. Typical values: `1.2..=2.0`.
/// - `b` controls **document length normalization**. `0.0` disables length
///   normalization; `1.0` applies full normalization. Typical value: `0.75`.
///
/// Values outside their natural ranges are clamped at scoring time
/// (`k1` to `>= 0`, `b` to `[0, 1]`) to avoid producing `NaN`/`inf` scores.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BM25Params {
    /// Term-frequency saturation factor.
    ///
    /// Higher values make repeated occurrences of a term contribute more to
    /// the score. Typical values are in the `1.2..=2.0` range.
    pub k1: f32,
    /// Document-length normalization factor.
    ///
    /// `0.0` disables length normalization; `1.0` applies full normalization.
    /// The usual BM25 default is `0.75`.
    pub b: f32,
}

impl Default for BM25Params {
    /// Returns default BM25 parameters (`k1 = 1.2`, `b = 0.75`) which work well
    /// for most use cases.
    fn default() -> Self {
        BM25Params { k1: 1.2, b: 0.75 }
    }
}

/// Top-level configuration of a [`BM25Index`].
///
/// * `bm25` — the scoring parameters, see [`BM25Params`].
/// * `bucket_overload_size` — the soft upper bound, in bytes of the serialized
///   CBOR payload, of a single bucket. When inserting a new token would push a
///   bucket past this limit the token is routed to a fresh bucket instead.
///   Smaller values produce more, smaller buckets (cheaper incremental flushes
///   but more I/O per full reload); larger values do the opposite.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BM25Config {
    /// BM25 scoring parameters used for all query scoring.
    pub bm25: BM25Params,
    /// Maximum size of a bucket before creating a new one
    /// When a bucket's stored data exceeds this size,
    /// a new bucket should be created for new data
    pub bucket_overload_size: usize,
}

impl Default for BM25Config {
    /// Returns a default configuration with [`BM25Params::default`] and a
    /// 512 KiB bucket size limit.
    fn default() -> Self {
        BM25Config {
            bm25: BM25Params::default(),
            bucket_overload_size: 1024 * 512,
        }
    }
}

/// Type alias for posting values: (bucket id, Vec<(document_id, token_frequency)>)
/// - bucket_id: The bucket where this posting is stored
/// - Vec<(document_id, token_frequency)>: List of documents and their term frequencies
pub type PostingValue = (u32, UniqueVec<(u64, usize)>);

/// Index metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BM25Metadata {
    /// Index name.
    pub name: String,

    /// BM25 algorithm parameters
    pub config: BM25Config,

    /// Index statistics.
    pub stats: BM25Stats,
}

/// Index statistics.
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct BM25Stats {
    /// Last insertion timestamp (unix ms).
    pub last_inserted: u64,

    /// Last deletion timestamp (unix ms).
    pub last_deleted: u64,

    /// Last saved timestamp (unix ms).
    pub last_saved: u64,

    /// Updated version for the index. It will be incremented when the index is updated.
    pub version: u64,

    /// Number of elements in the index.
    pub num_elements: u64,

    /// Number of search operations performed.
    pub search_count: u64,

    /// Number of insert operations performed.
    pub insert_count: u64,

    /// Number of delete operations performed.
    pub delete_count: u64,

    /// Maximum bucket ID currently in use
    pub max_bucket_id: u32,

    /// Maximum document ID currently in use
    pub max_document_id: u64,

    /// Average number of tokens per document
    pub avg_doc_tokens: f32,
}

/// Serializable BM25 index structure (owned version).
#[derive(Clone, Serialize, Deserialize)]
struct BM25IndexOwned {
    // postings: DashMap<String, PostingValue>,
    metadata: BM25Metadata,
}

#[derive(Clone, Serialize)]
struct BM25IndexRef<'a> {
    // postings: &'a DashMap<String, PostingValue>,
    metadata: &'a BM25Metadata,
}

// Helper structure for serialization and deserialization of bucket
#[derive(Debug, Clone, Serialize, Deserialize)]
struct BucketOwned {
    #[serde(rename = "p")]
    postings: FxHashMap<String, PostingValue>,

    #[serde(rename = "d")]
    doc_tokens: FxHashMap<u64, usize>,
}

// Reference structure for serializing bucket
#[derive(Serialize)]
struct BucketRef<'a> {
    #[serde(rename = "p")]
    postings: &'a FxHashMap<&'a String, dashmap::mapref::one::Ref<'a, String, PostingValue>>,

    #[serde(rename = "d")]
    doc_tokens: &'a FxHashMap<u64, usize>,
}

impl<T> BM25Index<T>
where
    T: Tokenizer + Clone,
{
    /// Creates a new empty BM25 index with the given tokenizer and optional config.
    ///
    /// # Arguments
    ///
    /// * `name` - Name of the index
    /// * `tokenizer` - Tokenizer to use for processing text
    /// * `config` - Optional BM25 configuration parameters
    ///
    /// # Returns
    ///
    /// * `BM25Index` - A new instance of the BM25 index
    pub fn new(name: String, tokenizer: T, config: Option<BM25Config>) -> Self {
        let config = config.unwrap_or_default();
        let stats = BM25Stats {
            version: 1,
            ..Default::default()
        };
        BM25Index {
            name: name.clone(),
            tokenizer,
            config: config.clone(),
            doc_tokens: DashMap::new(),
            postings: DashMap::new(),
            buckets: DashMap::from_iter([(0, Bucket::default())]),
            metadata: RwLock::new(BM25Metadata {
                name,
                config,
                stats,
            }),
            max_bucket_id: AtomicU32::new(0),
            max_document_id: AtomicU64::new(0),
            avg_doc_tokens: RwLock::new(0.0),
            total_tokens: AtomicU64::new(0),
            search_count: AtomicU64::new(0),
            last_saved_version: AtomicU64::new(0),
        }
    }

    /// Loads a complete index (metadata and all buckets) in one call.
    ///
    /// This is a convenience wrapper around [`load_metadata`](Self::load_metadata)
    /// followed by [`load_buckets`](Self::load_buckets).
    ///
    /// # Arguments
    ///
    /// * `tokenizer` — tokenizer to attach to the loaded index. It does not
    ///   need to be identical to the one originally used, but queries will
    ///   only be meaningful if the tokenization is compatible.
    /// * `metadata` — reader positioned at the start of the CBOR metadata blob.
    /// * `f` — async function invoked with each bucket id in `0..=max_bucket_id`;
    ///   return `Ok(Some(bytes))` for present buckets or `Ok(None)` to skip.
    ///
    /// # Returns
    ///
    /// The fully-loaded index, or a [`BM25Error`] if metadata could not be
    /// parsed or a bucket failed to load.
    pub async fn load_all<R: Read, F>(tokenizer: T, metadata: R, f: F) -> Result<Self, BM25Error>
    where
        F: AsyncFnMut(u32) -> Result<Option<Vec<u8>>, BoxError>,
    {
        let mut index = Self::load_metadata(tokenizer, metadata)?;
        index.load_buckets(f).await?;
        Ok(index)
    }

    /// Loads only the index metadata, returning an empty shell.
    ///
    /// The returned index contains the correct configuration, statistics and
    /// id watermarks, but no postings or `doc_tokens`. Call
    /// [`load_buckets`](Self::load_buckets) afterwards to populate the inverted
    /// index (possibly on demand, or only for a subset of buckets).
    pub fn load_metadata<R: Read>(tokenizer: T, r: R) -> Result<Self, BM25Error> {
        let index: BM25IndexOwned =
            cbor2::from_reader(r).map_err(|err| BM25Error::Serialization {
                name: "unknown".to_string(),
                source: err.into(),
            })?;
        let max_bucket_id = AtomicU32::new(index.metadata.stats.max_bucket_id);
        let max_document_id = AtomicU64::new(index.metadata.stats.max_document_id);
        let search_count = AtomicU64::new(index.metadata.stats.search_count);
        let avg_doc_tokens = RwLock::new(index.metadata.stats.avg_doc_tokens);
        let last_saved_version = AtomicU64::new(index.metadata.stats.version);

        Ok(BM25Index {
            name: index.metadata.name.clone(),
            tokenizer,
            config: index.metadata.config.clone(),
            doc_tokens: DashMap::new(),
            postings: DashMap::new(),
            buckets: DashMap::from_iter([(0, Bucket::default())]),
            metadata: RwLock::new(index.metadata),
            max_bucket_id,
            max_document_id,
            avg_doc_tokens,
            search_count,
            last_saved_version,
            total_tokens: AtomicU64::new(0),
        })
    }

    /// Populates the inverted index from previously persisted buckets.
    ///
    /// Intended to be called right after [`load_metadata`](Self::load_metadata).
    /// `f` is invoked once per bucket id in `0..=max_bucket_id`; returning
    /// `Ok(None)` leaves that bucket empty, which is useful for partial loads
    /// (e.g. lazy-loading buckets on first access).
    ///
    /// After this call, `total_tokens` and `avg_doc_tokens` are recomputed
    /// from the documents that were actually loaded.
    ///
    /// Posting entries that reference a document with no token count in any
    /// loaded bucket are pruned and the affected buckets are marked dirty, so
    /// the next [`flush`](Self::flush) persists the cleanup. Buckets are
    /// written self-contained (a bucket's `doc_tokens` cover every document
    /// referenced by its postings), so such entries can only be leftovers of
    /// documents removed with non-original text; documents from buckets that
    /// were skipped via `Ok(None)` are unaffected.
    pub async fn load_buckets<F>(&mut self, mut f: F) -> Result<(), BM25Error>
    where
        F: AsyncFnMut(u32) -> Result<Option<Vec<u8>>, BoxError>,
    {
        let mut doc_token_lengths: FxHashMap<u64, usize> = self
            .doc_tokens
            .iter()
            .map(|entry| (*entry.key(), *entry.value()))
            .collect();

        for i in 0..=self.max_bucket_id.load(Ordering::Relaxed) {
            let data = f(i).await.map_err(|err| BM25Error::Generic {
                name: self.name.clone(),
                source: err,
            })?;
            if let Some(data) = data {
                let bucket: BucketOwned =
                    cbor2::from_reader(&data[..]).map_err(|err| BM25Error::Serialization {
                        name: self.name.clone(),
                        source: err.into(),
                    })?;

                let mut b = Bucket {
                    size: data.len(),
                    ..Default::default()
                };
                if !bucket.doc_tokens.is_empty() {
                    b.doc_ids = bucket.doc_tokens.keys().cloned().collect();
                    for (doc_id, token_count) in bucket.doc_tokens {
                        doc_token_lengths.insert(doc_id, token_count);
                    }
                }

                if !bucket.postings.is_empty() {
                    for (token, mut posting) in bucket.postings {
                        // The bucket file path is the source of truth for ownership.
                        // If a stale lower-numbered bucket is still present after a
                        // partial flush, later buckets win and the old bucket is
                        // marked dirty so the stale token is removed on the next flush.
                        posting.0 = i;
                        if let Some(previous) = self.postings.insert(token.clone(), posting) {
                            let previous_bucket_id = previous.0;
                            if previous_bucket_id != i
                                && let Some(mut previous_bucket) =
                                    self.buckets.get_mut(&previous_bucket_id)
                                && previous_bucket
                                    .tokens
                                    .swap_remove_if(|k| &token == k)
                                    .is_some()
                            {
                                let previous_size = cbor_serialized_size(&(&token, &previous)) + 2;
                                previous_bucket.size =
                                    previous_bucket.size.saturating_sub(previous_size);
                                previous_bucket.mark_dirty();
                            }
                        }

                        b.tokens.push(token);
                    }
                }

                self.buckets.insert(i, b);
            }
        }

        let mut doc_ids_by_bucket: FxHashMap<u32, FxHashSet<u64>> = FxHashMap::default();
        let mut loaded_doc_tokens: FxHashMap<u64, usize> = FxHashMap::default();
        let mut empty_tokens: Vec<(u32, String)> = Vec::new();
        let mut bucket_size_decrease: FxHashMap<u32, usize> = FxHashMap::default();

        for mut posting in self.postings.iter_mut() {
            let bucket_id = posting.0;
            let doc_ids = doc_ids_by_bucket.entry(bucket_id).or_default();
            // Prune entries whose document has no token length anywhere.
            // Buckets are self-contained (a bucket's doc_tokens cover every
            // document referenced by its postings), so after loading, an entry
            // without a token length can only be a stale leftover from a
            // remove() that was given non-original text. Dropping it here makes
            // the index self-healing on reload. Documents from buckets that
            // were intentionally skipped (partial load) are not affected.
            let mut removed_entries: Vec<(u64, usize)> = Vec::new();
            posting.1.retain(|entry| {
                if let Some(token_count) = doc_token_lengths.get(&entry.0) {
                    loaded_doc_tokens.insert(entry.0, *token_count);
                    doc_ids.insert(entry.0);
                    true
                } else {
                    removed_entries.push(*entry);
                    false
                }
            });

            if !removed_entries.is_empty() {
                let size_decrease = if posting.1.is_empty() {
                    empty_tokens.push((bucket_id, posting.key().clone()));
                    cbor_serialized_size(&(posting.key(), (bucket_id, &removed_entries))) + 2
                } else {
                    removed_entries
                        .iter()
                        .map(|entry| cbor_serialized_size(entry) + 2)
                        .sum()
                };
                *bucket_size_decrease.entry(bucket_id).or_default() += size_decrease;
            }
        }

        for (bucket_id, token) in empty_tokens {
            self.postings.remove(&token);
            if let Some(mut bucket) = self.buckets.get_mut(&bucket_id) {
                bucket.tokens.swap_remove_if(|k| k == &token);
            }
        }

        for (bucket_id, size_decrease) in bucket_size_decrease {
            if let Some(mut bucket) = self.buckets.get_mut(&bucket_id) {
                bucket.size = bucket.size.saturating_sub(size_decrease);
                bucket.mark_dirty();
            }
        }

        self.doc_tokens.clear();
        self.doc_tokens.extend(loaded_doc_tokens);

        let bucket_ids: Vec<u32> = self.buckets.iter().map(|b| *b.key()).collect();
        for bucket_id in bucket_ids {
            if let Some(mut bucket) = self.buckets.get_mut(&bucket_id) {
                let doc_ids = doc_ids_by_bucket.remove(&bucket_id).unwrap_or_default();
                if bucket.doc_ids != doc_ids {
                    bucket.doc_ids = doc_ids;
                    bucket.mark_dirty();
                }
            }
        }

        let total_tokens: usize = self.doc_tokens.iter().map(|r| *r.value()).sum();
        self.total_tokens
            .store(total_tokens as u64, Ordering::Relaxed);

        let doc_count = self.doc_tokens.len();
        let avg = if doc_count == 0 {
            0.0
        } else {
            total_tokens as f32 / doc_count as f32
        };
        *self.avg_doc_tokens.write() = avg;

        Ok(())
    }

    /// Returns the number of documents in the index
    pub fn len(&self) -> usize {
        self.doc_tokens.len()
    }

    /// Returns whether the index is empty
    pub fn is_empty(&self) -> bool {
        self.doc_tokens.is_empty()
    }

    /// Returns the index name
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns the index metadata
    pub fn metadata(&self) -> BM25Metadata {
        let mut metadata = self.metadata.read().clone();
        self.refresh_live_stats(&mut metadata.stats);
        metadata
    }

    /// Gets current statistics about the index
    ///
    /// # Returns
    ///
    /// * `IndexStats` - Current statistics
    pub fn stats(&self) -> BM25Stats {
        let mut stats = self.metadata.read().stats.clone();
        self.refresh_live_stats(&mut stats);
        stats
    }

    /// Overlays the live atomic/lock-protected counters onto a snapshot of the
    /// persisted statistics so callers always observe up-to-date values.
    fn refresh_live_stats(&self, stats: &mut BM25Stats) {
        stats.search_count = self.search_count.load(Ordering::Relaxed);
        stats.num_elements = self.doc_tokens.len() as u64;
        stats.max_bucket_id = self.max_bucket_id.load(Ordering::Relaxed);
        stats.max_document_id = self.max_document_id.load(Ordering::Relaxed);
        stats.avg_doc_tokens = *self.avg_doc_tokens.read();
    }

    /// Inserts a document into the index.
    ///
    /// The text is tokenized with a clone of the index's tokenizer; token
    /// frequencies and the document length (total token count) are then used
    /// to update both the posting list and the running `avg_doc_tokens`
    /// statistic. Updates to buckets are staged and then applied in a second
    /// phase so that at most one bucket is marked dirty per affected bucket.
    ///
    /// # Arguments
    ///
    /// * `id` — unique, caller-assigned document identifier.
    /// * `text` — document text to index.
    /// * `now_ms` — wall-clock time in milliseconds, stored in
    ///   `stats.last_inserted`.
    ///
    /// # Errors
    ///
    /// * [`BM25Error::TokenizeFailed`] if the tokenizer produces no tokens.
    /// * [`BM25Error::AlreadyExists`] if `id` is already present.
    ///
    /// # Concurrency
    ///
    /// Safe to call concurrently with other `insert`/`remove`/`search` calls.
    pub fn insert(&self, id: u64, text: &str, now_ms: u64) -> Result<(), BM25Error> {
        // Tokenize the document
        let token_freqs = {
            let mut tokenizer = self.tokenizer.clone();
            collect_tokens(&mut tokenizer, text, None)
        };

        // Count token frequencies
        if token_freqs.is_empty() {
            return Err(BM25Error::TokenizeFailed {
                name: self.name.clone(),
                id,
                text: text.to_string(),
            });
        }

        // Phase 1: Update the postings collection
        let bucket_id = self.max_bucket_id.load(Ordering::Acquire);
        let tokens: usize = token_freqs.values().sum();
        // buckets_to_update: FxHashMap<bucketid, FxHashMap<token, size_increase>>
        let mut buckets_to_update: FxHashMap<u32, FxHashMap<String, usize>> = FxHashMap::default();
        match self.doc_tokens.entry(id) {
            dashmap::Entry::Occupied(_) => {
                return Err(BM25Error::AlreadyExists {
                    name: self.name.clone(),
                    id,
                });
            }
            dashmap::Entry::Vacant(v) => {
                v.insert(tokens);
                let _ = self.max_document_id.fetch_max(id, Ordering::Relaxed);

                {
                    // Recalculate average document length using consistent snapshots.
                    // Reading `doc_tokens.len()` AFTER inserting the entry above ensures
                    // it reflects this insertion. Concurrent inserts may make the avg
                    // briefly approximate, but it converges as updates settle.
                    let new_total = self
                        .total_tokens
                        .fetch_add(tokens as u64, Ordering::Relaxed)
                        + tokens as u64;
                    let doc_count = self.doc_tokens.len().max(1) as f32;
                    *self.avg_doc_tokens.write() = new_total as f32 / doc_count;
                }

                // Update inverted index
                for (token, freq) in token_freqs {
                    match self.postings.entry(token.clone()) {
                        dashmap::Entry::Occupied(mut entry) => {
                            let val = (id, freq);
                            let e = entry.get_mut();
                            // `push` is a no-op when the exact (doc, freq) pair is
                            // already present (a stale entry left by a remove() with
                            // non-original text). Don't count its size again, but
                            // still mark the bucket dirty below so the refreshed
                            // doc_tokens snapshot gets persisted.
                            let size_increase = if e.1.push(val) {
                                cbor_serialized_size(&val) + 2
                            } else {
                                0
                            };
                            let b = buckets_to_update.entry(e.0).or_default();
                            b.insert(token, size_increase);
                        }
                        dashmap::Entry::Vacant(entry) => {
                            // Create new posting
                            let val = (bucket_id, vec![(id, freq)].into());
                            let size_increase =
                                cbor_serialized_size(&(&token, (bucket_id, &[(id, freq)]))) + 2;
                            entry.insert(val);
                            let b = buckets_to_update.entry(bucket_id).or_default();
                            b.insert(token, size_increase);
                        }
                    };
                }
            }
        }

        // Phase 2: Update bucket states
        // tokens_to_migrate: (old_bucket_id, token, size)
        let mut tokens_to_migrate: Vec<(u32, String, usize)> = Vec::new();
        for (bid, val) in buckets_to_update {
            let mut bucket = self.buckets.entry(bid).or_default();
            // Mark as dirty, needs to be persisted
            bucket.mark_dirty();
            let mut bucket_contains_doc = false;
            for (token, size) in val {
                if bucket.tokens.contains(&token) {
                    // Token already tracked in this bucket; just account for the new posting entry.
                    bucket.size += size;
                    bucket_contains_doc = true;
                } else if bucket.tokens.is_empty()
                    || bucket.size + size < self.config.bucket_overload_size
                {
                    bucket.tokens.push(token);
                    bucket.size += size;
                    bucket_contains_doc = true;
                } else {
                    tokens_to_migrate.push((bid, token, size));
                }
            }
            if bucket_contains_doc {
                bucket.doc_ids.insert(id);
            }
        }

        // Phase 3: Create new buckets if needed
        if !tokens_to_migrate.is_empty() {
            let mut next_bucket_id = self.max_bucket_id.fetch_add(1, Ordering::Release) + 1;

            for (old_bucket_id, token, size) in tokens_to_migrate {
                if let Some(mut posting) = self.postings.get_mut(&token) {
                    posting.0 = next_bucket_id;
                }

                if let Some(mut ob) = self.buckets.get_mut(&old_bucket_id)
                    && ob.tokens.swap_remove_if(|k| &token == k).is_some()
                {
                    ob.size = ob.size.saturating_sub(size);
                    ob.mark_dirty();
                }

                let mut next_new_bucket = false;
                {
                    let mut nb = self.buckets.entry(next_bucket_id).or_default();

                    if nb.tokens.is_empty() || nb.size + size < self.config.bucket_overload_size {
                        // Bucket has enough space, update directly
                        nb.mark_dirty();
                        nb.size += size;
                        nb.tokens.push(token.clone());
                        nb.doc_ids.insert(id);
                    } else {
                        // Bucket doesn't have enough space, need to migrate to the next bucket
                        next_new_bucket = true;
                    }
                }

                if next_new_bucket {
                    next_bucket_id = self.max_bucket_id.fetch_add(1, Ordering::Release) + 1;
                    // update the posting's bucket_id again
                    if let Some(mut posting) = self.postings.get_mut(&token) {
                        posting.0 = next_bucket_id;
                    }
                    let mut nb = self.buckets.entry(next_bucket_id).or_default();
                    nb.mark_dirty();
                    nb.size += size;
                    nb.tokens.push(token.clone());
                    nb.doc_ids.insert(id);
                }
            }
        }

        self.update_metadata(|m| {
            m.stats.version += 1;
            m.stats.last_inserted = now_ms;
            m.stats.insert_count += 1;
        });

        Ok(())
    }

    /// Removes a document from the index.
    ///
    /// The caller must provide the *original text* that was used on
    /// [`insert`](Self::insert); it is re-tokenized to identify which posting
    /// lists should drop this document. If the text does not match, postings
    /// may retain stale entries — searches still skip them because scoring
    /// filters by `doc_tokens` membership, and the stale entries are pruned
    /// the next time the index is loaded via
    /// [`load_buckets`](Self::load_buckets).
    ///
    /// # Arguments
    ///
    /// * `id` — identifier of the document to remove.
    /// * `text` — original text of the document.
    /// * `now_ms` — wall-clock time, stored in `stats.last_deleted`.
    ///
    /// # Returns
    ///
    /// * `true` if a document with the given id was found and removed.
    /// * `false` otherwise.
    pub fn remove(&self, id: u64, text: &str, now_ms: u64) -> bool {
        let removed_tokens = match self.doc_tokens.remove(&id) {
            Some((_k, v)) => v,
            None => return false,
        };

        {
            // Recalculate average document length
            let prev_total = self
                .total_tokens
                .fetch_sub(removed_tokens as u64, Ordering::Relaxed);
            let new_total = prev_total.saturating_sub(removed_tokens as u64);
            let remaining = self.doc_tokens.len();
            let new_avg = if remaining == 0 {
                0.0
            } else {
                new_total as f32 / remaining as f32
            };
            *self.avg_doc_tokens.write() = new_avg;
        }

        // Tokenize the document
        let token_freqs = {
            let mut tokenizer = self.tokenizer.clone();
            collect_tokens(&mut tokenizer, text, None)
        };

        // buckets_to_update: FxHashMap<bucketid, FxHashMap<token, size_decrease>>
        let mut buckets_to_update: FxHashMap<u32, FxHashMap<String, usize>> = FxHashMap::default();
        // Remove from inverted index
        let mut maybe_empty_tokens: Vec<String> = Vec::new();
        for (token, _) in token_freqs {
            if let Some(mut posting) = self.postings.get_mut(&token) {
                // Remove every entry for this document. Duplicates can exist
                // when a previous remove() was given non-original text and the
                // document was re-inserted afterwards.
                let mut removed_vals: Vec<(u64, usize)> = Vec::new();
                while let Some(val) = posting.1.swap_remove_if(|&(idx, _)| idx == id) {
                    removed_vals.push(val);
                }
                if removed_vals.is_empty() {
                    continue;
                }

                let size_decrease = if posting.1.is_empty() {
                    maybe_empty_tokens.push(token.clone());
                    cbor_serialized_size(&(&token, (posting.0, &removed_vals))) + 2
                } else {
                    removed_vals
                        .iter()
                        .map(|val| cbor_serialized_size(val) + 2)
                        .sum()
                };
                let b = buckets_to_update.entry(posting.0).or_default();
                b.insert(token, size_decrease);
            }
        }

        // Drop empty postings atomically: a concurrent insert may have appended
        // a new entry after the guard above was released, in which case the
        // posting must survive. `remove_if` re-checks under the shard lock.
        let mut removed_postings: FxHashSet<String> =
            FxHashSet::with_capacity_and_hasher(maybe_empty_tokens.len(), FxBuildHasher);
        for token in maybe_empty_tokens {
            if self
                .postings
                .remove_if(&token, |_, posting| posting.1.is_empty())
                .is_some()
            {
                removed_postings.insert(token);
            }
        }

        for (bucket_id, val) in buckets_to_update {
            if let Some(mut b) = self.buckets.get_mut(&bucket_id) {
                // Mark as dirty, needs to be persisted
                b.mark_dirty();
                for (token, size_decrease) in val {
                    b.size = b.size.saturating_sub(size_decrease);
                    if removed_postings.contains(&token) {
                        b.tokens.swap_remove_if(|k| &token == k);
                    }
                }
                b.doc_ids.remove(&id);
            }
        }

        // Other buckets may still reference this document in their serialized
        // doc_tokens (e.g. stale postings left by a remove() with non-original
        // text); mark them dirty so the next flush drops the reference.
        // Read-scan first to avoid write-locking every shard on each remove.
        let stale_buckets: Vec<u32> = self
            .buckets
            .iter()
            .filter(|bucket| bucket.doc_ids.contains(&id))
            .map(|bucket| *bucket.key())
            .collect();
        for bucket_id in stale_buckets {
            if let Some(mut bucket) = self.buckets.get_mut(&bucket_id)
                && bucket.doc_ids.remove(&id)
            {
                bucket.mark_dirty();
            }
        }

        self.update_metadata(|m| {
            m.stats.version += 1;
            m.stats.last_deleted = now_ms;
            m.stats.delete_count += 1;
        });

        true
    }

    /// Searches the index and returns the highest-scoring documents.
    ///
    /// The query is tokenized with the index's tokenizer. Multiple tokens are
    /// treated as a disjunction (OR). Use [`search_advanced`](Self::search_advanced)
    /// for boolean expressions with `AND` / `OR` / `NOT` and parentheses.
    ///
    /// # Arguments
    ///
    /// * `query` — raw query text.
    /// * `top_k` — maximum number of results to return; `0` yields an empty vector.
    /// * `params` — override the default [`BM25Params`] for this call only.
    ///
    /// # Returns
    ///
    /// A vector of `(document_id, score)` pairs sorted by descending score.
    pub fn search(&self, query: &str, top_k: usize, params: Option<BM25Params>) -> Vec<(u64, f32)> {
        self.search_count.fetch_add(1, Ordering::Relaxed);
        if top_k == 0 {
            return Vec::new();
        }

        let params = params.as_ref().unwrap_or(&self.config.bm25);
        let scored_docs = self.score_term(query.trim(), params);

        Self::top_k_results(scored_docs, top_k)
    }

    /// Searches the index with a boolean query expression.
    ///
    /// Unlike [`search`](Self::search), the query string is first parsed by
    /// [`QueryType::parse`] and may contain `AND`, `OR`, `NOT` operators and
    /// parentheses. Operator precedence is `OR < AND < NOT`; multiple bare
    /// terms default to `OR`.
    ///
    /// # Arguments
    ///
    /// * `query` — e.g. `"(hello AND world) OR (rust AND NOT java)"`.
    /// * `top_k` — maximum number of results to return.
    /// * `params` — optional BM25 parameters override.
    ///
    /// # Returns
    ///
    /// A vector of `(document_id, score)` pairs sorted by descending score.
    pub fn search_advanced(
        &self,
        query: &str,
        top_k: usize,
        params: Option<BM25Params>,
    ) -> Vec<(u64, f32)> {
        self.search_count.fetch_add(1, Ordering::Relaxed);
        if top_k == 0 {
            return Vec::new();
        }

        let query_expr = QueryType::parse(query);
        let params = params.as_ref().unwrap_or(&self.config.bm25);
        let scored_docs = self.execute_query(&query_expr, params, false);

        Self::top_k_results(scored_docs, top_k)
    }

    /// Extracts the top-k results from scored documents using partial sorting.
    /// Uses `select_nth_unstable_by` for O(n + k·log(k)) instead of O(n·log(n)).
    fn top_k_results(scored_docs: FxHashMap<u64, f32>, top_k: usize) -> Vec<(u64, f32)> {
        if top_k == 0 || scored_docs.is_empty() {
            return Vec::new();
        }

        let mut results: Vec<(u64, f32)> = scored_docs.into_iter().collect();
        if results.len() > top_k {
            results.select_nth_unstable_by(top_k - 1, Self::compare_scored_docs);
            results.truncate(top_k);
        }
        results.sort_unstable_by(Self::compare_scored_docs);
        results
    }

    fn compare_scored_docs(a: &(u64, f32), b: &(u64, f32)) -> std::cmp::Ordering {
        b.1.partial_cmp(&a.1)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| a.0.cmp(&b.0))
    }

    /// Execute a query expression, returning a mapping of document IDs to scores
    fn execute_query(
        &self,
        query: &QueryType,
        params: &BM25Params,
        negated_not: bool,
    ) -> FxHashMap<u64, f32> {
        match query {
            QueryType::Term(term) => self.score_term(term, params),
            QueryType::And(subqueries) => self.score_and(subqueries, params),
            QueryType::Or(subqueries) => self.score_or(subqueries, params),
            QueryType::Not(subquery) => self.score_not(subquery, params, negated_not),
        }
    }

    /// Scores a single term (or multi-term query text) using BM25.
    /// Accumulates scores directly without intermediate allocations.
    fn score_term(&self, term: &str, params: &BM25Params) -> FxHashMap<u64, f32> {
        if self.postings.is_empty() {
            return FxHashMap::default();
        }

        // Be defensive against invalid params to avoid NaNs/inf in ranking.
        let defaults = BM25Params::default();
        let k1 = if params.k1.is_finite() {
            params.k1.max(0.0)
        } else {
            defaults.k1
        };
        let b = if params.b.is_finite() {
            params.b.clamp(0.0, 1.0)
        } else {
            defaults.b
        };

        let mut tokenizer = self.tokenizer.clone();
        let query_terms = collect_tokens(&mut tokenizer, term, None);
        if query_terms.is_empty() {
            return FxHashMap::default();
        }

        let doc_count = self.doc_tokens.len() as f32;
        if doc_count == 0.0 {
            return FxHashMap::default();
        }

        let mut scores: FxHashMap<u64, f32> =
            FxHashMap::with_capacity_and_hasher(self.doc_tokens.len().min(1000), FxBuildHasher);
        let avg_doc_tokens = *self.avg_doc_tokens.read();
        let avg_doc_tokens = avg_doc_tokens.max(1.0);

        // Per-token dedup buffer, reused across query terms so a multi-term
        // query does not reallocate a fresh map for every term.
        let mut valid: FxHashMap<u64, (f32, f32)> = FxHashMap::default();
        for query_token in query_terms.keys() {
            if let Some(postings) = self.postings.get(query_token) {
                // Single-pass: collect doc_id -> (tf, doc_len) for valid documents
                // in one sweep over the postings.
                // Filter out deleted / not-loaded documents:
                // `remove()` depends on the caller providing original text; if they don't,
                // postings can become stale. Also, when only part of buckets are loaded,
                // postings might contain docs missing in `doc_tokens`.
                // Keyed by doc_id so a stale duplicate entry (left by a remove()
                // with non-original text followed by a re-insert) cannot be
                // scored twice or inflate the document frequency; the newest
                // (last) entry wins.
                valid.clear();
                valid.reserve(postings.1.len());
                for (doc_id, token_freq) in postings.1.iter() {
                    if let Some(v) = self.doc_tokens.get(doc_id) {
                        valid.insert(*doc_id, (*token_freq as f32, *v as f32));
                    }
                }

                if valid.is_empty() {
                    continue;
                }

                // Classic Okapi BM25: ln(1 + (N - df + 0.5)/(df + 0.5))
                let df = valid.len() as f32;
                let idf = ((doc_count - df + 0.5) / (df + 0.5) + 1.0).ln();

                // Compute BM25 score for each valid document. `drain` empties the
                // map while keeping its allocation for the next query term.
                for (doc_id, (tf, doc_len)) in valid.drain() {
                    let tf_component =
                        (tf * (k1 + 1.0)) / (tf + k1 * (1.0 - b + b * doc_len / avg_doc_tokens));
                    *scores.entry(doc_id).or_default() += idf * tf_component;
                }
            }
        }

        scores
    }

    /// Scores an OR query
    fn score_or(&self, subqueries: &[Box<QueryType>], params: &BM25Params) -> FxHashMap<u64, f32> {
        if subqueries.is_empty() {
            return FxHashMap::default();
        }
        if subqueries.len() == 1 {
            return self.execute_query(&subqueries[0], params, false);
        }

        // Execute all subqueries and merge results
        let mut result = FxHashMap::default();
        for subquery in subqueries {
            let sub_result = self.execute_query(subquery, params, false);

            for (doc_id, score) in sub_result {
                *result.entry(doc_id).or_insert(0.0) += score;
            }
        }

        result
    }

    /// Scores an AND query
    fn score_and(&self, subqueries: &[Box<QueryType>], params: &BM25Params) -> FxHashMap<u64, f32> {
        if subqueries.is_empty() {
            return FxHashMap::default();
        }
        if subqueries.len() == 1 {
            return self.execute_query(&subqueries[0], params, false);
        }

        // Evaluate non-NOT subqueries first so that a leading NOT does not
        // force building the full complement document set: `NOT a AND b`
        // takes the same cheap path as `b AND NOT a`. The result is
        // order-independent (intersection then subtraction).
        let (positives, negatives): (Vec<&QueryType>, Vec<&QueryType>) = subqueries
            .iter()
            .map(|q| q.as_ref())
            .partition(|q| !matches!(q, QueryType::Not(_)));

        let mut result = if let Some((&first, _)) = positives.split_first() {
            self.execute_query(first, params, false)
        } else {
            // All subqueries are NOT: start from the complement of the first
            // and subtract the rest below.
            self.execute_query(negatives[0], params, false)
        };

        // Intersect the remaining positive subqueries, merging scores.
        for &subquery in positives.iter().skip(1) {
            if result.is_empty() {
                return result;
            }
            let sub_result = self.execute_query(subquery, params, false);

            // Keep only documents present in both results, summing their scores
            // in a single pass over the (already intersected, smaller) result.
            result.retain(|doc_id, score| {
                if let Some(sub_score) = sub_result.get(doc_id) {
                    *score += *sub_score;
                    true
                } else {
                    false
                }
            });
        }

        // Subtract documents matching the negated subqueries.
        let skip_negatives = if positives.is_empty() { 1 } else { 0 };
        for &subquery in negatives.iter().skip(skip_negatives) {
            if result.is_empty() {
                return result;
            }
            let excluded = self.execute_query(subquery, params, true);
            for doc_id in excluded.keys() {
                result.remove(doc_id);
            }
        }

        result
    }

    /// Scores a NOT query.
    ///
    /// The subquery is always evaluated in normal (non-negated) mode;
    /// `negated_not` only selects what to return: the matching documents
    /// (the AND caller subtracts them) or their complement. Evaluating the
    /// subquery with the parent's negation flag would mis-handle double
    /// negation such as `a AND NOT (NOT b)`.
    fn score_not(
        &self,
        subquery: &QueryType,
        params: &BM25Params,
        negated_not: bool,
    ) -> FxHashMap<u64, f32> {
        let exclude = self.execute_query(subquery, params, false);
        if negated_not {
            return exclude;
        }

        let mut result = FxHashMap::default();
        for entry in self.doc_tokens.iter() {
            let doc_id = *entry.key();
            if !exclude.contains_key(&doc_id) {
                result.insert(doc_id, 0.0);
            }
        }
        result
    }

    /// Persists metadata and every currently-dirty bucket.
    ///
    /// This is a convenience wrapper that calls [`store_metadata`](Self::store_metadata)
    /// followed by [`store_dirty_buckets`](Self::store_dirty_buckets). The
    /// closure `f` is invoked once per dirty bucket with `(bucket_id, cbor_bytes)`;
    /// returning `Ok(false)` from `f` aborts the bucket loop without producing
    /// an error (useful for co-operative shutdown).
    ///
    /// # Arguments
    ///
    /// * `metadata` — writer that receives the CBOR-encoded metadata blob.
    /// * `now_ms` — wall-clock time stored in `stats.last_saved`.
    /// * `f` — async function used to persist each dirty bucket.
    ///
    /// # Returns
    ///
    /// * `Ok(true)` if anything — metadata, buckets, or both — was written.
    /// * `Ok(false)` if the index is already fully persisted.
    /// * `Err` on serialization or I/O failure.
    pub async fn flush<W: Write, F>(
        &self,
        metadata: W,
        now_ms: u64,
        f: F,
    ) -> Result<bool, BM25Error>
    where
        F: AsyncFnMut(u32, &[u8]) -> Result<bool, BoxError>,
    {
        let meta_saved = self.store_metadata(metadata, now_ms)?;
        let has_dirty = self.has_dirty_buckets();
        if !meta_saved && !has_dirty {
            return Ok(false);
        }

        self.store_dirty_buckets(f).await?;
        Ok(meta_saved || has_dirty)
    }

    /// Returns whether there are dirty buckets pending persistence.
    pub fn has_dirty_buckets(&self) -> bool {
        self.buckets.iter().any(|b| b.is_dirty())
    }

    /// Returns whether metadata has a newer logical version than the last
    /// serialized metadata snapshot.
    pub fn has_pending_metadata_flush(&self) -> bool {
        let current_version = { self.metadata.read().stats.version };
        self.last_saved_version.load(Ordering::Acquire) < current_version
    }

    /// Repacks all tokens into a minimal set of buckets.
    ///
    /// Over the lifetime of an index — especially before bug fixes that tuned
    /// the bucket splitting logic — repeated inserts and removes can leave
    /// behind many under-filled buckets. `compact_buckets` estimates each
    /// posting's serialized CBOR size and performs a Best-Fit-Decreasing bin
    /// packing with [`BM25Config::bucket_overload_size`] as the bin capacity.
    ///
    /// After compaction:
    ///
    /// * bucket ids are reassigned to a contiguous `0..new_count` range;
    /// * every resulting bucket is marked dirty so the next
    ///   [`flush`](Self::flush) will rewrite the full on-disk layout.
    ///
    /// The operation runs in `O(n log n)` over the number of distinct tokens
    /// and is safe to call at any time from a single thread; it should not be
    /// interleaved with concurrent writes.
    ///
    /// # Returns
    ///
    /// `(old_bucket_count, new_bucket_count)`.
    pub fn compact_buckets(&self) -> (usize, usize) {
        let old_count = self.buckets.len();
        if old_count <= 1 {
            return (old_count, old_count);
        }

        // Step 1: Estimate each token's serialized contribution.
        let mut token_sizes: Vec<(String, usize)> = self
            .postings
            .iter()
            .map(|entry| {
                let size = cbor_serialized_size(&(entry.key(), entry.value())) + 2;
                (entry.key().clone(), size)
            })
            .collect();

        if token_sizes.is_empty() {
            self.buckets.clear();
            self.buckets.insert(
                0,
                Bucket {
                    dirty_version: 1,
                    ..Default::default()
                },
            );
            self.max_bucket_id.store(0, Ordering::Relaxed);
            self.update_metadata(|m| {
                m.stats.version += 1;
            });
            return (old_count, 1);
        }

        // Step 2: Sort by size descending for better packing.
        token_sizes.sort_unstable_by_key(|b| std::cmp::Reverse(b.1));

        // Step 3: Best-fit-decreasing bin packing in O(n log n).
        // `by_remaining` maps remaining-capacity -> bin indices. We pick the bin with the
        // smallest remaining capacity that still fits the token (best fit), which keeps
        // bucket count low without scanning all bins per token.
        let limit = self.config.bucket_overload_size;
        // Each bin: (accumulated_size, tokens)
        let mut bins: Vec<(usize, Vec<String>)> = Vec::new();
        // remaining_capacity -> bin indices with that capacity
        let mut by_remaining: std::collections::BTreeMap<usize, Vec<usize>> =
            std::collections::BTreeMap::new();

        for (token, size) in token_sizes {
            // Find smallest remaining capacity >= size + 1 (preserve `<` limit semantics).
            let needed = size.saturating_add(1);
            let chosen = by_remaining
                .range_mut(needed..)
                .next()
                .and_then(|(_, idxs)| idxs.pop().map(|i| (i, idxs.is_empty())));

            match chosen {
                Some((idx, bucket_now_empty)) => {
                    let old_remaining = limit.saturating_sub(bins[idx].0);
                    if bucket_now_empty {
                        by_remaining.remove(&old_remaining);
                    }
                    bins[idx].0 += size;
                    bins[idx].1.push(token);
                    let new_remaining = limit.saturating_sub(bins[idx].0);
                    by_remaining.entry(new_remaining).or_default().push(idx);
                }
                None => {
                    let idx = bins.len();
                    bins.push((size, vec![token]));
                    let new_remaining = limit.saturating_sub(size);
                    by_remaining.entry(new_remaining).or_default().push(idx);
                }
            }
        }

        // Step 4: Rebuild buckets.
        self.buckets.clear();
        let new_count = bins.len();
        let max_id = new_count.saturating_sub(1) as u32;

        for (i, (size, tokens)) in bins.into_iter().enumerate() {
            let bucket_id = i as u32;

            // Update posting references and collect doc_ids.
            let mut doc_ids = FxHashSet::default();
            for token in &tokens {
                if let Some(mut posting) = self.postings.get_mut(token) {
                    posting.0 = bucket_id;
                    for (doc_id, _) in posting.1.iter() {
                        doc_ids.insert(*doc_id);
                    }
                }
            }

            self.buckets.insert(
                bucket_id,
                Bucket {
                    dirty_version: 1,
                    saved_version: 0,
                    size,
                    tokens: tokens.into(),
                    doc_ids,
                },
            );
        }

        self.max_bucket_id.store(max_id, Ordering::Relaxed);
        self.update_metadata(|m| {
            m.stats.version += 1;
        });

        (old_count, new_count)
    }

    /// Stores the index metadata to a writer in CBOR format.
    ///
    /// # Arguments
    ///
    /// * `w` - Any type implementing the [`Write`] trait
    /// * `now_ms` - Current timestamp in milliseconds
    ///
    /// # Returns
    ///
    /// * `Result<bool, BM25Error>` - true if the metadata was saved, false if the version was not updated
    pub fn store_metadata<W: Write>(&self, w: W, now_ms: u64) -> Result<bool, BM25Error> {
        let current_version = { self.metadata.read().stats.version };
        if self.last_saved_version.load(Ordering::Relaxed) >= current_version {
            return Ok(false);
        }

        let mut meta = self.metadata();
        let prev_saved_version = self
            .last_saved_version
            .fetch_max(meta.stats.version, Ordering::Relaxed);
        if prev_saved_version >= meta.stats.version {
            // No need to save if the version is not updated
            return Ok(false);
        }

        meta.stats.last_saved = now_ms.max(meta.stats.last_saved);

        if let Err(err) = cbor2::to_writer(&BM25IndexRef { metadata: &meta }, w) {
            // Serialization failed: revert only if this call still owns the claimed version.
            let _ = self.last_saved_version.compare_exchange(
                meta.stats.version,
                prev_saved_version,
                Ordering::Relaxed,
                Ordering::Relaxed,
            );
            return Err(BM25Error::Serialization {
                name: self.name.clone(),
                source: err.into(),
            });
        }

        self.update_metadata(|m| {
            m.stats.last_saved = meta.stats.last_saved.max(m.stats.last_saved);
        });

        Ok(true)
    }

    /// Stores dirty buckets to persistent storage using the provided async function.
    /// Serializes each dirty bucket synchronously and releases all DashMap locks
    /// before making async persistence calls to minimize lock contention.
    pub async fn store_dirty_buckets<F>(&self, mut f: F) -> Result<(), BM25Error>
    where
        F: AsyncFnMut(u32, &[u8]) -> Result<bool, BoxError>,
    {
        // Collect dirty bucket IDs to avoid holding iter locks during async calls
        let dirty_buckets: Vec<(u32, u64)> = self
            .buckets
            .iter()
            .filter(|b| b.is_dirty())
            .map(|b| (*b.key(), b.dirty_version))
            .collect();

        let mut buf = Vec::with_capacity(4096);
        for (bucket_id, snapshot_version) in dirty_buckets {
            // Serialize within a scoped block to release all DashMap locks before async call
            {
                let bucket = match self.buckets.get(&bucket_id) {
                    Some(b) if b.is_dirty() => b,
                    _ => continue,
                };

                let mut referenced_doc_ids = FxHashSet::default();
                let postings: FxHashMap<_, _> = bucket
                    .tokens
                    .iter()
                    .filter_map(|k| {
                        let posting = self.postings.get(k)?;
                        if posting.0 != bucket_id {
                            return None;
                        }
                        for (doc_id, _) in posting.1.iter() {
                            referenced_doc_ids.insert(*doc_id);
                        }
                        Some((k, posting))
                    })
                    .collect();

                let doc_tokens: FxHashMap<_, _> = referenced_doc_ids
                    .iter()
                    .filter_map(|id| self.doc_tokens.get(id).map(|v| (*id, *v)))
                    .collect();

                buf.clear();
                cbor2::to_writer(
                    &BucketRef {
                        postings: &postings,
                        doc_tokens: &doc_tokens,
                    },
                    &mut buf,
                )
                .map_err(|err| BM25Error::Serialization {
                    name: self.name.clone(),
                    source: err.into(),
                })?;
            } // All DashMap Ref/RefMut guards dropped here

            let conti = f(bucket_id, &buf).await.map_err(|err| BM25Error::Generic {
                name: self.name.clone(),
                source: err,
            })?;

            // Use version-based dirty tracking: only mark as saved up to the snapshot version.
            // If another write incremented dirty_version after our snapshot, the bucket
            // will remain dirty and be re-persisted on the next flush.
            if let Some(mut b) = self.buckets.get_mut(&bucket_id) {
                b.saved_version = b.saved_version.max(snapshot_version);
            }

            if !conti {
                return Ok(());
            }
        }

        Ok(())
    }

    /// Gets the number of tokens for a document by its ID
    pub fn get_doc_tokens(&self, id: u64) -> Option<usize> {
        self.doc_tokens.get(&id).map(|v| *v)
    }

    /// Updates the index metadata
    ///
    /// # Arguments
    ///
    /// * `f` - Function that modifies the metadata
    fn update_metadata<F>(&self, f: F)
    where
        F: FnOnce(&mut BM25Metadata),
    {
        let mut metadata = self.metadata.write();
        f(&mut metadata);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use std::io::{self, Write};
    use std::sync::Arc;
    use tokio::sync::Mutex;

    struct FailingWriter;

    impl Write for FailingWriter {
        fn write(&mut self, _buf: &[u8]) -> io::Result<usize> {
            Err(io::Error::other("writer failed"))
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    // 创建一个简单的测试索引
    fn create_test_index() -> BM25Index<TokenizerChain> {
        let index = BM25Index::new("anda_db_tfs_bm25".to_string(), default_tokenizer(), None);

        // 添加一些测试文档
        index
            .insert(1, "The quick brown fox jumps over the lazy dog", 0)
            .unwrap();
        index
            .insert(2, "A fast brown fox runs past the lazy dog", 0)
            .unwrap();
        index.insert(3, "The lazy dog sleeps all day", 0).unwrap();
        index
            .insert(4, "Quick brown foxes are rare in the wild", 0)
            .unwrap();

        index
    }

    fn encode_bucket_owned(
        postings: FxHashMap<String, PostingValue>,
        doc_tokens: FxHashMap<u64, usize>,
    ) -> Vec<u8> {
        let mut buf = Vec::new();
        cbor2::to_writer(
            &BucketOwned {
                postings,
                doc_tokens,
            },
            &mut buf,
        )
        .unwrap();
        buf
    }

    #[test]
    fn test_insert() {
        let index = create_test_index();
        assert_eq!(index.len(), 4);

        // 测试添加新文档
        index
            .insert(5, "A new document about cats and dogs", 0)
            .unwrap();
        assert_eq!(index.len(), 5);

        // 测试添加已存在的文档ID
        let result = index.insert(3, "This should fail", 0);
        assert!(matches!(
            result,
            Err(BM25Error::AlreadyExists { id: 3, .. })
        ));

        // 测试添加空文档
        let result = index.insert(6, "", 0);
        assert!(matches!(
            result,
            Err(BM25Error::TokenizeFailed { id: 6, .. })
        ));
    }

    #[test]
    fn test_metadata_accessors_empty_compaction_and_writer_error_paths() {
        let load_result: Result<BM25Index<_>, _> =
            BM25Index::load_metadata(default_tokenizer(), &b"not cbor"[..]);
        assert!(matches!(load_result, Err(BM25Error::Serialization { .. })));

        let index = BM25Index::new("empty_bm25".to_string(), default_tokenizer(), None);
        assert_eq!(index.name(), "empty_bm25");
        assert_eq!(index.len(), 0);
        assert!(index.is_empty());
        assert!(index.has_pending_metadata_flush());
        assert_eq!(index.metadata().name, "empty_bm25");

        index.buckets.insert(1, Bucket::default());
        let (old_count, new_count) = index.compact_buckets();
        assert_eq!((old_count, new_count), (2, 1));
        assert_eq!(index.max_bucket_id.load(Ordering::Relaxed), 0);
        assert!(index.has_dirty_buckets());

        let mut writer = FailingWriter;
        assert!(matches!(
            index.store_metadata(&mut writer, 123),
            Err(BM25Error::Serialization { .. })
        ));
    }

    #[test]
    fn test_remove() {
        let index = create_test_index();
        assert_eq!(index.len(), 4);

        // 测试移除存在的文档
        let removed = index.remove(2, "A fast brown fox runs past the lazy dog", 0);
        assert!(removed);
        assert_eq!(index.len(), 3);

        // 测试移除不存在的文档
        let removed = index.remove(99, "This document doesn't exist", 0);
        assert!(!removed);
        assert_eq!(index.len(), 3);
    }

    #[tokio::test]
    async fn test_load_reconciles_duplicate_token_bucket_ownership() {
        let index = BM25Index::new(
            "duplicate_token_load".to_string(),
            default_tokenizer(),
            Some(BM25Config {
                bm25: BM25Params::default(),
                bucket_overload_size: 64,
            }),
        );
        index.insert(1, "alpha", 0).unwrap();

        let mut initial_metadata = Vec::new();
        let mut stale_bucket0 = Vec::new();
        index
            .flush(&mut initial_metadata, 1, async |bucket_id, data| {
                if bucket_id == 0 {
                    stale_bucket0 = data.to_vec();
                }
                Ok(true)
            })
            .await
            .unwrap();
        assert!(!stale_bucket0.is_empty());

        let mut metadata = index.metadata();
        metadata.stats.version += 1;
        metadata.stats.max_bucket_id = 1;
        let mut metadata_buf = Vec::new();
        cbor2::to_writer(
            &BM25IndexRef {
                metadata: &metadata,
            },
            &mut metadata_buf,
        )
        .unwrap();

        let mut newer_postings = FxHashMap::default();
        newer_postings.insert("alpha".to_string(), (1, vec![(1, 1)].into()));
        let newer_bucket1 = encode_bucket_owned(newer_postings, FxHashMap::from_iter([(1, 1)]));

        let mut loaded = BM25Index::load_metadata(default_tokenizer(), &metadata_buf[..]).unwrap();
        loaded
            .load_buckets(async |bucket_id| match bucket_id {
                0 => Ok(Some(stale_bucket0.clone())),
                1 => Ok(Some(newer_bucket1.clone())),
                _ => Ok(None),
            })
            .await
            .unwrap();

        assert_eq!(loaded.postings.get("alpha").unwrap().0, 1);
        assert!(
            !loaded
                .buckets
                .get(&0)
                .unwrap()
                .tokens
                .contains(&"alpha".to_string())
        );
        assert!(loaded.has_dirty_buckets());

        let mut repaired_buckets: HashMap<u32, Vec<u8>> = HashMap::new();
        loaded
            .store_dirty_buckets(async |bucket_id, data| {
                repaired_buckets.insert(bucket_id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        let repaired_bucket0: BucketOwned =
            cbor2::from_reader(&repaired_buckets.get(&0).unwrap()[..]).unwrap();
        assert!(repaired_bucket0.postings.is_empty());
        assert!(repaired_bucket0.doc_tokens.is_empty());

        let reloaded =
            BM25Index::load_all(
                default_tokenizer(),
                &metadata_buf[..],
                async |id| match id {
                    0 => Ok(repaired_buckets.get(&0).cloned()),
                    1 => Ok(Some(newer_bucket1.clone())),
                    _ => Ok(None),
                },
            )
            .await
            .unwrap();

        let results = reloaded.search("alpha", 10, None);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 1);
        assert!(!reloaded.has_dirty_buckets());
    }

    #[test]
    fn test_remove_with_wrong_text_does_not_leak_into_search() {
        let index = create_test_index();

        // remove() currently relies on caller providing the original text.
        // Even if postings are not fully cleaned, search must not return deleted documents.
        let removed = index.remove(2, "totally different text", 0);
        assert!(removed);
        assert_eq!(index.len(), 3);

        let results = index.search("fox", 10, None);
        assert!(!results.iter().any(|(id, _)| *id == 2));
    }

    #[tokio::test]
    async fn test_remove_with_wrong_text_does_not_resurrect_after_reload() {
        let config = BM25Config {
            bm25: BM25Params::default(),
            bucket_overload_size: 64,
        };
        let index = BM25Index::new(
            "remove_wrong_text_reload".to_string(),
            default_tokenizer(),
            Some(config),
        );
        let terms = [
            "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", "india",
            "juliet", "kilo", "lima",
        ];
        let text = terms.join(" ");
        index.insert(1, &text, 0).unwrap();
        assert!(index.stats().max_bucket_id > 0);

        let mut metadata: Vec<u8> = Vec::new();
        let mut buckets: HashMap<u32, Vec<u8>> = HashMap::new();
        index
            .flush(&mut metadata, 1, async |id: u32, data: &[u8]| {
                buckets.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        assert!(index.remove(1, "wrong text", 2));

        metadata.clear();
        index
            .flush(&mut metadata, 3, async |id: u32, data: &[u8]| {
                buckets.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        let loaded_index = BM25Index::load_all(default_tokenizer(), &metadata[..], async |id| {
            Ok(buckets.get(&id).cloned())
        })
        .await
        .unwrap();

        assert_eq!(loaded_index.len(), 0);
        for term in terms {
            assert!(
                loaded_index.search(term, 10, None).is_empty(),
                "removed document was found after reload for term '{term}'"
            );
        }

        // Loading prunes stale posting entries of deleted documents entirely,
        // and the resulting cleanup is flushed on the next store.
        assert!(
            loaded_index.postings.is_empty(),
            "stale postings must be pruned on load"
        );
        assert!(loaded_index.has_dirty_buckets());

        let mut metadata2: Vec<u8> = Vec::new();
        loaded_index
            .flush(&mut metadata2, 4, async |id: u32, data: &[u8]| {
                buckets.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();
        for data in buckets.values() {
            let bucket: BucketOwned = cbor2::from_reader(&data[..]).unwrap();
            assert!(bucket.postings.is_empty());
            assert!(bucket.doc_tokens.is_empty());
        }
    }

    #[test]
    fn test_search() {
        let index = create_test_index();

        // 测试基本搜索功能
        let results = index.search("fox", 10, None);
        assert_eq!(results.len(), 3); // 应该找到3个包含"fox"的文档

        // 检查结果排序 - 文档1和2应该排在前面，因为它们都包含"fox"
        assert!(results.iter().any(|(id, _)| *id == 1));
        assert!(results.iter().any(|(id, _)| *id == 2));
        assert!(results.iter().any(|(id, _)| *id == 4));

        // 测试多词搜索
        let results = index.search("quick fox dog", 10, None);
        assert!(results[0].0 == 1); // 文档1应该排在最前面，因为它同时包含"quick", "fox", "dog"

        // 测试top_k限制
        let results = index.search("dog", 2, None);
        assert_eq!(results.len(), 2); // 应该只返回2个结果，尽管有3个文档包含"dog"

        // 测试空查询
        let results = index.search("", 10, None);
        assert_eq!(results.len(), 0);

        // 测试无匹配查询
        let results = index.search("elephant giraffe", 10, None);
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_search_top_k_zero_returns_empty() {
        let index = create_test_index();

        let basic = index.search("fox", 0, None);
        assert!(basic.is_empty());

        let advanced = index.search_advanced("fox OR dog", 0, None);
        assert!(advanced.is_empty());
    }

    #[test]
    fn test_empty_index() {
        let tokenizer = default_tokenizer();
        let index = BM25Index::new("anda_db_tfs_bm25".to_string(), tokenizer, None);

        assert_eq!(index.len(), 0);
        assert!(index.is_empty());

        // 测试空索引的搜索
        let results = index.search("test", 10, None);
        assert_eq!(results.len(), 0);
    }

    #[tokio::test]
    async fn test_serialization() {
        let index = create_test_index();

        // 创建临时文件
        let mut metadata: Vec<u8> = Vec::new();
        let mut buckets: HashMap<u32, Vec<u8>> = HashMap::new();

        // 保存索引
        index
            .flush(&mut metadata, 0, async |id: u32, data: &[u8]| {
                buckets.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        // 加载索引
        let tokenizer = default_tokenizer();
        let loaded_index = BM25Index::load_all(tokenizer, &metadata[..], async |id| {
            Ok(buckets.get(&id).cloned())
        })
        .await
        .unwrap();

        // 验证加载的索引
        assert_eq!(loaded_index.len(), index.len());

        // 验证搜索结果
        let mut original_results = index.search("fox", 10, None);
        let mut loaded_results = loaded_index.search("fox", 10, None);

        assert_eq!(original_results.len(), loaded_results.len());
        original_results.sort_by_key(|a| a.0);
        loaded_results.sort_by_key(|a| a.0);
        // 比较文档ID和分数（允许浮点数有小误差）
        for i in 0..original_results.len() {
            assert_eq!(original_results[i].0, loaded_results[i].0);
            assert!((original_results[i].1 - loaded_results[i].1).abs() < 0.001);
        }
    }

    #[tokio::test]
    async fn test_flush_persists_dirty_buckets_even_if_metadata_unchanged() {
        let index = create_test_index();
        index.insert(99, "new fox document", 1).unwrap();

        let mut metadata_buf = Vec::new();
        assert!(index.store_metadata(&mut metadata_buf, 2).unwrap());
        assert!(index.has_dirty_buckets());

        let writes = Arc::new(Mutex::new(0usize));
        let writes_clone = writes.clone();
        let mut metadata_buf2 = Vec::new();
        let saved = index
            .flush(&mut metadata_buf2, 3, async move |_, _| {
                let mut g = writes_clone.lock().await;
                *g += 1;
                Ok(true)
            })
            .await
            .unwrap();

        assert!(saved);
        assert!(*writes.lock().await > 0);
        assert!(!index.has_dirty_buckets());
    }

    #[test]
    fn test_bm25_params() {
        // 使用默认参数
        let default_index = create_test_index();

        // 搜索相同的查询
        let default_results = default_index.search("fox", 10, None);
        let custom_results = default_index.search("fox", 10, Some(BM25Params { k1: 1.5, b: 0.75 }));

        // 验证结果数量相同但分数不同
        assert_eq!(default_results.len(), custom_results.len());

        // 至少有一个文档的分数应该不同
        let mut scores_different = false;
        for i in 0..default_results.len() {
            if (default_results[i].1 - custom_results[i].1).abs() > 0.001 {
                scores_different = true;
                break;
            }
        }
        assert!(scores_different);
    }

    #[test]
    fn test_invalid_bm25_params_do_not_produce_non_finite_scores() {
        let index = create_test_index();

        let results = index.search(
            "fox",
            10,
            Some(BM25Params {
                k1: f32::NAN,
                b: f32::INFINITY,
            }),
        );

        assert!(!results.is_empty());
        assert!(results.iter().all(|(_, score)| score.is_finite()));
    }

    #[test]
    fn test_search_advanced() {
        let index = create_test_index();

        // 测试简单的 Term 查询
        let results = index.search_advanced("fox", 10, None);
        assert_eq!(results.len(), 3); // 应该找到3个包含"fox"的文档

        // 测试 AND 查询
        let results = index.search_advanced("fox AND lazy", 10, None);
        assert_eq!(results.len(), 2); // 文档1和2同时包含"fox"和"lazy"
        assert!(results.iter().any(|(id, _)| *id == 1));
        assert!(results.iter().any(|(id, _)| *id == 2));

        // 测试 OR 查询
        let results = index.search_advanced("quick OR fast", 10, None);
        assert_eq!(results.len(), 3); // 文档1包含"quick"，文档2包含"fast"，文档4包含"quick"
        assert!(results.iter().any(|(id, _)| *id == 1));
        assert!(results.iter().any(|(id, _)| *id == 2));
        assert!(results.iter().any(|(id, _)| *id == 4));

        // 测试 NOT 查询
        let results = index.search_advanced("dog AND NOT lazy", 10, None);
        assert_eq!(results.len(), 0); // 所有包含"dog"的文档也包含"lazy"

        // 测试复杂的嵌套查询
        let results = index.search_advanced("(quick OR fast) AND fox", 10, None);
        assert_eq!(results.len(), 3); // 文档1、2和4

        // 测试更复杂的嵌套查询
        let results = index.search_advanced("(brown AND fox) AND NOT (rare OR sleeps)", 10, None);
        assert_eq!(results.len(), 2); // 文档1和2，排除了包含"rare"的文档4和包含"sleeps"的文档3
        assert!(results.iter().any(|(id, _)| *id == 1));
        assert!(results.iter().any(|(id, _)| *id == 2));

        // 测试空查询
        let results = index.search_advanced("", 10, None);
        assert_eq!(results.len(), 0);

        // 测试无匹配查询
        let results = index.search_advanced("elephant AND giraffe", 10, None);
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_search_advanced_with_parentheses() {
        let index = create_test_index();

        // 测试带括号的复杂查询
        let results = index.search_advanced("(fox AND quick) OR (dog AND sleeps)", 10, None);
        assert_eq!(results.len(), 3); // 文档1, 3, 4
        assert!(results.iter().any(|(id, _)| *id == 1));
        assert!(results.iter().any(|(id, _)| *id == 3));
        assert!(results.iter().any(|(id, _)| *id == 4));

        // 测试多层嵌套括号
        let results = index.search_advanced(
            "((brown AND fox) OR (lazy AND sleeps)) AND NOT rare",
            10,
            None,
        );
        assert_eq!(results.len(), 3); // 文档1、2和3，排除了包含"rare"的文档4
        assert!(results.iter().any(|(id, _)| *id == 1));
        assert!(results.iter().any(|(id, _)| *id == 2));
        assert!(results.iter().any(|(id, _)| *id == 3));

        // 测试带括号的 NOT 查询
        let results = index.search_advanced("dog AND NOT (quick OR fast)", 10, None);
        assert_eq!(results.len(), 1); // 只有文档3，因为它包含"dog"但不包含"quick"或"fast"
        assert_eq!(results[0].0, 3);
    }

    #[test]
    fn test_search_advanced_score_ordering() {
        let index = create_test_index();

        // 测试分数排序 - 包含更多匹配词的文档应该排在前面
        let results = index.search_advanced("quick OR fox OR dog", 10, None);
        assert!(results.len() >= 3);

        // 文档1应该排在最前面，因为它同时包含所有三个词
        assert_eq!(results[0].0, 1);

        // 测试 top_k 限制
        let results = index.search_advanced("dog", 2, None);
        assert_eq!(results.len(), 2); // 应该只返回2个结果，尽管有3个文档包含"dog"
    }

    #[test]
    fn test_search_vs_search_advanced() {
        let index = create_test_index();

        // 对于简单查询，search 和 search_advanced 应该返回相似的结果
        let simple_results = index.search("fox", 10, None);
        let advanced_results = index.search_advanced("fox", 10, None);

        assert_eq!(simple_results.len(), advanced_results.len());

        // 检查文档ID是否匹配（不检查分数，因为实现可能略有不同）
        let simple_ids: Vec<u64> = simple_results.iter().map(|(id, _)| *id).collect();
        let advanced_ids: Vec<u64> = advanced_results.iter().map(|(id, _)| *id).collect();

        assert_eq!(simple_ids.len(), advanced_ids.len());
        for id in simple_ids {
            assert!(advanced_ids.contains(&id));
        }

        // 测试多词查询 - search 将它们视为 OR，search_advanced 也应该如此
        let simple_results = index.search("quick fox", 10, None);
        let advanced_results = index.search_advanced("quick OR fox", 10, None);

        // 检查文档ID是否匹配
        let simple_ids: Vec<u64> = simple_results.iter().map(|(id, _)| *id).collect();
        let advanced_ids: Vec<u64> = advanced_results.iter().map(|(id, _)| *id).collect();

        assert_eq!(simple_ids.len(), advanced_ids.len());
        for id in simple_ids {
            assert!(advanced_ids.contains(&id));
        }
    }

    #[test]
    fn test_search_not_alone() {
        let index = create_test_index();
        // NOT fox => 返回所有不含 fox 的文档 (文档3)
        let results = index.search_advanced("NOT fox", 10, None);
        let ids: Vec<u64> = results.iter().map(|(id, _)| *id).collect();
        assert_eq!(ids, vec![3]);
    }

    #[test]
    fn test_double_negation_inside_and() {
        let index = create_test_index();

        // dog AND NOT (NOT lazy) === dog AND lazy => 文档1、2、3
        let results = index.search_advanced("dog AND NOT (NOT lazy)", 10, None);
        let mut ids: Vec<u64> = results.iter().map(|(id, _)| *id).collect();
        ids.sort_unstable();
        assert_eq!(ids, vec![1, 2, 3]);

        // NOT (NOT fox) === fox => 文档1、2、4
        let results = index.search_advanced("NOT (NOT fox)", 10, None);
        let mut ids: Vec<u64> = results.iter().map(|(id, _)| *id).collect();
        ids.sort_unstable();
        assert_eq!(ids, vec![1, 2, 4]);
    }

    #[test]
    fn test_not_first_in_and_matches_not_last() {
        let index = create_test_index();

        // NOT lazy AND fox === fox AND NOT lazy => 只有文档4
        let a = index.search_advanced("NOT lazy AND fox", 10, None);
        let b = index.search_advanced("fox AND NOT lazy", 10, None);
        let mut ids_a: Vec<u64> = a.iter().map(|(id, _)| *id).collect();
        let mut ids_b: Vec<u64> = b.iter().map(|(id, _)| *id).collect();
        ids_a.sort_unstable();
        ids_b.sort_unstable();
        assert_eq!(ids_a, vec![4]);
        assert_eq!(ids_a, ids_b);
    }

    #[test]
    fn test_and_with_only_not_subqueries() {
        let index = create_test_index();

        // NOT fox AND NOT rare => 不含 fox 也不含 rare 的文档 (文档3)
        let results = index.search_advanced("NOT fox AND NOT rare", 10, None);
        let ids: Vec<u64> = results.iter().map(|(id, _)| *id).collect();
        assert_eq!(ids, vec![3]);
    }

    #[test]
    fn test_reinsert_after_remove_with_wrong_text() {
        let index = BM25Index::new("reinsert".to_string(), default_tokenizer(), None);
        index.insert(1, "dog dog cat", 0).unwrap();

        // Remove with non-original text: the "dog"/"cat" postings keep stale entries.
        assert!(index.remove(1, "bird", 0));

        // Re-insert the same id with a different "dog" frequency; the stale
        // posting entry must not be double-counted nor inflate df.
        index.insert(1, "dog dog dog mouse", 0).unwrap();

        let results = index.search("dog", 10, None);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 1);
        assert!(
            results[0].1.is_finite() && results[0].1 > 0.0,
            "score must be positive, got {}",
            results[0].1
        );

        // Removing with the correct text must clear all entries for the doc,
        // including stale duplicates, and drop the now-empty posting.
        assert!(index.remove(1, "dog dog dog mouse", 0));
        assert!(index.search("dog", 10, None).is_empty());
        assert!(index.postings.get("dog").is_none());
    }

    #[test]
    fn test_concurrent_insert_remove_shared_token() {
        use std::thread;

        // Regression test: remove() must not drop a posting that a concurrent
        // insert just appended to (the empty-check and the removal must be
        // atomic). Two writers share the token "shared"; the reader-side
        // assertion in thread B would fail if the posting got lost.
        let index = Arc::new(BM25Index::new(
            "concurrent_shared".to_string(),
            default_tokenizer(),
            None,
        ));

        const ITERS: usize = 500;
        let a = {
            let index = index.clone();
            thread::spawn(move || {
                for _ in 0..ITERS {
                    index.insert(2, "shared alpha", 0).unwrap();
                    assert!(index.remove(2, "shared alpha", 0));
                }
            })
        };
        let b = {
            let index = index.clone();
            thread::spawn(move || {
                for _ in 0..ITERS {
                    index.insert(3, "shared beta", 0).unwrap();
                    let results = index.search("shared", 10, None);
                    assert!(
                        results.iter().any(|(id, _)| *id == 3),
                        "doc 3 must stay searchable while it exists"
                    );
                    assert!(index.remove(3, "shared beta", 0));
                }
            })
        };

        a.join().unwrap();
        b.join().unwrap();

        // After all churn the index must still accept and find new documents.
        index.insert(10, "shared final", 0).unwrap();
        let results = index.search("shared", 10, None);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, 10);
    }

    #[tokio::test]
    async fn test_serialization_with_buckets() {
        // 创建一个带有小桶大小的索引，强制触发分桶
        let tokenizer = default_tokenizer();
        let config = BM25Config {
            bm25: BM25Params::default(),
            bucket_overload_size: 100, // 非常小的桶大小，强制分桶
        };
        let index = BM25Index::new(
            "test_bucket_serialization".to_string(),
            tokenizer,
            Some(config),
        );

        // 添加大量文档，确保触发分桶
        let test_docs = vec![
            (
                1,
                "The quick brown fox jumps over the lazy dog in the forest",
            ),
            (2, "A fast brown fox runs past the lazy dog near the river"),
            (3, "The lazy dog sleeps all day under the warm sun"),
            (4, "Quick brown foxes are rare in the wild mountain regions"),
            (5, "Many foxes hunt at night when the moon is bright"),
            (6, "Dogs and cats are common pets in modern households"),
            (7, "Wild animals like foxes and wolves roam the countryside"),
            (8, "The forest is home to many different species of animals"),
            (9, "Lazy afternoon naps are enjoyed by both dogs and cats"),
            (
                10,
                "Quick movements help foxes catch their prey efficiently",
            ),
        ];

        for (id, text) in test_docs {
            index.insert(id, text, 0).unwrap();
        }

        // 验证确实创建了多个桶
        let original_stats = index.stats();
        println!(
            "Original index has {} buckets",
            original_stats.max_bucket_id + 1
        );
        assert!(original_stats.max_bucket_id > 0, "应该创建了多个桶");

        // 创建存储映射
        let mut metadata: Vec<u8> = Vec::new();
        let mut buckets: HashMap<u32, Vec<u8>> = HashMap::new();

        // 保存索引
        index
            .flush(&mut metadata, 100, async |id: u32, data: &[u8]| {
                println!("Saving bucket {}, size: {}", id, data.len());
                buckets.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        // 验证保存了正确数量的桶
        println!("Saved {} document buckets", buckets.len());
        assert!(buckets.len() > 1, "应该保存了多个文档桶");

        // 验证每个桶的内容
        for (bucket_id, data) in &buckets {
            let bucket: BucketOwned = cbor2::from_reader(&data[..]).unwrap();
            println!("Document bucket {bucket_id} {:?}", bucket.doc_tokens);
            assert!(!bucket.postings.is_empty());

            // 验证倒排索引结构
            for (term, (bucket_ref, doc_list)) in bucket.postings {
                assert_eq!(
                    bucket_ref, *bucket_id,
                    "术语 {} 的桶引用应该指向当前桶",
                    term
                );
                assert!(!doc_list.is_empty(), "术语 {} 的文档列表不应该为空", term);

                for (doc_id, freq) in doc_list.iter() {
                    assert!(*freq > 0, "文档 {} 中术语 {} 的频率应该大于0", doc_id, term);
                }
            }

            // 验证文档token数量的合理性
            for (doc_id, token_count) in bucket.doc_tokens {
                assert!(token_count > 0, "文档 {} 的token数量应该大于0", doc_id);
            }
        }

        // 加载索引
        let tokenizer2 = default_tokenizer();
        let loaded_index = BM25Index::load_all(tokenizer2, &metadata[..], async |id| {
            println!("Loading for bucket {}", id);
            Ok(buckets.get(&id).cloned())
        })
        .await
        .unwrap();

        // 验证加载的索引基本信息
        assert_eq!(loaded_index.len(), index.len(), "文档数量应该一致");

        let loaded_stats = loaded_index.stats();
        assert_eq!(
            loaded_stats.max_bucket_id, original_stats.max_bucket_id,
            "最大桶ID应该一致"
        );
        assert_eq!(
            loaded_stats.max_document_id, original_stats.max_document_id,
            "最大文档ID应该一致"
        );
        assert!(
            (loaded_stats.avg_doc_tokens - original_stats.avg_doc_tokens).abs() < 0.01,
            "平均文档token数应该基本一致"
        );

        // 验证每个文档的token数量
        for i in 1..=10 {
            let original_tokens = index.get_doc_tokens(i);
            let loaded_tokens = loaded_index.get_doc_tokens(i);
            assert_eq!(
                original_tokens, loaded_tokens,
                "文档 {} 的token数量应该一致",
                i
            );
        }

        // 验证多种搜索查询的结果一致性
        let test_queries = vec![
            "fox",
            "dog",
            "lazy",
            "quick brown",
            "fox AND dog",
            "brown OR lazy",
            "fox AND NOT lazy",
            "(quick OR fast) AND fox",
        ];

        for query in test_queries {
            println!("Testing query: {}", query);

            let original_results =
                if query.contains("AND") || query.contains("OR") || query.contains("NOT") {
                    index.search_advanced(query, 10, None)
                } else {
                    index.search(query, 10, None)
                };

            let loaded_results =
                if query.contains("AND") || query.contains("OR") || query.contains("NOT") {
                    loaded_index.search_advanced(query, 10, None)
                } else {
                    loaded_index.search(query, 10, None)
                };

            assert_eq!(
                original_results.len(),
                loaded_results.len(),
                "查询 '{}' 的结果数量应该一致",
                query
            );

            // 按文档ID排序后比较
            let mut orig_sorted = original_results.clone();
            let mut loaded_sorted = loaded_results.clone();
            orig_sorted.sort_by_key(|a| a.0);
            loaded_sorted.sort_by_key(|a| a.0);

            for i in 0..orig_sorted.len() {
                assert_eq!(
                    orig_sorted[i].0, loaded_sorted[i].0,
                    "查询 '{}' 的第 {} 个结果文档ID应该一致",
                    query, i
                );
                assert!(
                    (orig_sorted[i].1 - loaded_sorted[i].1).abs() < 0.001,
                    "查询 '{}' 的第 {} 个结果分数应该基本一致，原始: {}, 加载: {}",
                    query,
                    i,
                    orig_sorted[i].1,
                    loaded_sorted[i].1
                );
            }
        }

        // 验证倒排索引的完整性 - 检查一些关键词的倒排列表
        let key_terms = vec!["fox", "dog", "lazy", "brown", "quick"];
        for term in key_terms {
            let original_postings = index.postings.get(term);
            let loaded_postings = loaded_index.postings.get(term);

            match (original_postings, loaded_postings) {
                (Some(orig), Some(loaded)) => {
                    // 比较倒排列表内容
                    assert_eq!(
                        orig.1.len(),
                        loaded.1.len(),
                        "术语 '{}' 的倒排列表长度应该一致",
                        term
                    );

                    let mut orig_docs: Vec<_> = orig.1.iter().collect();
                    let mut loaded_docs: Vec<_> = loaded.1.iter().collect();
                    orig_docs.sort();
                    loaded_docs.sort();

                    for i in 0..orig_docs.len() {
                        assert_eq!(
                            orig_docs[i], loaded_docs[i],
                            "术语 '{}' 的第 {} 个倒排项应该一致",
                            term, i
                        );
                    }
                }
                (None, None) => {
                    // 都没有该术语，正常
                }
                _ => {
                    panic!("术语 '{}' 在原始索引和加载索引中的存在性不一致", term);
                }
            }
        }

        println!("所有分桶序列化测试通过！");

        {
            // 测试只加载部分桶的情况
            let tokenizer = default_tokenizer();
            let partial_index = BM25Index::load_all(tokenizer, &metadata[..], async |id| {
                // 只加载桶0的文档
                if id == 0 {
                    Ok(buckets.get(&id).cloned())
                } else {
                    Ok(None)
                }
            })
            .await
            .unwrap();

            // 部分加载会载入桶0 posting 需要的文档长度；如果桶0包含高频词，
            // 它可能覆盖全部文档，但搜索结果仍不应超过完整索引。
            assert!(partial_index.len() <= index.len());

            // 验证部分搜索结果
            let partial_results = partial_index.search("fox", 10, None);
            let full_results = index.search("fox", 10, None);

            // 部分结果应该是完整结果的子集
            assert!(partial_results.len() <= full_results.len());

            for (doc_id, _) in partial_results {
                assert!(
                    full_results.iter().any(|(id, _)| *id == doc_id),
                    "部分加载结果中的文档 {} 应该存在于完整结果中",
                    doc_id
                );
            }

            println!("加载部分分桶测试通过！");
        }
    }

    #[tokio::test]
    async fn test_partial_load_keeps_doc_tokens_with_existing_token_bucket() {
        let config = BM25Config {
            bm25: BM25Params::default(),
            bucket_overload_size: 64,
        };
        let index = BM25Index::new(
            "partial_load_doc_tokens".to_string(),
            default_tokenizer(),
            Some(config),
        );

        index.insert(1, "alpha bravo", 0).unwrap();
        let alpha_bucket = index.postings.get("alpha").unwrap().0;

        let filler_docs = [
            (2, "charlie delta echo foxtrot"),
            (3, "golf hotel india juliet"),
            (4, "kilo lima mike november"),
            (5, "oscar papa quebec romeo"),
            (6, "sierra tango uniform victor"),
            (7, "whiskey xray yankee zulu"),
        ];
        for (id, text) in filler_docs {
            index.insert(id, text, 0).unwrap();
        }
        assert!(index.stats().max_bucket_id > alpha_bucket);

        index.insert(99, "alpha alpha", 0).unwrap();
        assert_eq!(index.postings.get("alpha").unwrap().0, alpha_bucket);

        let mut metadata: Vec<u8> = Vec::new();
        let mut buckets: HashMap<u32, Vec<u8>> = HashMap::new();
        index
            .flush(&mut metadata, 1, async |id: u32, data: &[u8]| {
                buckets.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        let partial_index = BM25Index::load_all(default_tokenizer(), &metadata[..], async |id| {
            if id == alpha_bucket {
                Ok(buckets.get(&id).cloned())
            } else {
                Ok(None)
            }
        })
        .await
        .unwrap();

        assert_eq!(partial_index.get_doc_tokens(99), Some(2));
        let results = partial_index.search("alpha", 10, None);
        assert!(results.iter().any(|(id, _)| *id == 99));
    }

    #[test]
    fn test_no_excessive_small_buckets() {
        // Regression test: existing tokens in a bucket must NOT trigger migration,
        // otherwise each insert after the bucket reaches the limit creates many
        // tiny new buckets.
        let tokenizer = default_tokenizer();
        let config = BM25Config {
            bm25: BM25Params::default(),
            bucket_overload_size: 200, // small limit to trigger splits quickly
        };
        let index = BM25Index::new("small_bucket_test".to_string(), tokenizer, Some(config));

        // Insert many documents sharing common tokens
        let docs = vec![
            (1, "the quick brown fox"),
            (2, "the lazy brown dog"),
            (3, "the quick red cat"),
            (4, "a lazy brown fox jumps"),
            (5, "the brown dog runs fast"),
            (6, "a quick fox hunts at night"),
            (7, "the lazy cat sleeps all day"),
            (8, "brown dogs and brown cats"),
            (9, "quick movements help foxes"),
            (10, "the fast dog chases the fox"),
            (11, "lazy afternoons with brown dogs"),
            (12, "quick brown fox returns again"),
            (13, "the old brown dog rests"),
            (14, "a new quick fox appears"),
            (15, "brown and lazy describe the dog"),
        ];

        for (id, text) in &docs {
            index.insert(*id, text, 0).unwrap();
        }

        let stats = index.stats();
        let num_buckets = stats.max_bucket_id + 1;
        println!(
            "docs={}, buckets={}, max_bucket_id={}",
            docs.len(),
            num_buckets,
            stats.max_bucket_id
        );

        // With 15 short documents and 200-byte limit, we expect a modest number
        // of buckets — certainly not one per insert.
        assert!(
            (num_buckets as usize) < docs.len(),
            "Too many buckets ({num_buckets}) for {} documents — \
             existing tokens are likely being migrated incorrectly",
            docs.len()
        );

        // Verify all documents are still searchable
        for (id, text) in &docs {
            let first_word = text.split_whitespace().find(|w| w.len() > 2).unwrap();
            let results = index.search(first_word, 20, None);
            assert!(
                results.iter().any(|(rid, _)| *rid == *id),
                "doc {} not found when searching for '{}'",
                id,
                first_word
            );
        }
    }

    #[tokio::test]
    async fn test_compact_buckets() {
        // Simulate the real-world scenario: the configured limit is large, but the old
        // bucket-splitting bug created many tiny buckets anyway.
        // We build the index with a tiny limit (to generate fragmentation), then
        // serialize, reload with the correct large limit, and compact.
        let tokenizer = default_tokenizer();
        let small_config = BM25Config {
            bm25: BM25Params::default(),
            bucket_overload_size: 50, // tiny limit to force many buckets
        };
        let index = BM25Index::new("compact_test".to_string(), tokenizer, Some(small_config));

        let docs = vec![
            (1, "the quick brown fox jumps over the lazy dog"),
            (2, "a fast brown fox runs past the lazy dog"),
            (3, "the lazy dog sleeps all day long"),
            (4, "quick brown foxes are rare in the wild"),
            (5, "many foxes hunt at night when the moon is bright"),
            (6, "dogs and cats are common pets in modern households"),
            (7, "wild animals like foxes and wolves roam the countryside"),
            (8, "the forest is home to many different species of animals"),
        ];

        for (id, text) in &docs {
            index.insert(*id, text, 0).unwrap();
        }

        let bucket_count_before = index.stats().max_bucket_id + 1;
        println!("Before compact: {} buckets", bucket_count_before);
        assert!(
            bucket_count_before > 3,
            "should have many fragmented buckets"
        );

        // Serialize fragmented index
        let mut metadata_buf = Vec::new();
        let mut bucket_data: HashMap<u32, Vec<u8>> = HashMap::new();
        index
            .flush(&mut metadata_buf, 1, async |id: u32, data: &[u8]| {
                bucket_data.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        // Reload with the correct (large) bucket limit
        let mut loaded = BM25Index::load_metadata(default_tokenizer(), &metadata_buf[..]).unwrap();
        loaded.config.bucket_overload_size = 1024 * 512;
        loaded.metadata.write().config.bucket_overload_size = 1024 * 512;
        loaded
            .load_buckets(async |id| Ok(bucket_data.get(&id).cloned()))
            .await
            .unwrap();

        let bucket_count_loaded = loaded.stats().max_bucket_id + 1;
        assert_eq!(bucket_count_loaded, bucket_count_before);

        // Capture search results before compaction
        let queries = ["fox", "dog", "lazy brown", "quick OR fast"];
        let results_before: Vec<Vec<(u64, f32)>> = queries
            .iter()
            .map(|q| {
                if q.contains("OR") {
                    loaded.search_advanced(q, 20, None)
                } else {
                    loaded.search(q, 20, None)
                }
            })
            .collect();

        // Compact!
        let (old, new) = loaded.compact_buckets();
        println!("Compacted: {} -> {} buckets", old, new);
        assert!(
            new < old,
            "compaction should reduce bucket count significantly"
        );
        assert!(
            new <= 3,
            "with 512K limit all postings should fit in very few buckets, got {}",
            new,
        );

        // Verify search results are unchanged
        for (i, q) in queries.iter().enumerate() {
            let results_after = if q.contains("OR") {
                loaded.search_advanced(q, 20, None)
            } else {
                loaded.search(q, 20, None)
            };
            assert_eq!(
                results_before[i].len(),
                results_after.len(),
                "query '{}' result count changed after compaction",
                q
            );

            let mut before_sorted = results_before[i].clone();
            let mut after_sorted = results_after.clone();
            before_sorted.sort_by_key(|a| a.0);
            after_sorted.sort_by_key(|a| a.0);
            for j in 0..before_sorted.len() {
                assert_eq!(before_sorted[j].0, after_sorted[j].0);
                assert!(
                    (before_sorted[j].1 - after_sorted[j].1).abs() < 0.001,
                    "query '{}' scores diverged for doc {}",
                    q,
                    before_sorted[j].0
                );
            }
        }

        // Verify flush + reload works after compaction
        let mut metadata_buf2 = Vec::new();
        let mut bucket_data2: HashMap<u32, Vec<u8>> = HashMap::new();
        loaded
            .flush(&mut metadata_buf2, 200, async |id: u32, data: &[u8]| {
                bucket_data2.insert(id, data.to_vec());
                Ok(true)
            })
            .await
            .unwrap();

        let final_loaded =
            BM25Index::load_all(default_tokenizer(), &metadata_buf2[..], async |id| {
                Ok(bucket_data2.get(&id).cloned())
            })
            .await
            .unwrap();
        assert_eq!(final_loaded.len(), loaded.len());

        for q in &queries {
            let orig = if q.contains("OR") {
                loaded.search_advanced(q, 20, None)
            } else {
                loaded.search(q, 20, None)
            };
            let reloaded = if q.contains("OR") {
                final_loaded.search_advanced(q, 20, None)
            } else {
                final_loaded.search(q, 20, None)
            };
            assert_eq!(
                orig.len(),
                reloaded.len(),
                "query '{}' mismatch after reload",
                q
            );
        }
    }
}