use std::ops::Bound;

use anyhow::{Context, Result, anyhow, bail};
use tantivy::{
    Term,
    query::{
        AllQuery, BooleanQuery, EmptyQuery, FuzzyTermQuery, Occur, PhraseQuery, Query, RangeQuery,
        TermQuery, TermSetQuery,
    },
    schema::{Field, IndexRecordOption, Schema, Type},
};
use tici_common::codec::number::NumberDecoder;

use crate::{
    data_types::codec::encode_decimal_val,
    reader::query_result::SearchRows,
    utils::{
        key::ShardingKeyRange,
        schema::{
            AnalyzedToken, ColumnInfoExt, HANDLE_STORE_NAME, IndexMeta, PhraseAnalysis, SortField,
            VERSION_STORE_NAME, parse_column_id,
        },
    },
};

pub const BM25_SCORE_FIELD_NAME: &str = "column_-2050";

#[derive(Clone)]
pub struct DocField {
    pub name: String,
    pub field: Field,
    pub value_type: Type,
}

impl std::fmt::Debug for DocField {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}/{:?}/{:?}", self.name, self.field, self.value_type)
    }
}

impl DocField {
    pub(crate) fn bm25_score() -> Self {
        Self {
            name: BM25_SCORE_FIELD_NAME.to_string(),
            field: Field::from_field_id(0),
            value_type: Type::F64,
        }
    }

    pub(crate) fn is_bm25_score(&self) -> bool {
        self.name == BM25_SCORE_FIELD_NAME
    }
}

pub trait TantivySchemaExt {
    fn handle_doc_field(&self) -> Result<DocField>;
    fn version_doc_field(&self) -> Result<DocField>;
    fn doc_field(&self, field_name: &str) -> Result<DocField>;
}

impl TantivySchemaExt for Schema {
    fn handle_doc_field(&self) -> Result<DocField> {
        self.doc_field(HANDLE_STORE_NAME)
    }

    fn version_doc_field(&self) -> Result<DocField> {
        self.doc_field(VERSION_STORE_NAME)
    }

    fn doc_field(&self, name: &str) -> Result<DocField> {
        if name == BM25_SCORE_FIELD_NAME {
            return Ok(DocField::bm25_score());
        }
        let field = self
            .get_field(name)
            .with_context(|| format!("Field '{}' not found, schema={:?}", name, self))?;
        let value_type = self.get_field_entry(field).field_type().value_type();
        Ok(DocField {
            name: name.to_string(),
            field,
            value_type,
        })
    }
}

#[derive(Debug)]
pub enum ResultType {
    Rows(Vec<DocField>),
    Count, /* Only support count(*) or count(pk), no nullable cloumns, so we don't need a
            * DocField for count. */
}

#[derive(Debug)]
pub struct TopKRequest {
    pub limit: usize,
    pub sort_field_names: Vec<String>,
    pub is_asc: Vec<bool>,
}

#[derive(Debug, Clone)]
pub struct TopKParams {
    pub limit: usize,
    pub sort_fields: Vec<SortField>,
    pub is_store_sort: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SearchExecutionPlan {
    StoreSortTopKRows,
    FastFieldTopKRows,
    FragScanRows,
    FragScanCount,
}

impl SearchExecutionPlan {
    pub fn as_str(self) -> &'static str {
        match self {
            SearchExecutionPlan::StoreSortTopKRows => "store_sort_topk",
            SearchExecutionPlan::FastFieldTopKRows => "fast_field_topk",
            SearchExecutionPlan::FragScanRows => "frag_scan_rows",
            SearchExecutionPlan::FragScanCount => "frag_scan_count",
        }
    }

    fn needs_expr_query_with_read_ranges(self) -> bool {
        matches!(
            self,
            SearchExecutionPlan::FastFieldTopKRows | SearchExecutionPlan::FragScanCount
        )
    }
}

#[derive(Debug)]
pub struct QueryInfo {
    expr_query: Box<dyn Query>,
    expr_query_with_read_ranges: Option<Box<dyn Query>>,
    read_ranges: Vec<ShardingKeyRange>,
    handle_field: DocField,
    result_type: ResultType,
    topk_params: Option<TopKParams>,
    execution_plan: SearchExecutionPlan,
}

impl QueryInfo {
    fn is_supported_fast_field_topk_type(sort_field: &SortField) -> bool {
        matches!(
            sort_field.field.value_type,
            Type::Str | Type::I64 | Type::U64 | Type::F64
        )
    }

    fn fast_field_topk_error(sort_field: &SortField) -> anyhow::Error {
        anyhow!(
            "sort field '{}' has unsupported type {:?} for fast field topk",
            sort_field.field.name,
            sort_field.field.value_type
        )
    }

    pub fn parse_expr(
        index_meta: &IndexMeta,
        schema: &Schema,
        expr: &tipb::Expr,
        read_ranges: &[ShardingKeyRange],
        return_field_names: &[String],
        topk_request: TopKRequest,
    ) -> Result<Self> {
        let expr_query = Self::expr_to_query(schema, expr, index_meta)?;

        // Build TopKParams from provided parameters if limit > 0
        let topk_params = if topk_request.limit == 0 {
            None
        } else {
            let mut sort_fields = Vec::with_capacity(topk_request.sort_field_names.len());
            for (i, name) in topk_request.sort_field_names.iter().enumerate() {
                let df = schema
                    .doc_field(name)
                    .with_context(|| format!("Sort field '{}' not found in schema", name))?;
                let asc = topk_request.is_asc[i];
                sort_fields.push(SortField {
                    field: df,
                    is_asc: asc,
                });
            }

            // determine if this matches stored sort fields -> use store sort optimization
            let is_store_sort = Self::sort_fields_match(index_meta.get_sort_fields(), &sort_fields);

            // directly construct TopKParams
            Some(TopKParams {
                limit: topk_request.limit,
                sort_fields,
                is_store_sort,
            })
        };

        let result_type = Self::build_result_type(schema, return_field_names)?;
        if matches!(result_type, ResultType::Count) && topk_params.is_some() {
            bail!("TopK is not supported for count result type");
        }
        let execution_plan =
            Self::build_execution_plan(&result_type, topk_params.as_ref(), schema)?;

        let expr_query_with_read_ranges = Self::build_expr_query_with_read_ranges(
            execution_plan,
            &expr_query,
            schema,
            read_ranges,
        )?;

        Ok(QueryInfo {
            expr_query,
            expr_query_with_read_ranges,
            read_ranges: read_ranges.to_vec(),
            handle_field: schema.handle_doc_field()?,
            result_type,
            topk_params,
            execution_plan,
        })
    }

    // For backward compatibility, return true when `query_sort_fields` is empty.
    fn sort_fields_match(
        stored_sort_fields: &[SortField],
        query_sort_fields: &[SortField],
    ) -> bool {
        if query_sort_fields.len() > stored_sort_fields.len() {
            return false;
        }

        for (stored, query) in stored_sort_fields.iter().zip(query_sort_fields.iter()) {
            if stored.field.name != query.field.name {
                return false;
            }

            if stored.is_asc != query.is_asc {
                return false;
            }
        }

        true
    }

