// Copyright 2026 PingCAP, Inc.

use std::{
    collections::HashSet,
    sync::{Arc, OnceLock},
};

use anyhow::{Result, anyhow, bail};
use kvenginepb::fts::FullTextIndexFieldDef;
use tantivy::{
    Term,
    query::{BooleanQuery, Occur, Query as TantivyQuery, TermQuery},
    schema::{Field, Schema},
    tokenizer::TokenizerManager,
};
use tidb_query_datatype::FieldTypeAccessor;

mod bm25;
pub use bm25::{Bm25Stats, PreparedBm25Stats};
mod expr;
mod tantivy_query;
pub use expr::{
    FTS_MATCH_EXPR_METADATA_VERSION, FtsMatchExpr, FtsMatchExprKind, FtsMatchFieldRef,
    FtsMatchPredicate, FtsMatchType, FtsScoreMode, parse_match_exprs,
};

/// Represent a Full Text Query. It is the parsed form of a query from the user.
#[derive(Clone)]
pub struct Query(Arc<QueryCore>);

/// Maps one logical TiDB column to one physical Tantivy field.
///
/// A column can occur more than once when the same source value is indexed
/// with multiple analyzers.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct QueryDocumentField {
    column_id: i64,
    field: Field,
    field_type: crate::PhysicalFieldType,
}

impl QueryDocumentField {
    #[inline]
    pub fn column_id(&self) -> i64 {
        self.column_id
    }

    #[inline]
    pub fn field(&self) -> Field {
        self.field
    }

    #[inline]
    pub fn field_type(&self) -> crate::PhysicalFieldType {
        self.field_type
    }
}

struct QueryPhysicalContext {
    schema: Schema,
    tokenizers: TokenizerManager,
    document_fields: Vec<QueryDocumentField>,
}

struct QueryCore {
    info: tipb::FtsQueryInfo,

    /// Executable query for indexed search and single-document evaluation.
    tantivy_query: Box<dyn TantivyQuery + 'static>,

    /// Physical schema and analyzers used to prepare external documents.
    physical_context: QueryPhysicalContext,

    /// Deduplicated terms that may request BM25 document frequencies.
    ///
    /// Terms are sorted by Tantivy's [`Term`] order. Each term's index is its
    /// canonical slot in BM25 document-frequency arrays and supports direct
    /// binary lookup.
    bm25_terms: Vec<Term>,

    /// Deduplicated physical fields that may request BM25 token totals.
    ///
    /// Fields are sorted by Tantivy's [`Field`] order for deterministic
    /// iteration.
    bm25_fields: Vec<Field>,

    /// Prepared shard BM25 statistics for this query (if available).
    ///
    /// When present, indexed search will use this as Tantivy's
    /// `Bm25StatisticsProvider` so that scores from different indexes are on
    /// the same scale.
    prepared_bm25: OnceLock<PreparedBm25Stats>,
}

/// Legacy `query_text` queries always build text terms from the query
/// tokenizer, so they can only address a text field. A field resolved to a
/// native physical type must fail here instead of silently matching nothing:
/// such a field is only reachable through a match expression.
fn ensure_text_field_layout(layout: &crate::SingleFieldLayout) -> Result<()> {
    if layout.field_type() != crate::PhysicalFieldType::Text {
        bail!(
            "FTS field {:?} has physical type {:?} and cannot be queried by legacy query text; \
             it requires a match expression",
            layout.field_name(),
            layout.field_type(),
        );
    }
    Ok(())
}

