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

use anyhow::{Result, anyhow, bail};
use serde_json::Number;
use tantivy::{
    Term,
    schema::{Field, OwnedValue},
};
use tidb_query_datatype::{Collation, FieldTypeAccessor, FieldTypeTp, codec::data_type::Real};
use tipb::ColumnInfo;

use crate::PhysicalFieldType;

/// Decodes one non-null native exact value from the index-side `ColumnBuffer`
/// format.
///
/// Text fields use the normal FTS text decoding path instead.
pub fn exact_value_from_column_buffer(column: &ColumnInfo, raw: &[u8]) -> Result<OwnedValue> {
    let expected = exact_physical_field_type(column)?;
    let value = match expected {
        PhysicalFieldType::I64 => {
            expect_raw_len(raw, 8, column)?;
            OwnedValue::I64(i64::from_le_bytes(raw.try_into().unwrap()))
        }
        PhysicalFieldType::U64 => {
            expect_raw_len(raw, 8, column)?;
            let value = match column.as_accessor().tp() {
                FieldTypeTp::Year => {
                    let value = i64::from_le_bytes(raw.try_into().unwrap());
                    u64::try_from(value).map_err(|_| {
                        anyhow!(
                            "exact FTS year column {} contains negative value {}",
                            column.get_column_id(),
                            value
                        )
                    })?
                }
                _ => u64::from_le_bytes(raw.try_into().unwrap()),
            };
            OwnedValue::U64(value)
        }
        PhysicalFieldType::F64 => {
            expect_raw_len(raw, 8, column)?;
            OwnedValue::F64(normalize_f64(f64::from_le_bytes(raw.try_into().unwrap()))?)
        }
        PhysicalFieldType::Bytes => bail!("exact bytes fields do not have a ColumnBuffer encoding"),
        PhysicalFieldType::Text => bail!("exact text fields use the FTS text decoding path"),
    };
    Ok(value)
}

/// Converts a JSON number into the native exact value selected by the physical
/// field type.
pub(crate) fn exact_value_from_query_number(
    field_type: PhysicalFieldType,
    number: &Number,
) -> Result<OwnedValue> {
    Ok(match field_type {
        PhysicalFieldType::I64 => OwnedValue::I64(
            number
                .as_i64()
                .ok_or_else(|| anyhow!("invalid exact i64 query number"))?,
        ),
        PhysicalFieldType::U64 => OwnedValue::U64(
            number
                .as_u64()
                .ok_or_else(|| anyhow!("invalid exact u64 query number"))?,
        ),
        PhysicalFieldType::F64 => OwnedValue::F64(normalize_f64(
            number
                .as_f64()
                .ok_or_else(|| anyhow!("invalid exact f64 query number"))?,
        )?),
        PhysicalFieldType::Bytes => bail!("exact bytes fields require a string query payload"),
        PhysicalFieldType::Text => bail!("text fields do not accept exact query values"),
    })
}

/// Converts a canonical exact value into its native Tantivy term.
pub(crate) fn exact_term_from_value(field: Field, value: OwnedValue) -> Result<Term> {
    Ok(match value {
        OwnedValue::I64(value) => Term::from_field_i64(field, value),
        OwnedValue::U64(value) => Term::from_field_u64(field, value),
        OwnedValue::F64(value) => Term::from_field_f64(field, value),
        OwnedValue::Bytes(value) => Term::from_field_bytes(field, &value),
        _ => bail!("non-exact Tantivy value cannot be converted into an exact term"),
    })
}

/// Resolves the Tantivy field type used by an `EXACT_VALUE_V1` field.
/// This function is shared with both value adapters so schema and encoding
/// cannot drift.
pub fn exact_physical_field_type(column: &ColumnInfo) -> Result<PhysicalFieldType> {
    let accessor = column.as_accessor();
    let tp = accessor.tp();
    Ok(match tp {
        FieldTypeTp::Tiny
        | FieldTypeTp::Short
        | FieldTypeTp::Int24
        | FieldTypeTp::Long
        | FieldTypeTp::LongLong => {
            if accessor.is_unsigned() {
                PhysicalFieldType::U64
            } else {
                PhysicalFieldType::I64
            }
        }
        // TiDB encodes BIT, YEAR and ENUM with unsigned integer semantics in
        // the proto / columnar path, so exact FTS can index them natively.
        FieldTypeTp::Bit | FieldTypeTp::Year | FieldTypeTp::Enum => PhysicalFieldType::U64,
        FieldTypeTp::Float | FieldTypeTp::Double => PhysicalFieldType::F64,
        FieldTypeTp::Date
        | FieldTypeTp::DateTime
        | FieldTypeTp::Timestamp
        | FieldTypeTp::NewDate => PhysicalFieldType::U64,
        FieldTypeTp::Duration => PhysicalFieldType::I64,
        FieldTypeTp::VarChar
        | FieldTypeTp::TinyBlob
        | FieldTypeTp::MediumBlob
        | FieldTypeTp::LongBlob
        | FieldTypeTp::Blob
        | FieldTypeTp::VarString
        | FieldTypeTp::String => {
            // ColumnInfo carries a collation id rather than a separate charset.
            // Binary collation identifies BINARY/VARBINARY/BLOB values; every
            // other supported collation represents character data.
            let collation = accessor
                .collation()
                .map_err(|err| anyhow!(err.to_string()))?;
            if collation == Collation::Binary {
                bail!("exact FTS does not support binary string columns");
            }
            PhysicalFieldType::Text
        }
        FieldTypeTp::Unspecified
        | FieldTypeTp::Null
        | FieldTypeTp::NewDecimal
        | FieldTypeTp::Json
        | FieldTypeTp::Set
        | FieldTypeTp::Geometry
        | FieldTypeTp::TiDbVectorFloat32 => {
            bail!("exact FTS does not support column type {:?}", tp)
        }
    })
}

