// Copyright 2026 PingCAP, Inc.

use anyhow::{Context, bail};
use clara_fts::Query;
use tantivy::{
    TantivyDocument,
    schema::{Field, OwnedValue},
};
use tidb_query_datatype::{FieldTypeAccessor, FieldTypeFlag, FieldTypeTp};
use tipb::ColumnInfo;

use crate::table::{
    Error, Result,
    columnar::Block,
    fts::util::{StringifyFn, build_stringify_fn, empty_stringify_fn},
    schema_file::Schema,
};

/// Column id reserved for the virtual score column appended when
/// `FtsQueryTypeWithScore` is requested. The column is expected to be the
/// last entry in the output schema with type `Float` and `NOT NULL`.
pub const VIRTUAL_SCORE_COLUMN_ID: i64 = -2050;

/// Ensures the provided logical `Schema` matches requirements implied by the
/// FTS query:
/// - For `FtsQueryTypeNoScore`, the schema must NOT contain the virtual score
///   column.
/// - For `FtsQueryTypeWithScore`, the last column must be the virtual score
///   column with type `Float` and `NOT NULL`.
pub fn validate_schema(schema: &Schema, query: &Query) -> anyhow::Result<()> {
    #[cfg(debug_assertions)]
    if query.info().get_query_type() == tipb::FtsQueryType::FtsQueryTypeNoScore {
        // FtsQueryTypeNoScore should not contain score column in schema
        if schema
            .columns
            .iter()
            .any(|c| c.get_column_id() == VIRTUAL_SCORE_COLUMN_ID)
        {
            bail!(
                "unexpected score col (id={}) at pos {} for FtsQueryTypeNoScore",
                VIRTUAL_SCORE_COLUMN_ID,
                schema
                    .columns
                    .iter()
                    .position(|c| c.get_column_id() == VIRTUAL_SCORE_COLUMN_ID)
                    .unwrap()
            );
        }
    }

    if query.info().get_query_type() == tipb::FtsQueryType::FtsQueryTypeWithScore {
        // FtsQueryTypeWithScore should contain score column in schema
        if schema.columns.is_empty() {
            bail!("unexpected empty schema for FtsQueryTypeWithScore");
        }
        let last_col = schema.columns.last().unwrap();
        check_score_col(last_col, false)?;
    }

    Ok(())
}

/// Checks whether the given column is a valid score column.
pub fn check_score_col(col: &ColumnInfo, require_nullable: bool) -> anyhow::Result<()> {
    if col.get_column_id() != VIRTUAL_SCORE_COLUMN_ID {
        bail!(
            "column id={} is not the virtual score column id={}",
            col.get_column_id(),
            VIRTUAL_SCORE_COLUMN_ID
        );
    }
    if col.get_tp() != FieldTypeTp::Float as i32 {
        bail!("score column must be float column, got {:?}", col.get_tp());
    }
    if col.as_accessor().flag().contains(FieldTypeFlag::NOT_NULL) == require_nullable {
        bail!(
            "score column nullable={} does not match required={}",
            !col.as_accessor().flag().contains(FieldTypeFlag::NOT_NULL),
            require_nullable
        );
    }
    Ok(())
}

fn check_fts_source_col(
    col: &ColumnInfo,
    expected_column_id: i64,
    query: &Query,
) -> anyhow::Result<()> {
    if col.get_column_id() != expected_column_id {
        bail!(
            "FTS source column id={} does not match query source column id={}",
            col.get_column_id(),
            expected_column_id
        );
    }
    let source_field_type = clara_fts::exact_physical_field_type(col)?;
    for binding in query.document_fields() {
        if binding.column_id() == expected_column_id && binding.field_type() != source_field_type {
            bail!(
                "FTS source column {} has physical type {:?}, but query field {:?} requires {:?}",
                expected_column_id,
                source_field_type,
                binding.field(),
                binding.field_type(),
            );
        }
    }
    Ok(())
}

pub(crate) fn validate_fts_source_score_schema(
    schema: &Schema,
    query: &Query,
    reader_name: &str,
) -> Result<()> {
    let expected_source_columns = document_column_ids(query);
    if schema.columns.len() != expected_source_columns.len() + 1 {
        return Err(Error::Other(format!(
            "{} expected {} source columns and one score column, got {} columns",
            reader_name,
            expected_source_columns.len(),
            schema.columns.len(),
        )));
    }
    for (column_idx, column_id) in expected_source_columns.iter().enumerate() {
        check_fts_source_col(&schema.columns[column_idx], *column_id, query)
            .map_err(|error| Error::Other(error.to_string()))?;
    }
    check_score_col(schema.columns.last().unwrap(), true)
        .map_err(|error| Error::Other(error.to_string()))
}