impl Query {
    /// Validates request metadata shared by every FTS query construction path.
    ///
    /// Legacy `query_text` requests identify their source through exactly one
    /// entry in `columns`. Match-expression requests identify every source
    /// through their expression metadata, so `columns` is not required for
    /// that representation.
    ///
    /// Physical index fields and their layout are validated separately by
    /// [`Query::new_with_index_layout`].
    pub fn validate_info(info: &tipb::FtsQueryInfo) -> Result<()> {
        match info.get_query_type() {
            tipb::FtsQueryType::FtsQueryTypeWithScore | tipb::FtsQueryType::FtsQueryTypeNoScore => {
            }
            _ => {
                bail!(
                    "unsupported fts_query.query_type {:?}",
                    info.get_query_type()
                );
            }
        }
        if !info.has_index_id() {
            bail!("missing fts_query.index_id");
        }
        if !info.get_match_expr().is_empty() {
            let match_expr = parse_match_exprs(info.get_match_expr())?;
            tantivy_query::validate_scored_match_expr(info, &match_expr)?;
            return Ok(());
        }
        if info.get_columns().len() != 1 {
            bail!(
                "legacy FTS query requires exactly one column, got {}",
                info.get_columns().len()
            );
        }
        let fts_col = &info.get_columns()[0];
        if !fts_col.has_column_id() {
            bail!("missing fts_query.column.column_id");
        }
        if !fts_col.as_accessor().is_string_like() {
            bail!(
                "legacy fts_query.column.tp must be string, but got {}",
                fts_col.get_tp()
            );
        }
        Ok(())
    }

    fn new_legacy_with_physical_analysis(
        info: tipb::FtsQueryInfo,
        selected_field: &crate::SingleFieldLayout,
        analysis: (crate::PhysicalIndexSchema, TokenizerManager),
    ) -> Result<Self> {
        if !info.get_match_expr().is_empty() {
            bail!("match_expr queries require index-layout resolution");
        }
        ensure_text_field_layout(selected_field)?;
        if info.get_query_type() == tipb::FtsQueryType::FtsQueryTypeWithScore
            && !selected_field.enable_bm25()
        {
            bail!(
                "FTS field {:?} is filter-only and cannot contribute score",
                selected_field.field_name()
            );
        }
        let mut tokenizer = crate::TOKENIZERS
            .get(info.get_query_tokenizer())
            .ok_or_else(|| anyhow!("unsupported tokenizer {}", info.get_query_tokenizer()))?;

        let mut subqueries = Vec::new();
        let mut bm25_terms = Vec::new();
        let mut seen_bm25_terms = HashSet::new();
        let field = selected_field.field();
        let record_option = if selected_field.enable_bm25() {
            tantivy::schema::IndexRecordOption::WithFreqs
        } else {
            tantivy::schema::IndexRecordOption::Basic
        };
        {
            let mut token_stream = tokenizer.token_stream(info.get_query_text());
            while token_stream.advance() {
                let token = token_stream.token_mut();
                let text = std::mem::take(&mut token.text);

                let term = tantivy::Term::from_field_text(field, &text);
                let term_query = TermQuery::new(term.clone(), record_option);
                let boxed: Box<dyn TantivyQuery + 'static> = Box::new(term_query);
                subqueries.push((Occur::Should, boxed));

                if info.get_query_type() == tipb::FtsQueryType::FtsQueryTypeWithScore
                    && seen_bm25_terms.insert(term.clone())
                {
                    bm25_terms.push(term);
                }
            }
        }
        bm25_terms.sort_unstable();
        let bm25_fields = (!bm25_terms.is_empty())
            .then_some(field)
            .into_iter()
            .collect();
        let (physical_schema, tokenizers) = analysis;
        let schema = physical_schema.tantivy_schema().clone();
        let column_id = info.get_columns()[0].get_column_id();

        Ok(Query(Arc::new(QueryCore {
            info,
            tantivy_query: Box::new(BooleanQuery::new(subqueries)),
            physical_context: QueryPhysicalContext {
                schema,
                tokenizers,
                document_fields: vec![QueryDocumentField {
                    column_id,
                    field,
                    field_type: selected_field.field_type(),
                }],
            },
            bm25_terms,
            bm25_fields,
            prepared_bm25: OnceLock::new(),
        })))
    }

