// Copyright 2026 PingCAP, Inc.

use std::collections::HashMap;

use anyhow::{Result, bail};
use tantivy::{TantivyError, Term, query::Bm25StatisticsProvider, schema::Field};

use crate::Query;

/// Accumulates BM25 statistics across multiple Tantivy indexes for a query.
///
/// This is typically used to build `PreparedBm25` and store it into `Query` so
/// that scores from different sources are comparable on the same scale.
///
/// The document count follows Tantivy's BM25 convention: `max_doc` (including
/// empty-placeholder documents), not `num_docs`.
///
/// Token totals are keyed by physical field, and document frequencies are
/// aligned to [`Query::bm25_terms`].
#[non_exhaustive]
pub struct Bm25Stats {
    /// Total number of documents (`max_doc` in Tantivy BM25 convention).
    pub(crate) total_docs: u64,
    /// Total indexed tokens for each scored physical field.
    pub(crate) total_tokens_by_field: HashMap<Field, u64>,
    /// Document frequencies aligned to [`Query::bm25_terms`].
    pub(crate) doc_freqs_by_term_idx: Vec<u64>,
}

impl Bm25Stats {
    /// Create an empty BM25 statistics instance.
    pub fn empty(query: &Query) -> Self {
        Self {
            total_docs: 0,
            total_tokens_by_field: query
                .bm25_fields()
                .iter()
                .copied()
                .map(|field| (field, 0))
                .collect(),
            doc_freqs_by_term_idx: vec![0u64; query.bm25_terms().len()],
        }
    }
}

/// Prepared BM25 statistics and per-term weights for a specific query.
///
/// This is computed by aggregating BM25 statistics across a chosen "stats
/// corpus" (typically multiple Tantivy indexes) so that scores produced by
/// different sources are directly comparable.
///
/// Document frequencies are aligned to [`Query::bm25_terms`].
pub struct PreparedBm25Stats {
    total_docs: u64,
    total_tokens_by_field: HashMap<Field, u64>,
    doc_freqs_by_term_idx: Vec<u64>,
}

impl PreparedBm25Stats {
    pub(super) fn new(
        query_fields: &[Field],
        query_terms: &[Term],
        mut stats: Bm25Stats,
    ) -> Result<Self> {
        if stats.doc_freqs_by_term_idx.len() != query_terms.len() {
            bail!(
                "document frequency count {} does not match query term count {}",
                stats.doc_freqs_by_term_idx.len(),
                query_terms.len()
            );
        }
        for &field in query_fields {
            if !stats.total_tokens_by_field.contains_key(&field) {
                bail!("missing BM25 token total for field {field:?}");
            }
        }

        // Tantivy's BM25 assumes `total_docs > 0` and `avgdl > 0`. In practice,
        // empty indexes and fields with no tokens are valid inputs. Sanitize
        // them before Tantivy builds per-term weights.
        let total_docs = stats.total_docs.max(1);
        for total_tokens in stats.total_tokens_by_field.values_mut() {
            if *total_tokens == 0 {
                // Tantivy computes avgdl as total_tokens / total_docs. Use
                // total_docs (not 1) so avgdl stays 1 when N > 1.
                *total_tokens = total_docs;
            }
        }
        for doc_freq in &mut stats.doc_freqs_by_term_idx {
            *doc_freq = (*doc_freq).min(total_docs);
        }

        Ok(Self {
            total_docs,
            total_tokens_by_field: stats.total_tokens_by_field,
            doc_freqs_by_term_idx: stats.doc_freqs_by_term_idx,
        })
    }
}

// Before persisted statistics are prepared, expose a valid empty corpus
// without populating the query's OnceLock. This leaves room for real index
// statistics to be installed before fallback evaluation begins.
impl Bm25StatisticsProvider for crate::Query {
    fn total_num_tokens(&self, field: Field) -> tantivy::Result<u64> {
        let Some(prepared) = self.prepared_bm25() else {
            return Ok(1);
        };
        prepared
            .total_tokens_by_field
            .get(&field)
            .copied()
            .ok_or_else(|| {
                TantivyError::InvalidArgument(format!(
                    "BM25 token statistics are not prepared for field {:?}",
                    field
                ))
            })
    }

