// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use std::collections::{HashMap, HashSet};

use anyhow::{Result, anyhow, bail};
use kvenginepb::fts::{FullTextAnalyzerType, FullTextIndexFieldDef};
use tantivy::{
    Term,
    query::{
        AllQuery, BooleanQuery, ConstScoreQuery, EmptyQuery, FuzzyTermQuery, Occur,
        PhrasePrefixQuery, Query as TantivyQuery, TermQuery,
    },
    schema::{Field, IndexRecordOption},
    tokenizer::{TextAnalyzer, TokenStream},
};

use super::{
    FtsMatchExpr, FtsMatchExprKind, FtsMatchPredicate, FtsMatchType, FtsScoreMode,
    QueryDocumentField,
};
use crate::physical_index_analysis::{ResolvedFieldKind, resolve_field_kind};

#[derive(Clone)]
struct ResolvedMatchField {
    field_layout: crate::SingleFieldLayout,
    index_field: FullTextIndexFieldDef,
}

struct MatchFieldResolver {
    fields: HashMap<(i64, String), ResolvedMatchField>,
}

impl MatchFieldResolver {
    fn new(
        tracked_index_fields: &[FullTextIndexFieldDef],
        index_layout: &crate::IndexLayout,
    ) -> Result<Self> {
        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()
            );
        }

        let mut fields = HashMap::with_capacity(tracked_index_fields.len());
        for (tracked_field, field_layout) in tracked_index_fields.iter().zip(index_layout.fields())
        {
            let key = (
                tracked_field.get_col_id(),
                logical_index_field_name(tracked_field)?.to_string(),
            );
            let resolved = ResolvedMatchField {
                field_layout: field_layout.clone(),
                index_field: tracked_field.clone(),
            };
            if fields.insert(key.clone(), resolved).is_some() {
                bail!(
                    "duplicate tracked FTS field for column {} and index_field {}",
                    key.0,
                    key.1
                );
            }
        }
        Ok(Self { fields })
    }

    fn resolve(&self, field_ref: &super::FtsMatchFieldRef) -> Result<&ResolvedMatchField> {
        self.fields
            .get(&(field_ref.col_id, field_ref.index_field.clone()))
            .ok_or_else(|| {
                anyhow!(
                    "missing tracked FTS field for column {} and index_field {}",
                    field_ref.col_id,
                    field_ref.index_field
                )
            })
    }
}

enum LoweredQuery {
    Empty,
    Query(Box<dyn TantivyQuery + 'static>),
}

impl LoweredQuery {
    fn into_query(self) -> Box<dyn TantivyQuery + 'static> {
        match self {
            LoweredQuery::Empty => Box::new(EmptyQuery),
            LoweredQuery::Query(query) => query,
        }
    }

    fn wrap_const_score(self, score: f32) -> Self {
        match self {
            LoweredQuery::Empty => LoweredQuery::Empty,
            LoweredQuery::Query(query) => {
                LoweredQuery::Query(Box::new(ConstScoreQuery::new(query, score)))
            }
        }
    }
}

/// Final result of lowering one logical match expression into Tantivy.
///
/// Unlike [`LoweringMetadata`], all metadata vectors are deduplicated and
/// sorted by their physical Tantivy keys. This deterministic representation is
/// stored in [`super::Query`] and shared by indexed search, external-document
/// construction, and BM25 statistics collection.
pub(crate) struct LoweredMatchExprQuery {
    /// Executable Tantivy query tree.
    pub(crate) query: Box<dyn TantivyQuery + 'static>,
    /// Source-column bindings sorted by physical Tantivy field.
    pub(crate) document_fields: Vec<QueryDocumentField>,
    /// Scored terms sorted by Tantivy term order for direct binary lookup.
    pub(crate) bm25_terms: Vec<Term>,
    /// Physical fields whose token totals may be requested by Tantivy weights,
    /// sorted and deduplicated.
    pub(crate) bm25_fields: Vec<Field>,
}