    /// Builds a complete executable query against a persisted physical index
    /// layout.
    ///
    /// Both indexed search and brute-force evaluation require this physical
    /// context, so `Query` does not expose an unresolved construction state.
    pub fn new_with_index_layout(
        info: tipb::FtsQueryInfo,
        tracked_index_fields: &[FullTextIndexFieldDef],
        index_layout: &crate::IndexLayout,
    ) -> Result<Self> {
        Self::validate_info(&info)?;
        if info.get_match_expr().is_empty() {
            if tracked_index_fields.len() != index_layout.fields().len() {
                bail!(
                    "tracked index field count {} does not match index layout field count {}",
                    tracked_index_fields.len(),
                    index_layout.fields().len()
                );
            }
            // `body` is only the legacy physical Tantivy field name. Its
            // tracked column id must match the query source column; otherwise
            // indexed rows would search one column while brute-force fallback
            // rows evaluate another.
            let column_id = info.get_columns()[0].get_column_id();
            let mut matches = tracked_index_fields
                .iter()
                .zip(index_layout.fields())
                .filter(|(index_field, _)| index_field.get_col_id() == column_id);
            let selected_field = matches
                .next()
                .map(|(_, field)| field)
                .ok_or_else(|| anyhow!("column {} is not part of the FTS index", column_id))?;
            if matches.next().is_some() {
                bail!(
                    "legacy FTS query column {} maps to multiple physical index fields",
                    column_id
                );
            }
            let analysis = crate::physical_index_analysis::build_physical_index_analysis(
                tracked_index_fields,
                index_layout,
            )?;
            return Self::new_legacy_with_physical_analysis(info, selected_field, analysis);
        }

        let match_expr = parse_match_exprs(info.get_match_expr())?;
        tantivy_query::validate_scored_match_expr(&info, &match_expr)?;
        Self::new_match_expr_with_index_layout(info, match_expr, tracked_index_fields, index_layout)
    }

    fn new_match_expr_with_index_layout(
        info: tipb::FtsQueryInfo,
        match_expr: FtsMatchExpr,
        tracked_index_fields: &[FullTextIndexFieldDef],
        index_layout: &crate::IndexLayout,
    ) -> Result<Self> {
        let (physical_schema, tokenizers) =
            crate::physical_index_analysis::build_physical_index_analysis(
                tracked_index_fields,
                index_layout,
            )?;
        let schema = physical_schema.tantivy_schema().clone();
        let lowered = tantivy_query::lower_match_expr_query(
            &info,
            &match_expr,
            tracked_index_fields,
            index_layout,
        )?;
        Ok(Query(Arc::new(QueryCore {
            info,
            tantivy_query: lowered.query,
            physical_context: QueryPhysicalContext {
                schema,
                tokenizers,
                document_fields: lowered.document_fields,
            },
            bm25_terms: lowered.bm25_terms,
            bm25_fields: lowered.bm25_fields,
            prepared_bm25: OnceLock::new(),
        })))
    }

    #[inline]
    pub fn info(&self) -> &tipb::FtsQueryInfo {
        &self.0.info
    }

    #[inline]
    fn physical_context(&self) -> &QueryPhysicalContext {
        &self.0.physical_context
    }

    /// Returns the physical Tantivy schema used to prepare external documents.
    #[inline]
    pub fn evaluation_schema(&self) -> &Schema {
        &self.physical_context().schema
    }

    /// Returns the per-index tokenizer registry paired with
    /// [`Self::evaluation_schema`].
    #[inline]
    pub fn document_tokenizers(&self) -> &TokenizerManager {
        &self.physical_context().tokenizers
    }

    /// Returns logical-column to physical-field bindings required to build an
    /// external Tantivy document.
    #[inline]
    pub fn document_fields(&self) -> &[QueryDocumentField] {
        self.physical_context().document_fields.as_slice()
    }

    /// Returns deduplicated, sorted terms that may request BM25 document
    /// frequencies.
    #[inline]
    pub fn bm25_terms(&self) -> &[Term] {
        &self.0.bm25_terms
    }

    /// Returns deduplicated, sorted fields that may request BM25 token totals.
    #[inline]
    pub fn bm25_fields(&self) -> &[Field] {
        &self.0.bm25_fields
    }

    /// Returns the executable Tantivy query.
    #[inline]
    pub fn tantivy_query(&self) -> &(dyn TantivyQuery + 'static) {
        self.0.tantivy_query.as_ref()
    }

    /// Returns the prepared BM25 statistics if available.
    #[inline]
    pub fn prepared_bm25(&self) -> Option<&PreparedBm25Stats> {
        self.0.prepared_bm25.get()
    }