    fn get_cmp_func(cmp: &str) -> fn(a: Term) -> Result<Box<dyn Query>> {
        match cmp {
            ">=" => |a| {
                Ok(Box::new(RangeQuery::new(
                    Bound::Included(a),
                    Bound::Unbounded,
                )))
            },
            ">" => |a| {
                Ok(Box::new(RangeQuery::new(
                    Bound::Excluded(a),
                    Bound::Unbounded,
                )))
            },
            "<=" => |a| {
                Ok(Box::new(RangeQuery::new(
                    Bound::Unbounded,
                    Bound::Included(a),
                )))
            },
            "<" => |a| {
                Ok(Box::new(RangeQuery::new(
                    Bound::Unbounded,
                    Bound::Excluded(a),
                )))
            },
            "=" => |a| Ok(Box::new(TermQuery::new(a, IndexRecordOption::Basic))),
            "!=" => |a| {
                let query = TermQuery::new(a, IndexRecordOption::Basic);
                Ok(Box::new(BooleanQuery::new(vec![
                    (Occur::Must, Box::new(AllQuery)),
                    (Occur::MustNot, Box::new(query)),
                ])))
            },
            _ => |_| Err(anyhow!("Unsupported cmp func")),
        }
    }

    fn get_type_func(
        tp: &str,
    ) -> fn(
        column_ref_expr: &tipb::Expr,
        val_expr: &tipb::Expr,
        schema: &Schema,
        index_mete: &IndexMeta,
    ) -> Result<Term> {
        match tp {
            "Int" => |column_ref_expr, val_expr, schema, _| {
                let field_name = str::from_utf8(column_ref_expr.get_val())?;
                let field = schema
                    .get_field(field_name)
                    .with_context(|| format!("Field '{}' not found in schema", field_name))?;

                let value = val_expr.get_val().read_i64()?;
                Ok(Term::from_field_i64(field, value))
            },
            "Real" => |column_ref_expr, val_expr, schema, _| {
                let field_name = str::from_utf8(column_ref_expr.get_val())?;
                let field = schema
                    .get_field(field_name)
                    .with_context(|| format!("Field '{}' not found in schema", field_name))?;

                let value = val_expr.get_val().read_f64()?;
                Ok(Term::from_field_f64(field, value))
            },
            "Decimal" => |column_ref_expr, val_expr, schema, index_meta| {
                let field_name = str::from_utf8(column_ref_expr.get_val())?;
                let field = schema
                    .get_field(field_name)
                    .with_context(|| format!("Field '{}' not found in schema", field_name))?;

                let v = str::from_utf8(val_expr.get_val())?;
                let column_info = index_meta
                    .get_column_infos()
                    .iter()
                    .find(|col| format!("column_{}", col.column_id) == field_name)
                    .with_context(|| format!("Column '{}' not found in index meta", field_name))?;
                let mut buf = vec![];
                encode_decimal_val(
                    &mut buf,
                    v,
                    column_info.column_length as u8,
                    column_info.decimal as u8,
                )?;
                Ok(Term::from_field_bytes(field, &buf))
            },
            "String" => |column_ref_expr, val_expr, schema, _| {
                let field_name = str::from_utf8(column_ref_expr.get_val())?;
                let field = schema
                    .get_field(field_name)
                    .with_context(|| format!("Field '{}' not found in schema", field_name))?;
                let field_type = schema.get_field_entry(field).field_type().value_type();
                match field_type {
                    Type::Str => {
                        let s = std::str::from_utf8(val_expr.get_val())?;
                        Ok(Term::from_field_text(field, s))
                    }
                    Type::Bytes => Ok(Term::from_field_bytes(field, val_expr.get_val())),
                    _ => bail!(
                        "String cmp requires Str/Bytes field, field='{}', field_type={:?}",
                        field_name,
                        field_type
                    ),
                }
            },
            "Time" => |column_ref_expr, val_expr, schema, _| {
                let field_name = str::from_utf8(column_ref_expr.get_val())?;
                let field = schema
                    .get_field(field_name)
                    .with_context(|| format!("Field '{}' not found in schema", field_name))?;

                let packed_u64 = val_expr.get_val().read_u64()?;
                Ok(Term::from_field_u64(field, packed_u64))
            },
            _ => |_, _, _, _| Err(anyhow!("Unsupported cmp type")),
        }
    }