/// Unordered metadata accumulated while recursively lowering a match
/// expression.
///
/// Repeated leaves can reference the same physical field or scored term, so
/// hash collections perform deduplication during traversal. Once lowering is
/// complete, [`lower_match_expr_query`] converts these collections into the
/// deterministic vectors exposed by [`LoweredMatchExprQuery`].
struct LoweringMetadata {
    /// Whether to retain terms recorded while lowering this query.
    /// No-score queries never require BM25 statistics.
    collect_bm25_terms: bool,
    /// One logical source-column binding for each physical Tantivy field.
    document_fields: HashMap<Field, QueryDocumentField>,
    /// Terms whose document frequencies may be requested by BM25 scorers.
    bm25_terms: HashSet<Term>,
}

impl LoweringMetadata {
    fn new(collect_bm25_terms: bool) -> Self {
        Self {
            collect_bm25_terms,
            document_fields: HashMap::new(),
            bm25_terms: HashSet::new(),
        }
    }

    fn record_document_field(&mut self, resolved: &ResolvedMatchField) {
        let field = resolved.field_layout.field();
        self.document_fields
            .entry(field)
            .or_insert(QueryDocumentField {
                column_id: resolved.index_field.get_col_id(),
                field,
                field_type: resolved.field_layout.field_type(),
            });
    }

    fn record_bm25_term(&mut self, term: Term) {
        if self.collect_bm25_terms {
            self.bm25_terms.insert(term);
        }
    }
}

pub(crate) fn validate_scored_match_expr(
    info: &tipb::FtsQueryInfo,
    expr: &FtsMatchExpr,
) -> Result<()> {
    if info.get_query_type() == tipb::FtsQueryType::FtsQueryTypeWithScore
        && !contains_positive_match(expr)
    {
        bail!("scored pure-negative FTS match_expr is not supported");
    }
    Ok(())
}

pub(crate) fn lower_match_expr_query(
    info: &tipb::FtsQueryInfo,
    expr: &FtsMatchExpr,
    tracked_index_fields: &[FullTextIndexFieldDef],
    index_layout: &crate::IndexLayout,
) -> Result<LoweredMatchExprQuery> {
    validate_scored_match_expr(info, expr)?;
    let resolver = MatchFieldResolver::new(tracked_index_fields, index_layout)?;
    let scoring_enabled = info.get_query_type() == tipb::FtsQueryType::FtsQueryTypeWithScore;
    let mut metadata = LoweringMetadata::new(scoring_enabled);
    let query =
        lower_expr(expr, info.get_query_type(), &resolver, false, &mut metadata)?.into_query();
    let mut document_fields = metadata.document_fields.into_values().collect::<Vec<_>>();
    document_fields.sort_unstable_by_key(QueryDocumentField::field);

    let mut bm25_terms = metadata.bm25_terms.into_iter().collect::<Vec<_>>();
    bm25_terms.sort_unstable();

    // Tantivy's ConstScoreQuery currently forwards its scoring-enabled context
    // to the child weight. A filter-only child may therefore request
    // total_num_tokens for any field referenced by the query, even though its
    // score is replaced by a constant. Use the complete, already deduplicated
    // and sorted document field set instead of deriving fields from collected
    // BM25 terms.
    //
    // TODO: Change Tantivy's ConstScoreQuery to construct its child weight with
    // scoring disabled, then collect only fields that contribute BM25 scores.
    let bm25_fields = if scoring_enabled {
        document_fields
            .iter()
            .map(QueryDocumentField::field)
            .collect()
    } else {
        Vec::new()
    };

    Ok(LoweredMatchExprQuery {
        query,
        document_fields,
        bm25_terms,
        bm25_fields,
    })
}

fn contains_positive_match(expr: &FtsMatchExpr) -> bool {
    match &expr.kind {
        FtsMatchExprKind::Match(_) => true,
        FtsMatchExprKind::And(children) | FtsMatchExprKind::Or(children) => {
            children.iter().any(contains_positive_match)
        }
        FtsMatchExprKind::Not(_) => false,
    }
}