fn normalize_f64(value: f64) -> Result<f64> {
    Real::new(value).map_err(|_| anyhow!("exact FTS does not support NaN"))?;
    if value == 0.0 {
        return Ok(0.0);
    }
    Ok(value)
}

fn expect_raw_len(raw: &[u8], expected: usize, column: &ColumnInfo) -> Result<()> {
    if raw.len() != expected {
        bail!(
            "exact FTS column {} type {} requires {} bytes, got {}",
            column.get_column_id(),
            column.get_tp(),
            expected,
            raw.len()
        );
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use tidb_query_datatype::Collation;

    use super::*;

    fn column(tp: FieldTypeTp, collation: Collation) -> ColumnInfo {
        let mut column = ColumnInfo::new();
        column.set_column_id(10);
        column.set_tp(tp as i32);
        column.set_collation(collation as i32);
        column
    }

    #[test]
    fn character_strings_resolve_to_text_and_skip_native_adapter() -> Result<()> {
        for tp in [
            FieldTypeTp::String,
            FieldTypeTp::VarChar,
            FieldTypeTp::VarString,
            FieldTypeTp::TinyBlob,
            FieldTypeTp::Blob,
            FieldTypeTp::MediumBlob,
            FieldTypeTp::LongBlob,
        ] {
            let text_column = column(tp, Collation::Utf8Mb4GeneralCi);
            assert_eq!(
                exact_physical_field_type(&text_column)?,
                PhysicalFieldType::Text
            );
            assert!(exact_value_from_column_buffer(&text_column, b"TiDB").is_err());

            let binary_column = column(tp, Collation::Binary);
            assert!(exact_physical_field_type(&binary_column).is_err());
            assert!(exact_value_from_column_buffer(&binary_column, b"TiDB").is_err());
        }
        Ok(())
    }

    #[test]
    fn canonicalizes_signed_zero_for_index_and_query_terms() -> Result<()> {
        let field = Field::from_field_id(0);
        let positive_zero_term = Term::from_field_f64(field, 0.0);
        for tp in [FieldTypeTp::Float, FieldTypeTp::Double] {
            let column = column(tp, Collation::Binary);
            for value in [-0.0f64, 0.0] {
                let OwnedValue::F64(indexed) =
                    exact_value_from_column_buffer(&column, &value.to_le_bytes())?
                else {
                    panic!("expected f64 exact value");
                };
                assert_eq!(indexed.to_bits(), 0.0f64.to_bits());

                let number = Number::from_f64(value).unwrap();
                assert_eq!(
                    exact_term_from_value(
                        field,
                        exact_value_from_query_number(PhysicalFieldType::F64, &number)?,
                    )?,
                    positive_zero_term,
                );
            }
        }
        Ok(())
    }

    #[test]
    fn enum_exact_values_resolve_and_decode_as_u64() -> Result<()> {
        let column = column(FieldTypeTp::Enum, Collation::Utf8Mb4GeneralCi);
        assert_eq!(exact_physical_field_type(&column)?, PhysicalFieldType::U64);

        let OwnedValue::U64(empty) = exact_value_from_column_buffer(&column, &0_u64.to_le_bytes())?
        else {
            panic!("expected u64 exact value for enum");
        };
        assert_eq!(empty, 0);

        let OwnedValue::U64(named) = exact_value_from_column_buffer(&column, &7_u64.to_le_bytes())?
        else {
            panic!("expected u64 exact value for enum");
        };
        assert_eq!(named, 7);
        Ok(())
    }

    #[test]
    fn decimal_and_json_exact_are_not_supported() {
        for tp in [FieldTypeTp::NewDecimal, FieldTypeTp::Json] {
            let column = column(tp, Collation::Binary);
            let err = exact_physical_field_type(&column).unwrap_err();
            assert!(
                err.to_string()
                    .contains(&format!("exact FTS does not support column type {tp:?}"))
            );
        }
    }
}