    /// Sets prepared BM25 statistics for this query if it has not been set.
    ///
    /// This is used to make scores produced by different sources directly
    /// comparable by forcing them to use the same global `(N, df, avgdl)`
    /// statistics.
    ///
    /// Returns `true` if the prepared statistics were stored by this call.
    pub fn prepare_bm25_once(&self, stats: Bm25Stats) -> Result<bool> {
        if self.prepared_bm25().is_some() {
            return Ok(false);
        }
        let prepared = PreparedBm25Stats::new(self.bm25_fields(), self.bm25_terms(), stats)?;
        Ok(self.0.prepared_bm25.set(prepared).is_ok())
    }
}

pub mod test_util {
    use kvenginepb::fts::{FullTextAnalyzerType, FullTextIndexFieldDef, NgramOption};

    use super::Query;

    pub const DEFAULT_FTS_TOP_K: u32 = 10;
    pub const DEFAULT_FTS_COLUMN_ID: i64 = 130;
    pub const DEFAULT_FTS_INDEX_ID: i64 = 42;
    pub const DEFAULT_FTS_TOKENIZER: &str = "STANDARD_V1";

    /// Builds an NGRAM full-text field definition for tests.
    pub fn ngram_index_field(
        col_id: i64,
        min_gram: u64,
        max_gram: u64,
        granularity: &str,
        lower_case: bool,
        enable_bm25: bool,
    ) -> FullTextIndexFieldDef {
        let mut field = FullTextIndexFieldDef::new();
        field.set_col_id(col_id);
        let option = field.mut_index_field_option();
        option.set_analyzer_type(FullTextAnalyzerType::FtsAnalyzerNgramV1);
        option.set_enable_bm25(enable_bm25);
        let mut ngram = NgramOption::new();
        ngram.set_min_gram(min_gram);
        ngram.set_max_gram(max_gram);
        ngram.set_granularity(granularity.to_string());
        ngram.set_lower_case(lower_case);
        option.set_ngram(ngram);
        field
    }

    pub struct PlainFtsQueryInfo {
        pub query: String,
        pub top_k: u32,
        pub column_id: i64,
        pub index_id: i64,
        pub tokenizer: String,
    }

    impl Default for PlainFtsQueryInfo {
        fn default() -> Self {
            Self {
                query: String::new(),
                top_k: DEFAULT_FTS_TOP_K,
                column_id: DEFAULT_FTS_COLUMN_ID,
                index_id: DEFAULT_FTS_INDEX_ID,
                tokenizer: DEFAULT_FTS_TOKENIZER.to_string(),
            }
        }
    }

    impl PlainFtsQueryInfo {
        pub fn as_query_info(&self) -> tipb::FtsQueryInfo {
            let mut info = tipb::FtsQueryInfo::new();
            info.set_query_text(self.query.clone());
            info.set_top_k(self.top_k);
            info.set_index_id(self.index_id);
            info.set_query_tokenizer(self.tokenizer.clone());

            // Create a ColumnInfo for the column_id
            let mut column_info = tipb::ColumnInfo::new();
            column_info.set_column_id(self.column_id);
            column_info.set_tp(0xfe); // FieldTypeTp::String
            info.mut_columns().push(column_info);

            info
        }
    }

    /// Creates a new query for searching with scoring.
    pub fn make_scored_query(plain_info: &PlainFtsQueryInfo) -> Query {
        let mut info = plain_info.as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeWithScore);
        make_legacy_query(info)
    }

    /// Creates a new query for searching without scoring.
    pub fn make_unscored_query(plain_info: &PlainFtsQueryInfo) -> Query {
        let mut info = plain_info.as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        make_legacy_query(info)
    }

    fn make_legacy_query(info: tipb::FtsQueryInfo) -> Query {
        let mut index_field =
            crate::physical_index_analysis::legacy_index_field(info.get_query_tokenizer()).unwrap();
        index_field.set_col_id(info.get_columns()[0].get_column_id());
        Query::new_with_index_layout(info, &[index_field], &crate::IndexLayout::legacy_body())
            .unwrap()
    }
}