    fn sig_to_type(sig: tipb::ScalarFuncSig) -> Result<&'static str> {
        match sig {
            tipb::ScalarFuncSig::LtInt
            | tipb::ScalarFuncSig::LeInt
            | tipb::ScalarFuncSig::GtInt
            | tipb::ScalarFuncSig::GeInt
            | tipb::ScalarFuncSig::EqInt
            | tipb::ScalarFuncSig::NeInt
            | tipb::ScalarFuncSig::InInt => Ok("Int"),
            tipb::ScalarFuncSig::LtReal
            | tipb::ScalarFuncSig::LeReal
            | tipb::ScalarFuncSig::GtReal
            | tipb::ScalarFuncSig::GeReal
            | tipb::ScalarFuncSig::EqReal
            | tipb::ScalarFuncSig::NeReal
            | tipb::ScalarFuncSig::InReal => Ok("Real"),
            tipb::ScalarFuncSig::LtDecimal
            | tipb::ScalarFuncSig::LeDecimal
            | tipb::ScalarFuncSig::GtDecimal
            | tipb::ScalarFuncSig::GeDecimal
            | tipb::ScalarFuncSig::EqDecimal
            | tipb::ScalarFuncSig::NeDecimal
            | tipb::ScalarFuncSig::InDecimal => Ok("Decimal"),
            tipb::ScalarFuncSig::LtString
            | tipb::ScalarFuncSig::LeString
            | tipb::ScalarFuncSig::GtString
            | tipb::ScalarFuncSig::GeString
            | tipb::ScalarFuncSig::EqString
            | tipb::ScalarFuncSig::NeString
            | tipb::ScalarFuncSig::InString => Ok("String"),
            tipb::ScalarFuncSig::LtTime
            | tipb::ScalarFuncSig::LeTime
            | tipb::ScalarFuncSig::GtTime
            | tipb::ScalarFuncSig::GeTime
            | tipb::ScalarFuncSig::EqTime
            | tipb::ScalarFuncSig::NeTime
            | tipb::ScalarFuncSig::InTime => Ok("Time"),
            _ => Err(anyhow!("Unsupported ScalarFuncSig {:?}", sig)),
        }
    }

    fn sig_to_cmp(sig: tipb::ScalarFuncSig) -> Result<&'static str> {
        match sig {
            tipb::ScalarFuncSig::LtInt
            | tipb::ScalarFuncSig::LtReal
            | tipb::ScalarFuncSig::LtDecimal
            | tipb::ScalarFuncSig::LtString
            | tipb::ScalarFuncSig::LtTime => Ok("<"),
            tipb::ScalarFuncSig::LeInt
            | tipb::ScalarFuncSig::LeReal
            | tipb::ScalarFuncSig::LeDecimal
            | tipb::ScalarFuncSig::LeString
            | tipb::ScalarFuncSig::LeTime => Ok("<="),
            tipb::ScalarFuncSig::GtInt
            | tipb::ScalarFuncSig::GtReal
            | tipb::ScalarFuncSig::GtDecimal
            | tipb::ScalarFuncSig::GtString
            | tipb::ScalarFuncSig::GtTime => Ok(">"),
            tipb::ScalarFuncSig::GeInt
            | tipb::ScalarFuncSig::GeReal
            | tipb::ScalarFuncSig::GeDecimal
            | tipb::ScalarFuncSig::GeString
            | tipb::ScalarFuncSig::GeTime => Ok(">="),
            tipb::ScalarFuncSig::EqInt
            | tipb::ScalarFuncSig::EqReal
            | tipb::ScalarFuncSig::EqDecimal
            | tipb::ScalarFuncSig::EqString
            | tipb::ScalarFuncSig::EqTime => Ok("="),
            tipb::ScalarFuncSig::NeInt
            | tipb::ScalarFuncSig::NeReal
            | tipb::ScalarFuncSig::NeDecimal
            | tipb::ScalarFuncSig::NeString
            | tipb::ScalarFuncSig::NeTime => Ok("!="),
            _ => Err(anyhow!("Unsupported ScalarFuncSig {:?}", sig)),
        }
    }

    fn reverse_cmp(cmp: &str) -> Result<&'static str> {
        match cmp {
            ">=" => Ok("<="),
            ">" => Ok("<"),
            "<=" => Ok(">="),
            "<" => Ok(">"),
            "=" => Ok("="),
            "!=" => Ok("!="),
            _ => Err(anyhow!("Unsupported cmp func")),
        }
    }

    fn get_fulltext_field_and_tokens(
        schema: &Schema,
        index_meta: &IndexMeta,
        column_names: &[&str],
        query_text: &str,
    ) -> Result<(Field, Vec<AnalyzedToken>)> {
        let mut col_ids = Vec::with_capacity(column_names.len());
        for column_name in column_names {
            col_ids.push(parse_column_id(column_name)?);
        }
        let ft_field = index_meta.get_fulltext_field(col_ids)?;
        let analyzed_tokens =
            index_meta.get_analyzed_tokens_with_positions(ft_field.analyzer_name(), query_text)?;
        let ft_field_name = ft_field.field_name();
        let field = schema
            .get_field(&ft_field_name)
            .with_context(|| format!("Field '{}' not found in schema", ft_field_name))?;
        Ok((field, analyzed_tokens))
    }

    fn get_fulltext_field_and_phrase_analysis(
        schema: &Schema,
        index_meta: &IndexMeta,
        column_names: &[&str],
        query_text: &str,
    ) -> Result<(Field, PhraseAnalysis)> {
        let mut col_ids = Vec::with_capacity(column_names.len());
        for column_name in column_names {
            col_ids.push(parse_column_id(column_name)?);
        }
        let ft_field = index_meta.get_fulltext_field(col_ids)?;
        let phrase_analysis =
            index_meta.analyze_phrase_tokens(ft_field.analyzer_name(), query_text)?;
        let ft_field_name = ft_field.field_name();
        let field = schema
            .get_field(&ft_field_name)
            .with_context(|| format!("Field '{}' not found in schema", ft_field_name))?;
        Ok((field, phrase_analysis))
    }

    fn get_fulltext_field_and_prefix_token(
        schema: &Schema,
        index_meta: &IndexMeta,
        column_names: &[&str],
        query_text: &str,
    ) -> Result<(Field, Option<String>)> {
        let mut col_ids = Vec::with_capacity(column_names.len());
        for column_name in column_names {
            col_ids.push(parse_column_id(column_name)?);
        }
        let ft_field = index_meta.get_fulltext_field(col_ids)?;
        let prefix_token = index_meta.analyze_prefix_token(ft_field.analyzer_name(), query_text)?;
        let ft_field_name = ft_field.field_name();
        let field = schema
            .get_field(&ft_field_name)
            .with_context(|| format!("Field '{}' not found in schema", ft_field_name))?;
        Ok((field, prefix_token))
    }

    fn parse_fts_match_args(expr: &tipb::Expr) -> Result<(&str, Vec<&str>)> {
        let children = expr.get_children();
        let sig_name = match expr.get_sig() {
            tipb::ScalarFuncSig::FtsMatchWord => "FtsMatchWord",
            tipb::ScalarFuncSig::FtsMatchPrefix => "FtsMatchPrefix",
            tipb::ScalarFuncSig::FtsMatchPhrase => "FtsMatchPhrase",
            sig => bail!("Unsupported FTS ScalarFuncSig {:?}", sig),
        };
        if children.len() < 2 {
            bail!(
                "{} requires at least 2 children (query text + 1 or more column names), but got {}",
                sig_name,
                children.len()
            );
        }

        let query_text = str::from_utf8(children[0].get_val())?;
        let column_names = children
            .iter()
            .skip(1)
            .map(|child| str::from_utf8(child.get_val()))
            .collect::<std::result::Result<Vec<_>, _>>()?;
        Ok((query_text, column_names))
    }

    fn expr_to_query(
        schema: &Schema,
        expr: &tipb::Expr,
        index_meta: &IndexMeta,
    ) -> Result<Box<dyn Query>> {
        match expr.get_tp() {
            tipb::ExprType::ScalarFunc => match expr.get_sig() {
                tipb::ScalarFuncSig::UnaryNotInt | tipb::ScalarFuncSig::UnaryNotReal => {
                    if expr.get_children().len() != 1 {
                        bail!(
                            "UnaryNot requires must 1 children, but got {}",
                            expr.get_children().len()
                        );
                    }
                    let child =
                        QueryInfo::expr_to_query(schema, &expr.get_children()[0], index_meta)?;

                    Ok(Box::new(BooleanQuery::new(vec![
                        (Occur::Must, Box::new(AllQuery)),
                        (Occur::MustNot, child),
                    ])))
                }
                tipb::ScalarFuncSig::LogicalAnd => {
                    if expr.get_children().len() != 2 {
                        bail!(
                            "LogicalAnd requires 2 children, but got {}",
                            expr.get_children().len()
                        );
                    }
                    let left =
                        QueryInfo::expr_to_query(schema, &expr.get_children()[0], index_meta)?;
                    let right =
                        QueryInfo::expr_to_query(schema, &expr.get_children()[1], index_meta)?;
                    Ok(Box::new(BooleanQuery::intersection(vec![left, right])))
                }
                tipb::ScalarFuncSig::LogicalOr => {
                    if expr.get_children().len() != 2 {
                        bail!(
                            "LogicalOr requires 2 children, but got {}",
                            expr.get_children().len()
                        );
                    }
                    let left =
                        QueryInfo::expr_to_query(schema, &expr.get_children()[0], index_meta)?;
                    let right =
                        QueryInfo::expr_to_query(schema, &expr.get_children()[1], index_meta)?;
                    Ok(Box::new(BooleanQuery::union(vec![left, right])))
                }
                tipb::ScalarFuncSig::FtsMatchWord => {
                    let (query_text, column_names) = Self::parse_fts_match_args(expr)?;
                    let (field, analyzed_tokens) = Self::get_fulltext_field_and_tokens(
                        schema,
                        index_meta,
                        &column_names,
                        query_text,
                    )?;
                    // `fts_match_word` implies an exact match against a SINGLE token.
                    // We return an EmptyQuery in two cases:
                    // 1. The token is filtered out by the analyzer (e.g., stop words like "is",
                    //    "the").
                    // 2. The text is tokenized into multiple terms (e.g., "user_id" -> ["user",
                    //    "id"]).
                    //
                    // Example:
                    // If Input = "user_id", and Analyzer output = ["user", "id"],
                    // then analyzed_tokens.len() is 2, which fails the "single word" constraint.
                    if analyzed_tokens.len() != 1 {
                        return Ok(Box::new(EmptyQuery));
                    }
                    let query_text = &analyzed_tokens[0].text;
                    Ok(Box::new(TermQuery::new(
                        Term::from_field_text(field, query_text),
                        IndexRecordOption::Basic,
                    )))
                }
                tipb::ScalarFuncSig::FtsMatchPrefix => {
                    let (query_text, column_names) = Self::parse_fts_match_args(expr)?;
                    let (field, Some(query_text)) = Self::get_fulltext_field_and_prefix_token(
                        schema,
                        index_meta,
                        &column_names,
                        query_text,
                    )?
                    else {
                        return Ok(Box::new(EmptyQuery));
                    };
                    // `fts_match_prefix` is similar to `fts_match_word`, except that it matches
                    // prefixes instead of exact matches.
                    Ok(Box::new(FuzzyTermQuery::new_prefix(
                        Term::from_field_text(field, &query_text),
                        0,
                        false,
                    )))
                }
                tipb::ScalarFuncSig::FtsMatchPhrase => {
                    let (query_text, column_names) = Self::parse_fts_match_args(expr)?;
                    let (field, phrase_analysis) = Self::get_fulltext_field_and_phrase_analysis(
                        schema,
                        index_meta,
                        &column_names,
                        query_text,
                    )?;
                    let analyzed_tokens = match phrase_analysis {
                        PhraseAnalysis::NoMatch => return Ok(Box::new(EmptyQuery)),
                        PhraseAnalysis::Tokens(tokens) => tokens,
                    };
                    if analyzed_tokens.len() == 1 {
                        return Ok(Box::new(TermQuery::new(
                            Term::from_field_text(field, &analyzed_tokens[0].text),
                            IndexRecordOption::Basic,
                        )));
                    }
                    let base_position = analyzed_tokens
                        .iter()
                        .map(|token| token.position)
                        .min()
                        .unwrap_or(0);
                    let terms = analyzed_tokens
                        .iter()
                        .map(|token| {
                            (
                                token.position - base_position,
                                Term::from_field_text(field, &token.text),
                            )
                        })
                        .collect::<Vec<_>>();
                    Ok(Box::new(PhraseQuery::new_with_offset(terms)))
                }

                tipb::ScalarFuncSig::LtInt
                | tipb::ScalarFuncSig::LeInt
                | tipb::ScalarFuncSig::GtInt
                | tipb::ScalarFuncSig::GeInt
                | tipb::ScalarFuncSig::EqInt
                | tipb::ScalarFuncSig::NeInt
                | tipb::ScalarFuncSig::LtReal
                | tipb::ScalarFuncSig::LeReal
                | tipb::ScalarFuncSig::GtReal
                | tipb::ScalarFuncSig::GeReal
                | tipb::ScalarFuncSig::EqReal
                | tipb::ScalarFuncSig::NeReal
                | tipb::ScalarFuncSig::LtDecimal
                | tipb::ScalarFuncSig::LeDecimal
                | tipb::ScalarFuncSig::GtDecimal
                | tipb::ScalarFuncSig::GeDecimal
                | tipb::ScalarFuncSig::EqDecimal
                | tipb::ScalarFuncSig::NeDecimal
                | tipb::ScalarFuncSig::LtString
                | tipb::ScalarFuncSig::LeString
                | tipb::ScalarFuncSig::GtString
                | tipb::ScalarFuncSig::GeString
                | tipb::ScalarFuncSig::EqString
                | tipb::ScalarFuncSig::NeString
                | tipb::ScalarFuncSig::LtTime
                | tipb::ScalarFuncSig::LeTime
                | tipb::ScalarFuncSig::GtTime
                | tipb::ScalarFuncSig::GeTime
                | tipb::ScalarFuncSig::EqTime
                | tipb::ScalarFuncSig::NeTime => {
                    let typ = Self::sig_to_type(expr.get_sig())?;
                    let mut cmp = Self::sig_to_cmp(expr.get_sig())?;
                    let (mut col_idx, mut val_idx) = (0, 1);
                    if expr.get_children()[0].get_tp() != tipb::ExprType::ColumnRef {
                        // Reverse the cmp if the column_ref is on the right side.
                        cmp = Self::reverse_cmp(cmp)?;
                        (col_idx, val_idx) = (1, 0);
                    }
                    let col_ref_expr = &expr.get_children()[col_idx];
                    let val_expr = &expr.get_children()[val_idx];
                    if col_ref_expr.get_tp() == tipb::ExprType::ColumnRef
                        && val_expr.get_tp() == tipb::ExprType::ColumnRef
                    {
                        bail!("ColumnRef cmp ColumnRef is not supported: {:?}", expr);
                    }
                    let term =
                        Self::get_type_func(typ)(col_ref_expr, val_expr, schema, index_meta)?;
                    Self::get_cmp_func(cmp)(term)
                }

                tipb::ScalarFuncSig::InInt
                | tipb::ScalarFuncSig::InReal
                | tipb::ScalarFuncSig::InDecimal
                | tipb::ScalarFuncSig::InString
                | tipb::ScalarFuncSig::InTime => {
                    let typ = Self::sig_to_type(expr.get_sig())?;
                    let col_ref_expr = &expr.get_children()[0];
                    let typ_func = Self::get_type_func(typ);
                    let mut terms: Vec<Term> = Vec::with_capacity(expr.get_children().len() - 1);
                    for val_expr in expr.get_children().iter().skip(1) {
                        let term = typ_func(col_ref_expr, val_expr, schema, index_meta)?;
                        terms.push(term);
                    }
                    Ok(Box::new(TermSetQuery::new(terms)))
                }
                _ => Err(anyhow!("Unsupported ScalarFuncSig {:?}", expr.get_sig())),
            },
            _ => Err(anyhow!("Unsupported ExprType {:?}", expr.get_tp())),
        }
    }

    fn read_ranges_to_query(
        read_ranges: &[ShardingKeyRange],
        handle_field: &DocField,
    ) -> Box<dyn Query> {
        if read_ranges.is_empty() {
            return Box::new(EmptyQuery);
        }
        let mut range_querys: Vec<Box<dyn Query>> = Vec::with_capacity(read_ranges.len());
        for read_range in read_ranges {
            let start = Term::from_field_bytes(handle_field.field, read_range.start().raw());
            let end = Term::from_field_bytes(handle_field.field, read_range.end().raw());
            let range_query = RangeQuery::new(Bound::Included(start), Bound::Excluded(end));
            range_querys.push(Box::new(range_query));
        }
        Box::new(BooleanQuery::union(range_querys))
    }

    pub(crate) fn intersect_query_and_read_ranges(
        query: Box<dyn Query>,
        schema: &Schema,
        read_ranges: &[ShardingKeyRange],
    ) -> Result<Box<dyn Query>> {
        let handle_field = schema.handle_doc_field()?;
        let read_ranges_query = Self::read_ranges_to_query(read_ranges, &handle_field);
        Ok(Box::new(BooleanQuery::intersection(vec![
            query,
            read_ranges_query,
        ])))
    }

    fn build_result_type(schema: &Schema, return_field_names: &[String]) -> Result<ResultType> {
        if return_field_names.is_empty() {
            Ok(ResultType::Count)
        } else {
            let mut return_fields = Vec::new();
            for name in return_field_names {
                return_fields.push(schema.doc_field(name)?);
            }
            Ok(ResultType::Rows(return_fields))
        }
    }

    fn build_execution_plan(
        result_type: &ResultType,
        topk_params: Option<&TopKParams>,
        schema: &Schema,
    ) -> Result<SearchExecutionPlan> {
        match result_type {
            ResultType::Count => Ok(SearchExecutionPlan::FragScanCount),
            ResultType::Rows(_) => {
                let Some(topk_params) = topk_params else {
                    return Ok(SearchExecutionPlan::FragScanRows);
                };
                if topk_params.is_store_sort {
                    return Ok(SearchExecutionPlan::StoreSortTopKRows);
                }
                if topk_params.sort_fields.len() == 1 {
                    let sort_field = &topk_params.sort_fields[0];
                    if sort_field.field.is_bm25_score() {
                        return Ok(SearchExecutionPlan::FragScanRows);
                    }
                    if schema.get_field_entry(sort_field.field.field).is_fast() {
                        if Self::is_supported_fast_field_topk_type(sort_field) {
                            return Ok(SearchExecutionPlan::FastFieldTopKRows);
                        }
                        return Err(Self::fast_field_topk_error(sort_field));
                    }
                }
                Ok(SearchExecutionPlan::FragScanRows)
            }
        }
    }

    fn build_expr_query_with_read_ranges(
        execution_plan: SearchExecutionPlan,
        expr_query: &dyn Query,
        schema: &Schema,
        read_ranges: &[ShardingKeyRange],
    ) -> Result<Option<Box<dyn Query>>> {
        if execution_plan.needs_expr_query_with_read_ranges() {
            Ok(Some(Self::intersect_query_and_read_ranges(
                expr_query.box_clone(),
                schema,
                read_ranges,
            )?))
        } else {
            Ok(None)
        }
    }

    pub fn ensure_is_count_query(&self) -> Result<()> {
        match &self.result_type {
            ResultType::Count => Ok(()),
            ResultType::Rows(_) => Err(anyhow!(
                "ResultType is invalid, expected Count, got {:?}",
                self.result_type
            )),
        }
    }

    pub fn ensure_is_rows_query(&self) -> Result<()> {
        match &self.result_type {
            ResultType::Count => Err(anyhow!(
                "ResultType is invalid, expected Rows, got {:?}",
                self.result_type
            )),
            ResultType::Rows(_) => Ok(()),
        }
    }

    pub fn expr_query(&self) -> &dyn Query {
        &self.expr_query
    }

    pub fn expr_query_with_read_ranges(&self) -> Result<&dyn Query> {
        self.expr_query_with_read_ranges
            .as_ref()
            .map(|b| b.as_ref())
            .ok_or_else(|| anyhow!("expr_query_with_read_ranges is None"))
    }

    pub fn in_read_ranges(&self, handle: &[u8]) -> bool {
        self.read_ranges.iter().any(|range| range.contains(handle))
    }

    pub fn read_ranges(&self) -> &[ShardingKeyRange] {
        &self.read_ranges
    }

    pub fn handle_field(&self) -> &DocField {
        &self.handle_field
    }

    pub fn new_search_rows(&self) -> Result<SearchRows> {
        match &self.result_type {
            ResultType::Count => Err(anyhow!("Count result type has no search rows")),
            ResultType::Rows(fields) => SearchRows::empty(fields),
        }
    }

    pub fn return_fields(&self) -> Result<&[DocField]> {
        match &self.result_type {
            ResultType::Count => Err(anyhow!("Count result type has no return fields")),
            ResultType::Rows(fields) => Ok(fields),
        }
    }

    pub fn topk_params(&self) -> Option<&TopKParams> {
        self.topk_params.as_ref()
    }

    pub fn execution_plan(&self) -> SearchExecutionPlan {
        self.execution_plan
    }

    pub fn query_all_for_test(
        index_meta: &IndexMeta,
        schema: &Schema,
        read_ranges: Vec<ShardingKeyRange>,
    ) -> Result<Self> {
        let return_fields = index_meta
            .get_column_infos()
            .iter()
            .map(|col| schema.doc_field(&col.store_name()))
            .collect::<Result<Vec<DocField>>>()?;
        let result_type = ResultType::Rows(return_fields);
        let execution_plan = Self::build_execution_plan(&result_type, None, schema)?;
        Ok(Self {
            result_type,
            expr_query: Box::new(tantivy::query::AllQuery),
            expr_query_with_read_ranges: None,
            read_ranges,
            handle_field: schema.handle_doc_field()?,
            topk_params: None,
            execution_plan,
        })
    }

    #[cfg(any(test, feature = "bench"))]
    pub fn query_ranges_for_test(
        index_meta: &IndexMeta,
        schema: &Schema,
        read_ranges: Vec<ShardingKeyRange>,
    ) -> Result<Self> {
        let return_fields = index_meta
            .get_column_infos()
            .iter()
            .map(|col| schema.doc_field(&col.store_name()))
            .collect::<Result<Vec<DocField>>>()?;

        let expr_query = Box::new(tantivy::query::AllQuery);
        let result_type = ResultType::Rows(return_fields);
        let execution_plan = Self::build_execution_plan(&result_type, None, schema)?;
        let expr_query_with_read_ranges = Self::build_expr_query_with_read_ranges(
            execution_plan,
            &*expr_query,
            schema,
            &read_ranges,
        )?;

        Ok(Self {
            result_type,
            expr_query,
            expr_query_with_read_ranges,
            read_ranges,
            handle_field: schema.handle_doc_field()?,
            topk_params: None,
            execution_plan,
        })
    }

    pub fn query_cols_for_test(
        index_meta: &IndexMeta,
        schema: &Schema,
        read_ranges: Vec<ShardingKeyRange>,
        return_col_names: Vec<&str>,
    ) -> Result<QueryInfo> {
        let return_fields = return_col_names
            .iter()
            .map(|col_name| -> Result<DocField> {
                if *col_name == VERSION_STORE_NAME {
                    return schema.doc_field(VERSION_STORE_NAME);
                }
                if *col_name == BM25_SCORE_FIELD_NAME {
                    return schema.doc_field(BM25_SCORE_FIELD_NAME);
                }
                let col_info = index_meta
                    .get_column_infos()
                    .iter()
                    .find(|col_info| &col_info.column_name == col_name)
                    .ok_or_else(|| {
                        anyhow::anyhow!("Column {} not found in index meta", col_name)
                    })?;
                schema.doc_field(&col_info.store_name())
            })
            .collect::<Result<Vec<DocField>>>()?;
        let result_type = ResultType::Rows(return_fields);
        let execution_plan = Self::build_execution_plan(&result_type, None, schema)?;
        Ok(Self {
            result_type,
            expr_query: Box::new(tantivy::query::AllQuery),
            expr_query_with_read_ranges: None,
            read_ranges,
            handle_field: schema.handle_doc_field()?,
            topk_params: None,
            execution_plan,
        })
    }

    #[cfg(any(test, feature = "bench"))]
    pub fn query_word(
        index_meta: &IndexMeta,
        schema: &Schema,
        read_ranges: &[ShardingKeyRange],
        field_name: String,
        value: String,
    ) -> Result<Self> {
        let expr = Self::match_word_expr(field_name, value);
        let return_field_names = index_meta
            .pk_columns()
            .iter()
            .map(|pk_col| pk_col.store_name())
            .collect::<Vec<String>>();
        Self::parse_expr(
            index_meta,
            schema,
            &expr,
            read_ranges,
            &return_field_names,
            TopKRequest {
                limit: 0,
                sort_field_names: vec![],
                is_asc: vec![],
            },
        )
    }

    #[cfg(any(test, feature = "bench"))]
    pub fn add_topk(
        &self,
        schema: &Schema,
        limit: usize,
        sort_fields: Vec<SortField>,
        is_store_sort: bool,
    ) -> Result<Self> {
        let result_type = match &self.result_type {
            ResultType::Count => bail!("TopK is not supported for count result type"),
            ResultType::Rows(fields) => {
                let mut new_fields = fields.clone();
                for sf in sort_fields.iter() {
                    new_fields.push(sf.field.clone());
                }
                ResultType::Rows(new_fields)
            }
        };

        let topk_params = TopKParams {
            limit,
            sort_fields,
            is_store_sort,
        };
        let execution_plan = Self::build_execution_plan(&result_type, Some(&topk_params), schema)?;
        let expr_query_with_read_ranges = Self::build_expr_query_with_read_ranges(
            execution_plan,
            &*self.expr_query,
            schema,
            &self.read_ranges,
        )?;

        Ok(QueryInfo {
            expr_query: self.expr_query.box_clone(),
            expr_query_with_read_ranges,
            read_ranges: self.read_ranges.clone(),
            handle_field: self.handle_field.clone(),
            result_type,
            topk_params: Some(topk_params),
            execution_plan,
        })
    }

    #[cfg(any(test, feature = "bench"))]
    fn match_word_expr(field_name: String, value: String) -> tipb::Expr {
        let mut expr: tipb::Expr = tipb::Expr::default();
        expr.set_tp(tipb::ExprType::ScalarFunc);
        expr.set_sig(tipb::ScalarFuncSig::FtsMatchWord);

        let mut c0 = tipb::Expr::default();
        c0.set_tp(tipb::ExprType::String);
        c0.set_sig(tipb::ScalarFuncSig::Unspecified);
        c0.set_val(value.into());
        expr.mut_children().push(c0);

        let mut c1 = tipb::Expr::default();
        c1.set_tp(tipb::ExprType::String);
        c1.set_sig(tipb::ScalarFuncSig::Unspecified);
        c1.set_val(field_name.into());
        expr.mut_children().push(c1);

        expr
    }

    #[cfg(any(test, feature = "bench"))]
    pub fn query_word_count(
        index_meta: &IndexMeta,
        schema: &Schema,
        read_ranges: &[ShardingKeyRange],
        field_name: String,
        value: String,
    ) -> Result<Self> {
        let expr = Self::match_word_expr(field_name, value);
        Self::parse_expr(
            index_meta,
            schema,
            &expr,
            read_ranges,
            &[],
            TopKRequest {
                limit: 0,
                sort_field_names: vec![],
                is_asc: vec![],
            },
        )
    }
}