struct BruteDocumentColumn {
    column_idx: usize,
    column_info: ColumnInfo,
    field_type: clara_fts::PhysicalFieldType,
    stringify: StringifyFn,
    fields: Vec<Field>,
}

/// Builds Tantivy documents from the physical source columns required by a
/// resolved FTS query.
///
/// NULL values and text decoding failures are materialized as explicit
/// non-searchable [`OwnedValue::Null`] fields. Successfully decoded text,
/// including an empty string, remains searchable. Unsupported text collations
/// follow index-build behavior and materialize as empty text.
pub(crate) struct BruteDocumentBuilder {
    columns: Vec<BruteDocumentColumn>,
}

impl BruteDocumentBuilder {
    pub(crate) fn new(schema: &Schema, query: &Query) -> anyhow::Result<Self> {
        let mut columns = Vec::<BruteDocumentColumn>::new();
        for binding in query.document_fields() {
            let column_idx = schema
                .columns
                .iter()
                .position(|column| column.get_column_id() == binding.column_id())
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "missing FTS source column {} in brute-force reader schema",
                        binding.column_id()
                    )
                })?;
            let column_info = &schema.columns[column_idx];
            let source_field_type = clara_fts::exact_physical_field_type(column_info)
                .with_context(|| {
                    format!(
                        "unsupported FTS source column {}",
                        column_info.get_column_id()
                    )
                })?;
            if binding.field_type() != source_field_type {
                bail!(
                    "FTS source column {} has physical type {:?}, but query field {:?} requires {:?}",
                    binding.column_id(),
                    source_field_type,
                    binding.field(),
                    binding.field_type(),
                );
            }
            if let Some(column) = columns
                .iter_mut()
                .find(|column| column.column_idx == column_idx)
            {
                column.fields.push(binding.field());
                continue;
            }
            let stringify =
                build_stringify_fn(column_info).unwrap_or_else(|_| empty_stringify_fn());
            columns.push(BruteDocumentColumn {
                column_idx,
                column_info: column_info.clone(),
                field_type: source_field_type,
                stringify,
                fields: vec![binding.field()],
            });
        }
        Ok(Self { columns })
    }

    pub(crate) fn build(&self, block: &Block, row: usize) -> anyhow::Result<TantivyDocument> {
        let mut document = TantivyDocument::new();
        for column in &self.columns {
            let raw = block.columns[column.column_idx].get_value(row);
            let value = match (raw, column.field_type) {
                (None, _) => OwnedValue::Null,
                (Some(raw), clara_fts::PhysicalFieldType::Text) => match (column.stringify)(raw) {
                    Ok(text) => OwnedValue::Str(text.into_owned()),
                    Err(_) => OwnedValue::Null,
                },
                (Some(raw), _) => {
                    clara_fts::exact_value_from_column_buffer(&column.column_info, raw)
                        .with_context(|| {
                            format!(
                                "failed to decode exact FTS source column {}",
                                column.column_info.get_column_id()
                            )
                        })?
                }
            };
            for &field in &column.fields {
                document.add_field_value(field, value.clone());
            }
        }
        Ok(document)
    }
}

pub(crate) fn document_column_ids(query: &Query) -> Vec<i64> {
    let mut column_ids = Vec::new();
    for field in query.document_fields() {
        if !column_ids.contains(&field.column_id()) {
            column_ids.push(field.column_id());
        }
    }
    column_ids
}

#[cfg(test)]
mod tests {
    use clara_fts::{
        BruteScoredSearcher, FTS_MATCH_EXPR_METADATA_VERSION, IndexLayout, Query,
        SingleFieldLayout, test_util::PlainFtsQueryInfo,
    };
    use kvenginepb::fts::{FullTextAnalyzerType, FullTextIndexFieldDef, FullTextIndexFieldOption};
    use tantivy::{query::DocumentEvaluation, schema::Field};
    use tidb_query_datatype::FieldTypeTp;

    use super::BruteDocumentBuilder;
    use crate::table::fts::{
        IntPk,
        test_util::{SchemaBuilder, make_exact_match_expr_query, new_block},
    };

    const TABLE_ID: i64 = 7;
    const TEXT_COL_ID: i64 = 100;

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

    fn exact_index_field(column_id: i64) -> FullTextIndexFieldDef {
        let mut field = FullTextIndexFieldDef::new();
        field.set_col_id(column_id);
        field
            .mut_index_field_option()
            .set_analyzer_type(FullTextAnalyzerType::FtsAnalyzerExactValueV1);
        field
    }