#[cfg(test)]
mod tests {
    use kvenginepb::fts::{FullTextAnalyzerType, FullTextIndexFieldDef, FullTextIndexFieldOption};
    use tantivy::{query::DocumentEvaluation, schema::Field};

    use super::{
        test_util::{PlainFtsQueryInfo, ngram_index_field},
        *,
    };

    fn field_ref(col_id: i64, index_field: &str) -> serde_json::Value {
        serde_json::json!({
            "col_id": col_id,
            "index_field": index_field,
        })
    }

    fn match_expr_metadata(
        score_mode: &str,
        query: &str,
        fields: Vec<serde_json::Value>,
    ) -> Vec<u8> {
        let metadata = if score_mode == "inherit" {
            serde_json::json!({
                "version": FTS_MATCH_EXPR_METADATA_VERSION,
                "query": query,
                "fields": fields,
            })
        } else {
            serde_json::json!({
                "version": FTS_MATCH_EXPR_METADATA_VERSION,
                "score_mode": score_mode,
                "query": query,
                "fields": fields,
            })
        };
        serde_json::to_vec(&metadata).unwrap()
    }

    fn match_leaf(
        sig: tipb::ScalarFuncSig,
        score_mode: &str,
        query: &str,
        fields: Vec<serde_json::Value>,
    ) -> tipb::Expr {
        let mut expr = tipb::Expr::new();
        expr.set_tp(tipb::ExprType::ScalarFunc);
        expr.set_sig(sig);
        expr.set_val(match_expr_metadata(score_mode, query, fields));
        expr
    }

    fn standard_index_field(col_id: i64, enable_bm25: bool) -> FullTextIndexFieldDef {
        let mut field = FullTextIndexFieldDef::new();
        field.set_col_id(col_id);
        let mut option = FullTextIndexFieldOption::new();
        option.set_analyzer_type(FullTextAnalyzerType::FtsAnalyzerStandardV1);
        option.set_enable_bm25(enable_bm25);
        field.set_index_field_option(option);
        field
    }

    fn doc_ids(results: &[crate::ScoredResult]) -> Vec<u32> {
        results.iter().map(|result| result.doc_id).collect()
    }

    #[test]
    fn test_query_constructor_rejects_invalid_request_metadata() {
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        let column_id = info.get_columns()[0].get_column_id();
        info.clear_index_id();

        let index_fields = vec![standard_index_field(column_id, true)];
        let index_layout = crate::IndexLayout::legacy_body();
        let error = Query::new_with_index_layout(info, &index_fields, &index_layout)
            .err()
            .unwrap();
        assert!(error.to_string().contains("missing fts_query.index_id"));
    }

    #[test]
    fn test_validate_info_rejects_invalid_match_expr() {
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr().push(tipb::Expr::new());

        let error = Query::validate_info(&info).unwrap_err();
        assert!(
            error
                .to_string()
                .contains("unsupported FTS match_expr type")
        );
    }

    #[test]
    fn test_legacy_query_requires_column_metadata() {
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.clear_columns();

        let index_fields = vec![standard_index_field(test_util::DEFAULT_FTS_COLUMN_ID, true)];
        let index_layout = crate::IndexLayout::legacy_body();
        let error = Query::new_with_index_layout(info, &index_fields, &index_layout)
            .err()
            .unwrap();
        assert!(
            error
                .to_string()
                .contains("legacy FTS query requires exactly one column, got 0")
        );
    }

    fn score_for_doc(results: &[crate::ScoredResult], doc_id: u32) -> f32 {
        results
            .iter()
            .find(|result| result.doc_id == doc_id)
            .unwrap()
            .score
    }

    #[test]
    fn test_legacy_query_exposes_physical_evaluation_context() -> Result<()> {
        let mut info = PlainFtsQueryInfo {
            query: "rust".to_string(),
            ..Default::default()
        }
        .as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);