#[cfg(test)]
mod tests {
    use tantivy::{Index, collector::Count, doc};
    use tici_common::index::IndexKey;
    use tici_grpc::{ColumnInfo, tidb_def::COLTYPE_VARCHAR};

    use super::*;
    use crate::{
        analysis::{
            analyzer::{AnalyzerDef, DEFAULT_ANALYZER_NAME},
            token_filter::TokenFilterType,
            tokenizer::TokenizerType,
        },
        test_utils::{TEST_ALL_TYPES_META, TEST_INDEX_META_NORMAL, test_get_inf_range},
        utils::schema::{FulltextField, IndexMeta, IndexMetaStorage, IndexMetaVersion},
    };

    const TEST_NGRAM_ANALYZER_NAME: &str = "test_ngram";

    #[test]
    fn test_bm25_score_doc_field_is_pseudo_field() -> Result<()> {
        let schema = Schema::builder().build();
        let field = schema.doc_field(BM25_SCORE_FIELD_NAME)?;

        assert!(field.is_bm25_score());
        assert_eq!(field.value_type, Type::F64);
        Ok(())
    }

    fn build_test_index_meta() -> IndexMeta {
        IndexMeta::try_from(IndexMetaStorage {
            version: IndexMetaVersion::LATEST,
            index_key: IndexKey {
                keyspace_id: 1,
                table_id: 1,
                index_id: 1,
            },
            created_ts: 0,
            columns: vec![
                ColumnInfo {
                    column_id: 1,
                    column_name: "title".to_string(),
                    r#type: COLTYPE_VARCHAR,
                    collation: "utf8mb4_bin".to_string(),
                    ..Default::default()
                },
                ColumnInfo {
                    column_id: 2,
                    column_name: "body".to_string(),
                    r#type: COLTYPE_VARCHAR,
                    collation: "utf8mb4_bin".to_string(),
                    ..Default::default()
                },
            ],
            index_column_count: 2,
            sort_columns: vec![],
            inverted_columns: vec![],
            fulltext_fields: vec![FulltextField::new(
                vec![1, 2],
                DEFAULT_ANALYZER_NAME.to_string(),
            )],
            sharding_columns: vec![],
            analyzer_defs: vec![],
            is_partition: false,
        })
        .unwrap()
    }