fn lower_expr(
    expr: &FtsMatchExpr,
    query_type: tipb::FtsQueryType,
    resolver: &MatchFieldResolver,
    inside_const_score: bool,
    metadata: &mut LoweringMetadata,
) -> Result<LoweredQuery> {
    let current_is_filter = query_type == tipb::FtsQueryType::FtsQueryTypeWithScore
        && expr.score_mode == FtsScoreMode::Filter;
    let child_inside_const_score = inside_const_score || current_is_filter;
    let score_contributes = query_type == tipb::FtsQueryType::FtsQueryTypeWithScore
        && !inside_const_score
        && expr.score_mode != FtsScoreMode::Filter;

    let lowered = match &expr.kind {
        FtsMatchExprKind::Match(predicate) => {
            lower_match(predicate, resolver, score_contributes, metadata)?
        }
        FtsMatchExprKind::And(children) => {
            let lowered_children = children
                .iter()
                .map(|child| {
                    lower_expr(
                        child,
                        query_type,
                        resolver,
                        child_inside_const_score,
                        metadata,
                    )
                })
                .collect::<Result<Vec<_>>>()?;
            lower_and(lowered_children)
        }
        FtsMatchExprKind::Or(children) => {
            let lowered_children = children
                .iter()
                .map(|child| {
                    lower_expr(
                        child,
                        query_type,
                        resolver,
                        child_inside_const_score,
                        metadata,
                    )
                })
                .collect::<Result<Vec<_>>>()?;
            lower_or(lowered_children)
        }
        FtsMatchExprKind::Not(child) => lower_not(lower_expr(
            child,
            query_type,
            resolver,
            child_inside_const_score,
            metadata,
        )?),
    };

    if current_is_filter && !inside_const_score {
        Ok(lowered.wrap_const_score(0.0))
    } else {
        Ok(lowered)
    }
}

fn lower_match(
    predicate: &FtsMatchPredicate,
    resolver: &MatchFieldResolver,
    score_contributes: bool,
    metadata: &mut LoweringMetadata,
) -> Result<LoweredQuery> {
    let mut field_queries = Vec::with_capacity(predicate.fields.len());
    for field_ref in &predicate.fields {
        let resolved = resolver.resolve(field_ref)?;
        field_queries.push(lower_field_match(
            predicate,
            resolved,
            score_contributes,
            metadata,
        )?);
    }
    Ok(lower_or(field_queries))
}

fn lower_field_match(
    predicate: &FtsMatchPredicate,
    resolved: &ResolvedMatchField,
    score_contributes: bool,
    metadata: &mut LoweringMetadata,
) -> Result<LoweredQuery> {
    metadata.record_document_field(resolved);
    let analyzer_type = analyzer_type(&resolved.index_field)?;
    // Analyzer and physical type jointly determine how the wire `query` scalar
    // is interpreted. EXACT_VALUE_V1 uses a literal text term on text fields,
    // JSON numbers on native numeric fields, and Base64 on bytes fields. This
    // is independent of enable_bm25, which only controls scoring.
    let field_kind = resolve_field_kind(analyzer_type, resolved.field_layout.field_type())
        .map_err(|err| {
            anyhow!(
                "invalid physical FTS field for column {}: {}",
                resolved.index_field.get_col_id(),
                err,
            )
        })?;
    if score_contributes && !resolved.field_layout.enable_bm25() {
        bail!(
            "FTS field {} for column {} is filter-only and cannot contribute score",
            logical_index_field_name(&resolved.index_field)?,
            resolved.index_field.get_col_id()
        );
    }

    if analyzer_type == FullTextAnalyzerType::FtsAnalyzerExactValueV1 {
        let field_type = match field_kind {
            ResolvedFieldKind::Text => crate::PhysicalFieldType::Text,
            ResolvedFieldKind::NativeExact(field_type) => field_type,
        };
        return lower_exact_field_match(
            predicate,
            resolved,
            field_type,
            score_contributes,
            metadata,
        );
    }
    debug_assert_eq!(field_kind, ResolvedFieldKind::Text);

    let query_text = predicate.query.as_str().ok_or_else(|| {
        anyhow!(
            "text FTS field for column {} requires a string query payload",
            resolved.index_field.get_col_id()
        )
    })?;
    let record_option = if resolved.field_layout.enable_bm25() && score_contributes {
        IndexRecordOption::WithFreqsAndPositions
    } else {
        IndexRecordOption::Basic
    };
    match analyzer_type {
        FullTextAnalyzerType::FtsAnalyzerPathHierarchyV1 => {
            match path_hierarchy_query_term(&resolved.index_field, query_text)? {
                PathHierarchyQueryTerm::MatchAll => Ok(LoweredQuery::Query(Box::new(AllQuery))),
                PathHierarchyQueryTerm::Empty => Ok(LoweredQuery::Empty),
                PathHierarchyQueryTerm::Term(term_text) => {
                    let term = Term::from_field_text(resolved.field_layout.field(), &term_text);
                    metadata.record_bm25_term(term.clone());
                    Ok(LoweredQuery::Query(Box::new(TermQuery::new(
                        term,
                        record_option,
                    ))))
                }
            }
        }
        _ => lower_non_path_field_match(
            predicate,
            query_text,
            resolved.field_layout.field(),
            &resolved.index_field,
            analyzer_type,
            record_option,
            score_contributes,
            metadata,
        ),
    }
}