    fn match_word(column_id: i64, text: &str) -> tipb::Expr {
        let mut expr = tipb::Expr::new();
        expr.set_tp(tipb::ExprType::ScalarFunc);
        expr.set_sig(tipb::ScalarFuncSig::FtsMatchWord);
        expr.set_val(
            serde_json::to_vec(&serde_json::json!({
                "version": FTS_MATCH_EXPR_METADATA_VERSION,
                "score_mode": "inherit",
                "query": text,
                "fields": [{
                    "col_id": column_id,
                    "index_field": "STANDARD_V1",
                }],
            }))
            .unwrap(),
        );
        expr
    }

    fn boolean(sig: tipb::ScalarFuncSig, left: tipb::Expr, right: tipb::Expr) -> tipb::Expr {
        let mut expr = tipb::Expr::new();
        expr.set_tp(tipb::ExprType::ScalarFunc);
        expr.set_sig(sig);
        expr.set_val(br#"{"version":1}"#.to_vec());
        expr.mut_children().push(left);
        expr.mut_children().push(right);
        expr
    }

    fn unary_not(child: tipb::Expr) -> tipb::Expr {
        let mut expr = tipb::Expr::new();
        expr.set_tp(tipb::ExprType::ScalarFunc);
        expr.set_sig(tipb::ScalarFuncSig::UnaryNotInt);
        expr.set_val(br#"{"version":1}"#.to_vec());
        expr.mut_children().push(child);
        expr
    }

    fn match_expr_query_info() -> tipb::FtsQueryInfo {
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr()
            .push(match_word(TEXT_COL_ID, "search"));
        info
    }

    fn legacy_query_info() -> tipb::FtsQueryInfo {
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info
    }

    #[test]
    fn validate_match_expr_accepts_empty_legacy_columns() {
        let mut info = match_expr_query_info();
        info.mut_columns().clear();

        Query::validate_info(&info).unwrap();
    }

    #[test]
    fn validate_match_expr_does_not_apply_legacy_column_validation() {
        let mut info = match_expr_query_info();
        info.mut_columns().push(tipb::ColumnInfo::new());

        Query::validate_info(&info).unwrap();
    }

    #[test]
    fn validate_match_expr_still_requires_index_id() {
        let mut info = match_expr_query_info();
        info.clear_index_id();
        info.mut_columns().clear();

        assert_eq!(
            Query::validate_info(&info).err().unwrap().to_string(),
            "missing fts_query.index_id"
        );
    }

    #[test]
    fn validate_legacy_query_keeps_single_string_column_contract() {
        Query::validate_info(&legacy_query_info()).unwrap();

        let mut empty = legacy_query_info();
        empty.mut_columns().clear();
        assert_eq!(
            Query::validate_info(&empty).err().unwrap().to_string(),
            "legacy FTS query requires exactly one column, got 0"
        );

        let mut multiple = legacy_query_info();
        multiple.mut_columns().push(tipb::ColumnInfo::new());
        assert_eq!(
            Query::validate_info(&multiple).err().unwrap().to_string(),
            "legacy FTS query requires exactly one column, got 2"
        );

        let mut missing_column_id = legacy_query_info();
        missing_column_id.mut_columns().clear();
        let mut string_column = tipb::ColumnInfo::new();
        string_column.set_tp(FieldTypeTp::VarString as i32);
        missing_column_id.mut_columns().push(string_column);
        assert_eq!(
            Query::validate_info(&missing_column_id)
                .err()
                .unwrap()
                .to_string(),
            "missing fts_query.column.column_id"
        );

        let mut non_string = legacy_query_info();
        non_string.mut_columns()[0].set_tp(FieldTypeTp::LongLong as i32);
        assert!(
            Query::validate_info(&non_string)
                .err()
                .unwrap()
                .to_string()
                .contains("legacy fts_query.column.tp must be string")
        );
    }

    #[test]
    fn builds_multi_column_tantivy_documents() {
        let first_column_id = 101;
        let second_column_id = 202;
        let schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column(first_column_id, FieldTypeTp::VarString)
            .column(second_column_id, FieldTypeTp::VarString)
            .schema();
        let block = new_block::<IntPk>(&schema, |row| {
            row(1, 10, false, |datum| {
                datum("rust language");
                datum("search engine");
            });
            row(2, 10, false, |datum| {
                datum("rust language");
                datum("database");
            });
        });

        let layout = IndexLayout::from_field_names(["0", "1"]);
        let index_fields = vec![
            standard_index_field(first_column_id),
            standard_index_field(second_column_id),
        ];
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr()
            .push(match_word(first_column_id, "rust"));
        info.mut_match_expr()
            .push(match_word(second_column_id, "search"));
        let query = Query::new_with_index_layout(info, &index_fields, &layout).unwrap();

        let builder = BruteDocumentBuilder::new(&schema, &query).unwrap();
        let mut searcher = BruteScoredSearcher::new(&query).unwrap();
        let mut matched_rows = Vec::new();
        for row in 0..block.length() {
            let document = builder.build(&block, row).unwrap();
            if let DocumentEvaluation::Match(_) = searcher.evaluate_document(&document).unwrap() {
                matched_rows.push(row as u32);
            }
        }
        assert_eq!(matched_rows, vec![0]);
    }

    #[test]
    fn exact_empty_text_is_distinct_from_null() {
        let schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column(TEXT_COL_ID, FieldTypeTp::VarString)
            .schema();
        let block = new_block::<IntPk>(&schema, |row| {
            row(1, 10, false, |datum| datum(""));
            row(2, 10, false, |datum| datum.null());
            row(3, 10, false, |datum| datum("non-empty"));
        });

        let index_fields = vec![exact_index_field(TEXT_COL_ID)];
        let index_layout = IndexLayout::new(vec![SingleFieldLayout::new_with_field_and_bm25(
            Field::from_field_id(0),
            "0",
            false,
        )]);
        let query =
            make_exact_match_expr_query("", TEXT_COL_ID, 1, &index_fields, &index_layout).unwrap();
        let builder = BruteDocumentBuilder::new(&schema, &query).unwrap();
        let mut searcher = BruteScoredSearcher::new(&query).unwrap();

        let empty_document = builder.build(&block, 0).unwrap();
        assert!(matches!(
            searcher.evaluate_document(&empty_document).unwrap(),
            DocumentEvaluation::Match(_)
        ));
        for row in 1..block.length() {
            let document = builder.build(&block, row).unwrap();
            assert_eq!(
                searcher.evaluate_document(&document).unwrap(),
                DocumentEvaluation::NoMatch
            );
        }
    }

    #[test]
    fn malformed_text_is_null_and_other_fields_still_match() {
        let malformed_column_id = 101;
        let valid_column_id = 202;
        let schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column(malformed_column_id, FieldTypeTp::VarString)
            .column(valid_column_id, FieldTypeTp::VarString)
            .schema();
        let block = new_block::<IntPk>(&schema, |row| {
            row(1, 10, false, |datum| {
                datum(&[0xff][..]);
                datum("search engine");
            });
        });

        let layout = IndexLayout::from_field_names(["0", "1"]);
        let index_fields = vec![
            standard_index_field(malformed_column_id),
            standard_index_field(valid_column_id),
        ];
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr().push(boolean(
            tipb::ScalarFuncSig::LogicalOr,
            match_word(malformed_column_id, "rust"),
            match_word(valid_column_id, "search"),
        ));
        let query = Query::new_with_index_layout(info, &index_fields, &layout).unwrap();

        let builder = BruteDocumentBuilder::new(&schema, &query).unwrap();
        let document = builder.build(&block, 0).unwrap();
        let mut searcher = BruteScoredSearcher::new(&query).unwrap();
        assert_eq!(
            searcher.evaluate_document(&document).unwrap(),
            DocumentEvaluation::Match(1.0)
        );
    }

    #[test]
    fn malformed_text_matches_not_filter_as_missing_field() {
        let schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column(TEXT_COL_ID, FieldTypeTp::VarString)
            .schema();
        let block = new_block::<IntPk>(&schema, |row| {
            row(1, 10, false, |datum| {
                datum(&[0xff][..]);
            });
        });

        let layout = IndexLayout::legacy_body();
        let index_fields = vec![standard_index_field(TEXT_COL_ID)];
        let mut info = PlainFtsQueryInfo::default().as_query_info();
        info.set_query_type(tipb::FtsQueryType::FtsQueryTypeNoScore);
        info.mut_match_expr()
            .push(unary_not(match_word(TEXT_COL_ID, "blocked")));
        let query = Query::new_with_index_layout(info, &index_fields, &layout).unwrap();

        let builder = BruteDocumentBuilder::new(&schema, &query).unwrap();
        let document = builder.build(&block, 0).unwrap();
        let mut searcher = BruteScoredSearcher::new(&query).unwrap();
        assert_eq!(
            searcher.evaluate_document(&document).unwrap(),
            DocumentEvaluation::Match(1.0)
        );
    }
}