    fn build_test_ngram_index_meta() -> IndexMeta {
        IndexMeta::try_from(IndexMetaStorage {
            version: IndexMetaVersion::LATEST,
            index_key: IndexKey {
                keyspace_id: 1,
                table_id: 1,
                index_id: 2,
            },
            created_ts: 0,
            columns: vec![ColumnInfo {
                column_id: 1,
                column_name: "title".to_string(),
                r#type: COLTYPE_VARCHAR,
                collation: "utf8mb4_bin".to_string(),
                ..Default::default()
            }],
            index_column_count: 1,
            sort_columns: vec![],
            inverted_columns: vec![],
            fulltext_fields: vec![FulltextField::new(
                vec![1],
                TEST_NGRAM_ANALYZER_NAME.to_string(),
            )],
            sharding_columns: vec![],
            analyzer_defs: vec![AnalyzerDef::new(
                TEST_NGRAM_ANALYZER_NAME.to_string(),
                TokenizerType::PreserveUnderscore,
                vec![
                    TokenFilterType::Ngram { min: 2, max: 2 },
                    TokenFilterType::LowerCaser,
                ],
            )],
            is_partition: false,
        })
        .unwrap()
    }

    fn count_fts_matches_for_columns(
        index_meta: &IndexMeta,
        sig: tipb::ScalarFuncSig,
        query: &str,
        docs: &[&str],
        col_ids: &[i64],
    ) -> usize {
        let schema = index_meta.get_tantivy_schema().clone();
        let expr = build_fts_expr(sig, query, col_ids);
        let query = QueryInfo::expr_to_query(&schema, &expr, index_meta).unwrap();

        let mut index = Index::create_in_ram(schema.clone());
        index.set_tokenizers(index_meta.get_analyzer_manager().clone());
        let mut writer = index.writer(50_000_000).unwrap();
        let field = schema
            .get_field(
                &index_meta
                    .get_fulltext_field(col_ids.to_vec())
                    .unwrap()
                    .field_name(),
            )
            .unwrap();
        for doc_text in docs {
            writer.add_document(doc!(field => *doc_text)).unwrap();
        }
        writer.commit().unwrap();

        let reader = index.reader().unwrap();
        let searcher = reader.searcher();
        searcher.search(query.as_ref(), &Count).unwrap()
    }