fn lower_exact_field_match(
    predicate: &FtsMatchPredicate,
    resolved: &ResolvedMatchField,
    field_type: crate::PhysicalFieldType,
    score_contributes: bool,
    metadata: &mut LoweringMetadata,
) -> Result<LoweredQuery> {
    if predicate.fields.len() != 1 {
        bail!(
            "exact field for column {} requires exactly one field",
            resolved.index_field.get_col_id()
        );
    }
    if predicate.match_type != FtsMatchType::Word {
        bail!("exact fields only support exact word matching");
    }
    let field = resolved.field_layout.field();
    if field_type == crate::PhysicalFieldType::Text {
        let text = predicate.query.as_str().ok_or_else(|| {
            anyhow!(
                "exact text field for column {} requires a string query payload",
                resolved.index_field.get_col_id()
            )
        })?;
        let term = Term::from_field_text(field, text);
        if score_contributes {
            metadata.record_bm25_term(term.clone());
        }
        return Ok(LoweredQuery::Query(Box::new(TermQuery::new(
            term,
            IndexRecordOption::Basic,
        ))));
    }

    let value = if field_type == crate::PhysicalFieldType::Bytes {
        let encoded = predicate.query.as_str().ok_or_else(|| {
            anyhow!(
                "exact bytes field for column {} requires a Base64 string query payload",
                resolved.index_field.get_col_id()
            )
        })?;
        let raw = base64::decode(encoded).map_err(|err| {
            anyhow!(
                "invalid exact query base64 for column {}: {}",
                resolved.index_field.get_col_id(),
                err
            )
        })?;
        tantivy::schema::OwnedValue::Bytes(raw)
    } else {
        let number = predicate.query.as_number().ok_or_else(|| {
            anyhow!(
                "exact {:?} field for column {} requires a JSON number query payload",
                field_type,
                resolved.index_field.get_col_id()
            )
        })?;
        crate::exact_value::exact_value_from_query_number(field_type, number)?
    };
    let term = crate::exact_value::exact_term_from_value(field, value)?;
    if score_contributes {
        metadata.record_bm25_term(term.clone());
    }
    Ok(LoweredQuery::Query(Box::new(TermQuery::new(
        term,
        IndexRecordOption::Basic,
    ))))
}

fn lower_and(children: Vec<LoweredQuery>) -> LoweredQuery {
    if children
        .iter()
        .any(|child| matches!(child, LoweredQuery::Empty))
    {
        return LoweredQuery::Empty;
    }
    let mut queries = children
        .into_iter()
        .filter_map(|child| match child {
            LoweredQuery::Empty => None,
            LoweredQuery::Query(query) => Some(query),
        })
        .collect::<Vec<_>>();
    if queries.is_empty() {
        return LoweredQuery::Query(Box::new(AllQuery));
    }
    if queries.len() == 1 {
        return LoweredQuery::Query(queries.pop().unwrap());
    }
    LoweredQuery::Query(Box::new(BooleanQuery::new(
        queries
            .into_iter()
            .map(|query| (Occur::Must, query))
            .collect(),
    )))
}