        let index_fields = vec![standard_index_field(test_util::DEFAULT_FTS_COLUMN_ID, true)];
        let query =
            Query::new_with_index_layout(info, &index_fields, &crate::IndexLayout::legacy_body())?;
        assert_eq!(
            query.evaluation_schema().get_field("body")?,
            crate::SINGLE_TEXT_FIELD
        );
        assert!(query.document_tokenizers().get("STANDARD_V1").is_some());
        assert_eq!(query.document_fields().len(), 1);
        assert_eq!(
            query.document_fields()[0].column_id(),
            test_util::DEFAULT_FTS_COLUMN_ID
        );
        assert_eq!(query.document_fields()[0].field(), crate::SINGLE_TEXT_FIELD);
        assert!(query.bm25_terms().is_empty());
        assert!(query.bm25_fields().is_empty());
        Ok(())
    }

    #[test]
    fn test_legacy_query_resolves_multi_field_physical_context() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["101", "202"]);
        let index_fields = vec![
            standard_index_field(101, true),
            standard_index_field(202, true),
        ];
        let mut info = PlainFtsQueryInfo {
            query: "rust".to_string(),
            column_id: 202,
            ..Default::default()
        }
        .as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);

        let query = Query::new_with_index_layout(info, &index_fields, &layout)?;
        assert_eq!(query.evaluation_schema().num_fields(), 2);
        assert!(query.document_tokenizers().get("STANDARD_V1").is_some());
        assert_eq!(query.document_fields().len(), 1);
        assert_eq!(query.document_fields()[0].column_id(), 202);
        assert_eq!(query.document_fields()[0].field(), Field::from_field_id(1));
        Ok(())
    }

    #[test]
    fn test_match_expr_query_is_complete_after_construction() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0"]);
        let index_fields = vec![standard_index_field(101, true)];
        let mut writer =
            crate::TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer.add_document_fields([(Field::from_field_id(0), "rust")])?;
        let reader = crate::IndexReader::from_tantivy_index(writer.finalize()?)?;

        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchWord,
            "inherit",
            "rust",
            vec![field_ref(101, "STANDARD_V1")],
        ));
        info.clear_columns();
        let query = Query::new_with_index_layout(info, &index_fields, &layout)?;
        assert!(query.info().get_columns().is_empty());
        assert_eq!(query.document_fields().len(), 1);
        assert_eq!(query.document_fields()[0].column_id(), 101);
        assert_eq!(reader.search_new(&query)?.count(), 1);
        Ok(())
    }

    mod exact_tests;

    #[test]
    fn test_match_expr_multi_field_brute_score_matches_index() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0", "1"]);
        let index_fields = vec![
            standard_index_field(101, true),
            standard_index_field(202, true),
        ];
        let mut writer =
            crate::TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer.add_document_fields([
            (Field::from_field_id(0), "rust language"),
            (Field::from_field_id(1), "search engine"),
        ])?;
        writer.add_document_fields([
            (Field::from_field_id(0), "rust language"),
            (Field::from_field_id(1), "database"),
        ])?;
        let reader = crate::IndexReader::from_tantivy_index(writer.finalize()?)?;

        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeWithScore);
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchWord,
            "inherit",
            "search",
            vec![field_ref(202, "STANDARD_V1")],
        ));
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchWord,
            "inherit",
            "rust",
            vec![field_ref(101, "STANDARD_V1")],
        ));
        let query = Query::new_with_index_layout(info, &index_fields, &layout)?;
        assert_eq!(query.document_fields().len(), 2);
        assert_eq!(
            query
                .document_fields()
                .iter()
                .map(|binding| binding.field().field_id())
                .collect::<Vec<_>>(),
            vec![0, 1]
        );
        assert_eq!(
            query
                .bm25_fields()
                .iter()
                .map(|field| field.field_id())
                .collect::<Vec<_>>(),
            vec![0, 1]
        );
        assert!(query.bm25_terms().is_sorted());
        assert_eq!(
            query
                .bm25_terms()
                .iter()
                .map(|term| term.field().field_id())
                .collect::<Vec<_>>(),
            vec![0, 1]
        );

        let mut stats = Bm25Stats::empty(&query);
        reader.accumulate_bm25_stats(&query, &mut stats)?;
        assert!(query.prepare_bm25_once(stats)?);
        let mut indexed_results = Vec::new();
        reader.search(&query, &mut indexed_results)?;

        let mut first_document = tantivy::TantivyDocument::new();
        first_document.add_text(Field::from_field_id(0), "rust language");
        first_document.add_text(Field::from_field_id(1), "search engine");
        let mut second_document = tantivy::TantivyDocument::new();
        second_document.add_text(Field::from_field_id(0), "rust language");
        second_document.add_text(Field::from_field_id(1), "database");
        let mut brute = crate::BruteScoredSearcher::new(&query)?;
        let first_evaluation = brute.evaluate_document(&first_document)?;
        let second_evaluation = brute.evaluate_document(&second_document)?;

        assert_eq!(indexed_results.len(), 1);
        let DocumentEvaluation::Match(brute_score) = first_evaluation else {
            panic!("first document should match");
        };
        assert_eq!(second_evaluation, DocumentEvaluation::NoMatch);
        assert!((indexed_results[0].score - brute_score).abs() < 0.001);
        Ok(())
    }

    #[test]
    fn test_match_expr_prefix_matches_with_brute_force() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0"]);
        let index_fields = vec![standard_index_field(101, true)];
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchPrefix,
            "inherit",
            "rus",
            vec![field_ref(101, "STANDARD_V1")],
        ));
        let query = Query::new_with_index_layout(info, &index_fields, &layout)?;

        let mut brute = crate::BruteScoredSearcher::new(&query)?;
        let mut matching_document = tantivy::TantivyDocument::new();
        matching_document.add_text(Field::from_field_id(0), "rust");
        let mut non_matching_document = tantivy::TantivyDocument::new();
        non_matching_document.add_text(Field::from_field_id(0), "database");
        assert_eq!(
            brute.evaluate_document(&matching_document)?,
            DocumentEvaluation::Match(1.0)
        );
        assert_eq!(
            brute.evaluate_document(&non_matching_document)?,
            DocumentEvaluation::NoMatch
        );
        Ok(())
    }

    #[test]
    fn test_match_expr_rejects_scored_multi_token_prefix() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0"]);
        let index_fields = vec![standard_index_field(101, true)];
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeWithScore);
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchPrefix,
            "inherit",
            "rust search",
            vec![field_ref(101, "STANDARD_V1")],
        ));

        let error = match Query::new_with_index_layout(info, &index_fields, &layout) {
            Ok(_) => panic!("expected scored multi-token prefix query to fail"),
            Err(error) => error,
        };
        assert!(
            error
                .to_string()
                .contains("scored multi-token FTS prefix query")
        );
        Ok(())
    }

    #[test]
    fn test_match_expr_ngram_terms_use_and() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0"]);
        let index_fields = vec![ngram_index_field(202, 3, 3, "WORD", false, true)];
        let mut writer =
            crate::TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer.add_document_fields([(Field::from_field_id(0), "database systems")])?;
        writer.add_document_fields([(Field::from_field_id(0), "database storage")])?;
        let reader = crate::IndexReader::from_tantivy_index(writer.finalize()?)?;

        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchWord,
            "inherit",
            "database systems",
            vec![field_ref(202, "NGRAM_V1")],
        ));
        let query = Query::new_with_index_layout(info, &index_fields, &layout)?;
        assert!(
            query
                .document_tokenizers()
                .get("NGRAM_V1_FIELD_0")
                .is_some()
        );

        let mut results = Vec::new();
        reader.search(&query, &mut results)?;
        assert_eq!(doc_ids(&results), vec![0]);
        Ok(())
    }

    #[test]
    fn test_match_expr_filter_leaf_uses_constant_zero_score() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0"]);
        let index_fields = vec![standard_index_field(101, true)];
        let mut writer =
            crate::TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer.add_document_fields([(Field::from_field_id(0), "rust search")])?;
        let reader = crate::IndexReader::from_tantivy_index(writer.finalize()?)?;

        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeWithScore);
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchWord,
            "filter",
            "rust",
            vec![field_ref(101, "STANDARD_V1")],
        ));
        let query = Query::new_with_index_layout(info, &index_fields, &layout)?;
        assert_eq!(query.bm25_terms().len(), 1);
        assert_eq!(query.bm25_fields().len(), 1);

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

        let mut results = Vec::new();
        reader.search(&query, &mut results)?;
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].doc_id, 0);
        assert_eq!(results[0].score, 0.0);

        let mut document = tantivy::TantivyDocument::new();
        document.add_text(Field::from_field_id(0), "rust search");
        let mut brute = crate::BruteScoredSearcher::new(&query)?;
        assert_eq!(
            brute.evaluate_document(&document)?,
            DocumentEvaluation::Match(0.0)
        );
        Ok(())
    }

    #[test]
    fn test_match_expr_global_bm25_stats_cover_multiple_fields() -> Result<()> {
        let layout = crate::IndexLayout::from_field_names(["0", "1"]);
        let index_fields = vec![
            standard_index_field(101, true),
            standard_index_field(202, true),
        ];

        let mut writer_a =
            crate::TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer_a.add_document_fields([
            (Field::from_field_id(0), "rust rust rust"),
            (Field::from_field_id(1), "other"),
        ])?;
        writer_a.add_document_fields([
            (Field::from_field_id(0), "other"),
            (Field::from_field_id(1), "other"),
        ])?;
        let reader_a = crate::IndexReader::from_tantivy_index(writer_a.finalize()?)?;

        let mut writer_b =
            crate::TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer_b.add_document_fields([
            (Field::from_field_id(0), "other"),
            (Field::from_field_id(1), "rust rust"),
        ])?;
        writer_b.add_document_fields([
            (Field::from_field_id(0), "other"),
            (Field::from_field_id(1), "other"),
        ])?;
        let reader_b = crate::IndexReader::from_tantivy_index(writer_b.finalize()?)?;

        let mut writer_combined =
            crate::TantivyIndexWriter::new_in_memory_with_index_fields(&index_fields, &layout)?;
        writer_combined.add_document_fields([
            (Field::from_field_id(0), "rust rust rust"),
            (Field::from_field_id(1), "other"),
        ])?;
        writer_combined.add_document_fields([
            (Field::from_field_id(0), "other"),
            (Field::from_field_id(1), "other"),
        ])?;
        writer_combined.add_document_fields([
            (Field::from_field_id(0), "other"),
            (Field::from_field_id(1), "rust rust"),
        ])?;
        writer_combined.add_document_fields([
            (Field::from_field_id(0), "other"),
            (Field::from_field_id(1), "other"),
        ])?;
        let reader_combined = crate::IndexReader::from_tantivy_index(writer_combined.finalize()?)?;

        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeWithScore);
        info.mut_match_expr().push(match_leaf(
            tipb::ScalarFuncSig::FtsMatchWord,
            "inherit",
            "rust",
            vec![field_ref(101, "STANDARD_V1"), field_ref(202, "STANDARD_V1")],
        ));
        let prepared_query = Query::new_with_index_layout(info.clone(), &index_fields, &layout)?;
        assert_eq!(prepared_query.bm25_terms().len(), 2);
        assert_eq!(prepared_query.bm25_fields().len(), 2);

        let mut stats = Bm25Stats::empty(&prepared_query);
        reader_a.accumulate_bm25_stats(&prepared_query, &mut stats)?;
        reader_b.accumulate_bm25_stats(&prepared_query, &mut stats)?;
        assert!(prepared_query.prepare_bm25_once(stats)?);

        let combined_query = Query::new_with_index_layout(info, &index_fields, &layout)?;
        let mut results_a = Vec::new();
        let mut results_b = Vec::new();
        let mut results_combined = Vec::new();
        reader_a.search(&prepared_query, &mut results_a)?;
        reader_b.search(&prepared_query, &mut results_b)?;
        reader_combined.search(&combined_query, &mut results_combined)?;

        let score_a = score_for_doc(&results_a, 0);
        let score_b = score_for_doc(&results_b, 0);
        let combined_score_a = score_for_doc(&results_combined, 0);
        let combined_score_b = score_for_doc(&results_combined, 2);
        assert!((score_a - combined_score_a).abs() < 1e-6);
        assert!((score_b - combined_score_b).abs() < 1e-6);
        Ok(())
    }
}