    fn count_phrase_matches_for_columns(
        index_meta: &IndexMeta,
        query: &str,
        docs: &[&str],
        col_ids: &[i64],
    ) -> usize {
        count_fts_matches_for_columns(
            index_meta,
            tipb::ScalarFuncSig::FtsMatchPhrase,
            query,
            docs,
            col_ids,
        )
    }

    fn count_phrase_matches(index_meta: &IndexMeta, query: &str, docs: &[&str]) -> usize {
        count_phrase_matches_for_columns(index_meta, query, docs, &[1])
    }

    fn count_prefix_matches(index_meta: &IndexMeta, query: &str, docs: &[&str]) -> usize {
        count_fts_matches_for_columns(
            index_meta,
            tipb::ScalarFuncSig::FtsMatchPrefix,
            query,
            docs,
            &[1, 2],
        )
    }

    fn build_fts_expr(sig: tipb::ScalarFuncSig, query: &str, col_ids: &[i64]) -> tipb::Expr {
        let mut expr = tipb::Expr::default();
        expr.set_tp(tipb::ExprType::ScalarFunc);
        expr.set_sig(sig);

        let mut query_expr = tipb::Expr::default();
        query_expr.set_tp(tipb::ExprType::String);
        query_expr.set_val(query.as_bytes().to_vec());
        expr.mut_children().push(query_expr);

        for col_id in col_ids {
            let mut col_expr = tipb::Expr::default();
            col_expr.set_tp(tipb::ExprType::String);
            col_expr.set_val(format!("column_{}", col_id).into_bytes());
            expr.mut_children().push(col_expr);
        }

        expr
    }