fn lower_or(children: Vec<LoweredQuery>) -> LoweredQuery {
    let mut queries = children
        .into_iter()
        .filter_map(|child| match child {
            LoweredQuery::Empty => None,
            LoweredQuery::Query(query) => Some(query),
        })
        .collect::<Vec<_>>();
    if queries.is_empty() {
        return LoweredQuery::Empty;
    }
    if queries.len() == 1 {
        return LoweredQuery::Query(queries.pop().unwrap());
    }
    LoweredQuery::Query(Box::new(BooleanQuery::new(
        queries
            .into_iter()
            .map(|query| (Occur::Should, query))
            .collect(),
    )))
}

fn lower_not(child: LoweredQuery) -> LoweredQuery {
    match child {
        LoweredQuery::Empty => LoweredQuery::Query(Box::new(AllQuery)),
        LoweredQuery::Query(query) => LoweredQuery::Query(Box::new(BooleanQuery::new(vec![
            (Occur::Must, Box::new(AllQuery)),
            (Occur::MustNot, query),
        ]))),
    }
}

fn lower_non_path_field_match(
    predicate: &FtsMatchPredicate,
    query_text: &str,
    field: tantivy::schema::Field,
    index_field: &FullTextIndexFieldDef,
    analyzer_type: FullTextAnalyzerType,
    record_option: IndexRecordOption,
    score_contributes: bool,
    metadata: &mut LoweringMetadata,
) -> Result<LoweredQuery> {
    let analyzed_terms = analyzed_query_terms(index_field, query_text)?;
    if analyzed_terms.is_empty() {
        return Ok(LoweredQuery::Empty);
    }

    match predicate.match_type {
        FtsMatchType::Word => {
            let lowered_queries = analyzed_terms
                .iter()
                .map(|term_text| {
                    let term = Term::from_field_text(field, term_text);
                    metadata.record_bm25_term(term.clone());
                    LoweredQuery::Query(Box::new(TermQuery::new(term, record_option)))
                })
                .collect::<Vec<_>>();
            Ok(match analyzer_type {
                FullTextAnalyzerType::FtsAnalyzerNgramV1 => lower_and(lowered_queries),
                _ => lower_or(lowered_queries),
            })
        }
        FtsMatchType::Prefix => {
            if score_contributes && analyzed_terms.len() > 1 {
                bail!(
                    "scored multi-token FTS prefix query is not supported with global BM25 statistics"
                );
            }
            Ok(LoweredQuery::Query(prefix_terms_query(
                field,
                analyzed_terms,
            )))
        }
    }
}

fn prefix_terms_query(
    field: tantivy::schema::Field,
    terms: Vec<String>,
) -> Box<dyn TantivyQuery + 'static> {
    if terms.len() == 1 {
        return Box::new(FuzzyTermQuery::new_prefix(
            Term::from_field_text(field, &terms[0]),
            0,
            true,
        ));
    }

    // Tantivy applies the phrase-prefix expansion limit to the segment term
    // dictionary during indexed search, but single-document evaluation can
    // only apply it to terms in the current document. For example, if the
    // segment's first 50 `c*` terms exclude `cb`, indexed search for `"x c*"`
    // will not match `"x cb"`, while fallback sees `cb` in that document and
    // may match it. This difference is accepted by the FTS fallback contract.
    Box::new(PhrasePrefixQuery::new(
        terms
            .into_iter()
            .map(|term_text| Term::from_field_text(field, &term_text))
            .collect(),
    ))
}