    fn total_num_docs(&self) -> tantivy::Result<u64> {
        Ok(self
            .prepared_bm25()
            .map_or(1, |prepared| prepared.total_docs))
    }

    fn doc_freq(&self, term: &Term) -> tantivy::Result<u64> {
        let Some(prepared) = self.prepared_bm25() else {
            return Ok(0);
        };
        match self.bm25_terms().binary_search(term) {
            Ok(term_idx) => Ok(prepared.doc_freqs_by_term_idx[term_idx]),
            Err(_) => Ok(0),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::query::test_util::{PlainFtsQueryInfo, make_scored_query};

    fn text_document(query: &Query, text: &str) -> tantivy::TantivyDocument {
        let mut document = tantivy::TantivyDocument::new();
        document.add_text(query.document_fields()[0].field(), text);
        document
    }

    #[test]
    fn unprepared_query_uses_empty_bm25_statistics_without_becoming_prepared() -> Result<()> {
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "rust".to_string(),
            ..Default::default()
        });
        let field = query.bm25_fields()[0];
        let term = &query.bm25_terms()[0];

        assert!(query.prepared_bm25().is_none());
        assert_eq!(Bm25StatisticsProvider::total_num_docs(&query)?, 1);
        assert_eq!(Bm25StatisticsProvider::total_num_tokens(&query, field)?, 1);
        assert_eq!(Bm25StatisticsProvider::doc_freq(&query, term)?, 0);
        assert!(query.prepared_bm25().is_none());

        let mut stats = Bm25Stats::empty(&query);
        stats.total_docs = 4;
        stats.total_tokens_by_field.insert(field, 8);
        stats.doc_freqs_by_term_idx[0] = 2;
        assert!(query.prepare_bm25_once(stats)?);
        assert_eq!(Bm25StatisticsProvider::total_num_docs(&query)?, 4);
        assert_eq!(Bm25StatisticsProvider::total_num_tokens(&query, field)?, 8);
        assert_eq!(Bm25StatisticsProvider::doc_freq(&query, term)?, 2);
        Ok(())
    }

    #[test]
    fn prepare_bm25_succeeds_with_empty_corpus() -> Result<()> {
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "rust".to_string(),
            ..Default::default()
        });
        assert!(!query.bm25_terms().is_empty());

        let stats = Bm25Stats::empty(&query);
        assert!(query.prepare_bm25_once(stats)?);

        assert!(query.prepared_bm25().is_some());

        let total_docs = Bm25StatisticsProvider::total_num_docs(&query)?;
        assert!(total_docs >= 1);
        let total_tokens =
            Bm25StatisticsProvider::total_num_tokens(&query, query.bm25_fields()[0])?;
        assert!(total_tokens >= total_docs);
        assert_eq!(
            Bm25StatisticsProvider::doc_freq(&query, &query.bm25_terms()[0])?,
            0
        );

        let mut brute = crate::BruteScoredSearcher::new(&query)?;
        let tantivy::query::DocumentEvaluation::Match(score) =
            brute.evaluate_document(&text_document(&query, "rust rust"))?
        else {
            panic!("prepared query should match the document");
        };
        assert!(score.is_finite());
        assert!(score > 0.0);
        Ok(())
    }

    #[test]
    fn prepare_bm25_succeeds_with_empty_query() -> Result<()> {
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: String::new(),
            ..Default::default()
        });
        assert!(query.bm25_terms().is_empty());

        let stats = Bm25Stats::empty(&query);
        assert!(query.prepare_bm25_once(stats)?);

        let total_docs = Bm25StatisticsProvider::total_num_docs(&query)?;
        assert!(total_docs >= 1);

        // An empty query has no BM25 terms, so its field is not included in
        // `bm25_fields` and no token total is prepared for it.
        let document_field = query.document_fields()[0].field();
        assert!(Bm25StatisticsProvider::total_num_tokens(&query, document_field).is_err());

        let term = Term::from_field_text(document_field, "rust");
        assert_eq!(Bm25StatisticsProvider::doc_freq(&query, &term)?, 0);

        let mut brute = crate::BruteScoredSearcher::new(&query)?;
        assert_eq!(
            brute.evaluate_document(&text_document(&query, "rust rust"))?,
            tantivy::query::DocumentEvaluation::NoMatch
        );
        Ok(())
    }
}