    #[test]
    fn test_expr_to_query_multi_column_word() {
        let index_meta = build_test_index_meta();
        let schema = index_meta.get_tantivy_schema();
        let expr = build_fts_expr(tipb::ScalarFuncSig::FtsMatchWord, "hello", &[1, 2]);

        assert!(QueryInfo::expr_to_query(schema, &expr, &index_meta).is_ok());
    }

    #[test]
    fn test_expr_to_query_multi_column_prefix() {
        let index_meta = build_test_index_meta();
        let schema = index_meta.get_tantivy_schema();
        let expr = build_fts_expr(tipb::ScalarFuncSig::FtsMatchPrefix, "hello", &[1, 2]);

        assert!(QueryInfo::expr_to_query(schema, &expr, &index_meta).is_ok());
    }

    #[test]
    fn test_standard_prefix_keeps_short_query_token() {
        let index_meta = build_test_index_meta();
        let count = count_prefix_matches(&index_meta, "he", &["hello", "HELIUM", "world"]);
        assert_eq!(count, 2);
    }

    #[test]
    fn test_standard_prefix_does_not_match_filtered_short_index_token() {
        let index_meta = build_test_index_meta();
        let count = count_prefix_matches(&index_meta, "he", &["he", "hello"]);
        assert_eq!(count, 1);
    }

    #[test]
    fn test_standard_prefix_still_normalizes_case() {
        let index_meta = build_test_index_meta();
        let count = count_prefix_matches(&index_meta, "HE", &["hello", "helium", "world"]);
        assert_eq!(count, 2);
    }

    #[test]
    fn test_standard_prefix_rejects_multiple_source_tokens() {
        let index_meta = build_test_index_meta();
        let count = count_prefix_matches(&index_meta, "he world", &["hello world", "hello"]);
        assert_eq!(count, 0);
    }

    #[test]
    fn test_standard_prefix_keeps_stopword_prefix() {
        let index_meta = build_test_index_meta();
        let count = count_prefix_matches(&index_meta, "the", &["theater", "theme", "the", "other"]);
        assert_eq!(count, 2);
    }

    #[test]
    fn test_expr_to_query_multi_column_phrase_order_insensitive() {
        let index_meta = build_test_index_meta();
        let schema = index_meta.get_tantivy_schema();
        let expr = build_fts_expr(tipb::ScalarFuncSig::FtsMatchPhrase, "hello world", &[2, 1]);

        assert!(QueryInfo::expr_to_query(schema, &expr, &index_meta).is_ok());
    }