fn analyzed_query_terms(
    index_field: &FullTextIndexFieldDef,
    query_text: &str,
) -> Result<Vec<String>> {
    let mut analyzer = query_text_analyzer(index_field)?;
    let mut token_stream = analyzer.token_stream(query_text);
    let mut terms = Vec::new();
    let mut seen = HashSet::new();
    while token_stream.advance() {
        let term = token_stream.token().text.clone();
        if seen.insert(term.clone()) {
            terms.push(term);
        }
    }
    Ok(terms)
}

fn query_text_analyzer(index_field: &FullTextIndexFieldDef) -> Result<TextAnalyzer> {
    let option = index_field
        .has_index_field_option()
        .then(|| index_field.get_index_field_option())
        .ok_or_else(|| {
            anyhow!(
                "missing index_field_option for fulltext index field with col_id {}",
                index_field.get_col_id()
            )
        })?;
    match option.get_analyzer_type() {
        FullTextAnalyzerType::FtsAnalyzerExactValueV1 => {
            bail!("EXACT_VALUE_V1 queries must use exact query lowering")
        }
        FullTextAnalyzerType::FtsAnalyzerStandardV1 => Ok(crate::standard_v1()),
        FullTextAnalyzerType::FtsAnalyzerMultilingualV1 => Ok(crate::multilingual_v1()),
        FullTextAnalyzerType::FtsAnalyzerNgramV1 => {
            Ok(crate::NgramV1Spec::from_option(option.get_ngram())?.query_analyzer())
        }
        FullTextAnalyzerType::FtsAnalyzerPathHierarchyV1 => Ok(TextAnalyzer::from(
            crate::PathHierarchyTokenizer::new(option.get_path_hierarchy().get_delimiter())?,
        )),
        FullTextAnalyzerType::FtsAnalyzerUnspecified => {
            bail!("missing analyzer type for fulltext index field")
        }
    }
}

fn analyzer_type(index_field: &FullTextIndexFieldDef) -> Result<FullTextAnalyzerType> {
    if !index_field.has_index_field_option() {
        bail!(
            "missing index_field_option for fulltext index field with col_id {}",
            index_field.get_col_id()
        );
    }
    Ok(index_field.get_index_field_option().get_analyzer_type())
}

fn logical_index_field_name(index_field: &FullTextIndexFieldDef) -> Result<&'static str> {
    match analyzer_type(index_field)? {
        FullTextAnalyzerType::FtsAnalyzerExactValueV1 => Ok("EXACT_VALUE_V1"),
        FullTextAnalyzerType::FtsAnalyzerStandardV1 => Ok("STANDARD_V1"),
        FullTextAnalyzerType::FtsAnalyzerMultilingualV1 => Ok("MULTILINGUAL_V1"),
        FullTextAnalyzerType::FtsAnalyzerNgramV1 => Ok("NGRAM_V1"),
        FullTextAnalyzerType::FtsAnalyzerPathHierarchyV1 => Ok("PATH_HIERARCHY_V1"),
        FullTextAnalyzerType::FtsAnalyzerUnspecified => {
            bail!("missing analyzer type for fulltext index field")
        }
    }
}

enum PathHierarchyQueryTerm {
    MatchAll,
    Empty,
    Term(String),
}

fn path_hierarchy_query_term(
    index_field: &FullTextIndexFieldDef,
    query_text: &str,
) -> Result<PathHierarchyQueryTerm> {
    let option = index_field.get_index_field_option();
    let delimiter = option.get_path_hierarchy().get_delimiter();
    if delimiter.is_empty() {
        bail!(
            "path hierarchy field with col_id {} requires a non-empty delimiter",
            index_field.get_col_id()
        );
    }
    if query_text == delimiter {
        return Ok(PathHierarchyQueryTerm::MatchAll);
    }

    let mut term = query_text.to_string();
    while term.len() > delimiter.len() && term.ends_with(delimiter) {
        let new_len = term.len() - delimiter.len();
        term.truncate(new_len);
    }
    if term.is_empty() {
        Ok(PathHierarchyQueryTerm::Empty)
    } else {
        Ok(PathHierarchyQueryTerm::Term(term))
    }
}