    #[test]
    fn test_expr_to_query_multi_column_requires_columns() {
        let index_meta = build_test_index_meta();
        let schema = index_meta.get_tantivy_schema();
        let expr = build_fts_expr(tipb::ScalarFuncSig::FtsMatchWord, "hello", &[]);

        let err = QueryInfo::expr_to_query(schema, &expr, &index_meta).unwrap_err();
        assert_eq!(
            err.to_string(),
            "FtsMatchWord requires at least 2 children (query text + 1 or more column names), but got 1"
        );
    }

    #[test]
    fn test_ngram_phrase_plain_term_matches_containing_token() {
        let index_meta = build_test_ngram_index_meta();
        let count = count_phrase_matches(&index_meta, "hello", &["hello", "ahello", "yellow"]);
        assert_eq!(count, 2);
    }

    #[test]
    fn test_standard_phrase_with_short_filtered_token_keeps_gap() {
        let index_meta = build_test_index_meta();
        let count =
            count_phrase_matches_for_columns(&index_meta, "abc de fgh", &["abc de fgh"], &[1, 2]);
        assert_eq!(count, 1);
    }

    #[test]
    fn test_standard_phrase_with_stopword_keeps_gap() {
        let index_meta = build_test_index_meta();
        let count = count_phrase_matches_for_columns(
            &index_meta,
            "hello the world",
            &["hello the world"],
            &[1, 2],
        );
        assert_eq!(count, 1);
    }

    #[test]
    fn test_ngram_phrase_exact_phrase_requires_token_boundary() {
        let index_meta = build_test_ngram_index_meta();
        let count = count_phrase_matches(
            &index_meta,
            "hello world",
            &["hello world", "helloworld", "helloXworld"],
        );
        assert_eq!(count, 1);
    }

    #[test]
    fn test_ngram_phrase_with_empty_token_group_returns_no_matches() {
        let index_meta = build_test_ngram_index_meta();
        let count = count_phrase_matches(&index_meta, "lo w", &["hello", "yellow", "hello world"]);
        assert_eq!(count, 0);
    }

    #[test]
    fn test_ngram_phrase_with_complete_cross_token_ngrams_still_matches() {
        let index_meta = build_test_ngram_index_meta();
        let count = count_phrase_matches(&index_meta, "lo wo", &["hello world", "yellow"]);
        assert_eq!(count, 1);
    }

    #[test]
    fn test_ngram_phrase_does_not_cross_filtered_short_token() {
        let index_meta = build_test_ngram_index_meta();
        let count = count_phrase_matches(&index_meta, "lo wo", &["hello world", "hello a world"]);
        assert_eq!(count, 1);
    }

    #[test]
    fn test_execution_plan_defaults_to_frag_scan_rows() {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let query_info =
            QueryInfo::query_all_for_test(&index_meta, &schema, vec![test_get_inf_range()])
                .unwrap();

        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanRows
        );
        assert!(query_info.expr_query_with_read_ranges().is_err());
    }

    #[test]
    fn test_execution_plan_uses_store_sort_topk_rows() {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let base_query =
            QueryInfo::query_all_for_test(&index_meta, &schema, vec![test_get_inf_range()])
                .unwrap();
        let query_info = base_query
            .add_topk(
                &schema,
                10,
                vec![SortField {
                    field: schema.doc_field("column_1").unwrap(),
                    is_asc: true,
                }],
                true,
            )
            .unwrap();

        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::StoreSortTopKRows
        );
        assert!(query_info.expr_query_with_read_ranges().is_err());
    }

    #[test]
    fn test_execution_plan_uses_fast_field_topk_rows() {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let base_query =
            QueryInfo::query_all_for_test(&index_meta, &schema, vec![test_get_inf_range()])
                .unwrap();
        let query_info = base_query
            .add_topk(
                &schema,
                10,
                vec![SortField {
                    field: schema.doc_field(VERSION_STORE_NAME).unwrap(),
                    is_asc: false,
                }],
                false,
            )
            .unwrap();

        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FastFieldTopKRows
        );
        assert!(query_info.expr_query_with_read_ranges().is_ok());
    }

    #[test]
    fn test_execution_plan_falls_back_for_multi_sort_topk() {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let base_query =
            QueryInfo::query_all_for_test(&index_meta, &schema, vec![test_get_inf_range()])
                .unwrap();
        let query_info = base_query
            .add_topk(
                &schema,
                10,
                vec![
                    SortField {
                        field: schema.doc_field("column_3").unwrap(),
                        is_asc: false,
                    },
                    SortField {
                        field: schema.doc_field("column_4").unwrap(),
                        is_asc: false,
                    },
                ],
                false,
            )
            .unwrap();

        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanRows
        );
        assert!(query_info.expr_query_with_read_ranges().is_err());
    }

    #[test]
    fn test_execution_plan_falls_back_for_non_fast_sort_field() {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let base_query =
            QueryInfo::query_all_for_test(&index_meta, &schema, vec![test_get_inf_range()])
                .unwrap();
        let query_info = base_query
            .add_topk(
                &schema,
                10,
                vec![SortField {
                    field: schema.doc_field("column_2").unwrap(),
                    is_asc: true,
                }],
                false,
            )
            .unwrap();

        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanRows
        );
        assert!(query_info.expr_query_with_read_ranges().is_err());
    }

    #[test]
    fn test_execution_plan_returns_error_for_unsupported_fast_sort_field_type() {
        let index_meta = TEST_ALL_TYPES_META.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let base_query =
            QueryInfo::query_all_for_test(&index_meta, &schema, vec![test_get_inf_range()])
                .unwrap();
        let err = base_query
            .add_topk(
                &schema,
                10,
                vec![SortField {
                    field: schema.doc_field("column_10").unwrap(),
                    is_asc: true,
                }],
                false,
            )
            .unwrap_err();
        assert_eq!(
            err.to_string(),
            "sort field 'column_10' has unsupported type Bytes for fast field topk"
        );
    }

    #[test]
    fn test_execution_plan_uses_frag_scan_count() {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let query_info = QueryInfo::query_word_count(
            &index_meta,
            &schema,
            &[test_get_inf_range()],
            "column_1".to_string(),
            "rwp09000".to_string(),
        )
        .unwrap();

        assert_eq!(
            query_info.execution_plan(),
            SearchExecutionPlan::FragScanCount
        );
        assert!(query_info.expr_query_with_read_ranges().is_ok());
    }

    #[test]
    fn test_parse_expr_rejects_topk_for_count_query() {
        let index_meta = TEST_INDEX_META_NORMAL.clone();
        let schema = index_meta.get_tantivy_schema().clone();
        let expr = build_fts_expr(tipb::ScalarFuncSig::FtsMatchWord, "rwp09000", &[1]);
        let err = QueryInfo::parse_expr(
            &index_meta,
            &schema,
            &expr,
            &[test_get_inf_range()],
            &[],
            TopKRequest {
                limit: 10,
                sort_field_names: vec!["column_1".to_string()],
                is_asc: vec![true],
            },
        )
        .unwrap_err();

        assert_eq!(
            err.to_string(),
            "TopK is not supported for count result type"
        );
    }
}