// Copyright 2026 PingCAP, Inc.

use std::{collections::HashMap, io::Cursor, mem::size_of, sync::Arc};

use anyhow::{Context, Result, bail};
use bytes::Bytes;
use clara_fts::{IndexReader as ClaraIndexReader, TantivyIndexWriter};
use quick_cache::{Weighter, sync::Cache};
use tikv_util::{
    config::{AbsoluteOrPercentSize, ReadableSize},
    sys::SysQuota,
};

use super::util::build_fts_value_fn;
use crate::table::{
    ChecksumType, SnapVersion,
    blobtable::blobtable::BlobTable,
    columnar::{Block, ColumnarReader, ColumnarRowTableReader, ColumnarTableReader},
    file::InMemFile,
    fts::{
        FtsCache, PackedFile, compact,
        packed_file::{EPackedFileLp, PackedFileBuilder, PackedFileBuilderOptions},
    },
    memtable::CfTable,
    schema_file::Schema,
    sstable::L0Table,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum FtsDeltaKind {
    /// A sealed (immutable) memtable.
    /// FIXME: Currently in read node we can only get one virtual memtable from
    /// the snapshot, so this kind is currently unused.
    Memtable,
    /// A row-format L0 (SST L0) that is not converted to columnar yet.
    RowL0,
    /// A columnar L0 file that is not covered by the official FTS index yet.
    ColumnarL0,
}

/// Configuration for [`FtsDeltaCache`].
#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct FtsDeltaCacheConfig {
    /// Maximum capacity of the cache.
    pub cache_cap: AbsoluteOrPercentSize,
}

impl Default for FtsDeltaCacheConfig {
    fn default() -> Self {
        Self {
            cache_cap: AbsoluteOrPercentSize::Abs(ReadableSize::mb(64)),
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
struct CacheKey {
    kind: FtsDeltaKind,
    source_id: u64,
    table_id: i64,
    index_id: i64,
}

#[derive(Clone)]
pub enum FtsDeltaInput {
    /// FIXME: Currently in read node we can only get one virtual memtable from
    /// the snapshot, so this kind is currently unused.
    Memtable {
        mem: CfTable,
        snap_version: SnapVersion,
    },
    RowL0 {
        l0: L0Table,
        snap_version: SnapVersion,
    },
    ColumnarL0 {
        file: crate::table::columnar::ColumnarFile,
        snap_version: SnapVersion,
    },
}

impl FtsDeltaInput {
    pub fn snap_version(&self) -> SnapVersion {
        match self {
            Self::Memtable { snap_version, .. }
            | Self::RowL0 { snap_version, .. }
            | Self::ColumnarL0 { snap_version, .. } => *snap_version,
        }
    }
}

#[derive(Clone)]
pub struct FtsDeltaBuildSpec {
    pub kind: FtsDeltaKind,
    pub source_id: u64,
    pub table_id: i64,
    pub index_id: i64,
    pub input: FtsDeltaInput,
}

impl FtsDeltaBuildSpec {
    #[inline]
    fn cache_key(&self) -> CacheKey {
        CacheKey {
            kind: self.kind,
            source_id: self.source_id,
            table_id: self.table_id,
            index_id: self.index_id,
        }
    }
}

/// One cached delta index entry.
///
/// The entry is aligned to the on-disk packed file layout so that we can reuse
/// the existing MVCC/has_newer_version logic implemented on [`EPackedFileLp`].
#[derive(Clone)]
pub struct FtsDeltaCacheEntry {
    lp: Arc<EPackedFileLp>,
    index_reader: Arc<ClaraIndexReader>,
}

impl FtsDeltaCacheEntry {
    #[inline]
    pub fn lp(&self) -> &EPackedFileLp {
        &self.lp
    }

    #[inline]
    pub fn index_reader(&self) -> &Arc<ClaraIndexReader> {
        &self.index_reader
    }

    #[inline]
    fn estimated_size(&self) -> u64 {
        self.lp.serialized_size() as u64 + size_of::<Self>() as u64
    }
}

#[derive(Clone)]
struct CacheWeighter;

impl Weighter<CacheKey, Arc<FtsDeltaCacheEntry>> for CacheWeighter {
    fn weight(&self, _key: &CacheKey, value: &Arc<FtsDeltaCacheEntry>) -> u64 {
        size_of::<CacheKey>() as u64 + value.estimated_size()
    }
}

type DeltaCacheInner = Cache<CacheKey, Arc<FtsDeltaCacheEntry>, CacheWeighter>;

#[derive(Clone)]
pub struct FtsDeltaSource {
    pub entry: Arc<FtsDeltaCacheEntry>,
    /// Ordering key used by the MVCC shadowing algorithm. Higher means newer.
    pub seq: u64,
}

/// A cache for per-source, per-(table_id,index_id) in-memory FTS indexes for
/// the "unindexed" part of FTS queries.
///
/// This cache is designed for immutable sources only (sealed memtables, row
/// L0 and untracked columnar L0). Writable memtables are intentionally
/// excluded.
#[derive(Clone)]
pub struct FtsDeltaCache {
    inner: Option<Arc<DeltaCacheInner>>,
}

impl FtsDeltaCache {
    /// Creates a cache with the given total weight capacity in bytes.
    ///
    /// A capacity of 0 disables the cache.
    pub fn new(capacity_bytes: u64) -> Self {
        if capacity_bytes == 0 {
            return Self::disabled();
        }

        let cache_shards = (SysQuota::cpu_cores_quota() as usize).max(1) * 8;
        // A very rough heuristic; entries can vary wildly but we only need a
        // reasonable starting point for allocation.
        let estimated_items_capacity = (capacity_bytes as usize / (4 << 20)).max(64);

        let opts = quick_cache::OptionsBuilder::new()
            .shards(cache_shards)
            .weight_capacity(capacity_bytes)
            .estimated_items_capacity(estimated_items_capacity)
            .build()
            .unwrap();
        let cache = Cache::with_options(
            opts,
            CacheWeighter,
            quick_cache::DefaultHashBuilder::default(),
            quick_cache::sync::DefaultLifecycle::default(),
        );

        Self {
            inner: Some(Arc::new(cache)),
        }
    }

    /// Creates a disabled cache instance that never stores any entry.
    pub fn disabled() -> Self {
        Self { inner: None }
    }

    #[inline]
    pub fn is_enabled(&self) -> bool {
        self.inner.is_some()
    }

    /// Returns a cached delta index entry if present.
    #[inline]
    pub fn get(
        &self,
        kind: FtsDeltaKind,
        source_id: u64,
        table_id: i64,
        index_id: i64,
    ) -> Option<Arc<FtsDeltaCacheEntry>> {
        let key = CacheKey {
            kind,
            source_id,
            table_id,
            index_id,
        };
        self.inner.as_ref().and_then(|cache| cache.get(&key))
    }

    /// Invalidates all cached entries for a source (across all table/index
    /// combinations).
    ///
    /// Currently unused, because Read Node always grabs a snapshot, it does not
    /// know what is really invalidated. But it is provided for future use
    /// cases.
    pub fn invalidate_memtable(&self, memtable_id: u64) {
        self.invalidate_source(FtsDeltaKind::Memtable, memtable_id);
    }

    /// Currently unused, because Read Node always grabs a snapshot, it does not
    /// know what is really invalidated. But it is provided for future use
    /// cases.
    pub fn invalidate_row_l0(&self, l0_id: u64) {
        self.invalidate_source(FtsDeltaKind::RowL0, l0_id);
    }

    /// Currently unused, because Read Node always grabs a snapshot, it does not
    /// know what is really invalidated. But it is provided for future use
    /// cases.
    pub fn invalidate_columnar_l0(&self, file_id: u64) {
        self.invalidate_source(FtsDeltaKind::ColumnarL0, file_id);
    }

    fn invalidate_source(&self, kind: FtsDeltaKind, source_id: u64) {
        let Some(cache) = &self.inner else {
            return;
        };
        cache.retain(|key, _| key.kind != kind || key.source_id != source_id);
    }

    /// Builds (or retrieves) a delta index entry for the given spec.
    ///
    /// This is a lazy, singleflight operation: concurrent callers for the same
    /// key will share the build.
    ///
    /// Returns `Ok(None)` when the cache is disabled or the source has no rows
    /// for this table.
    ///
    /// Returns `Err` when building the delta index fails.
    pub async fn get_or_build(
        &self,
        spec: FtsDeltaBuildSpec,
        schema: Schema,
        tracked_index_layout: &crate::table::fts::TrackedIndexLayout,
        blob_tables: Arc<HashMap<u64, BlobTable>>,
        encryption_key: Option<cloud_encryption::EncryptionKey>,
    ) -> Result<Option<Arc<FtsDeltaCacheEntry>>> {
        let Some(cache) = self.inner.as_ref() else {
            // FIXME: We should possibly still build the index even when the cache is
            // disabled.
            return Ok(None);
        };
        let key = spec.cache_key();
        match cache.get_value_or_guard_async(&key).await {
            Ok(entry) => Ok(Some(entry)),
            Err(guard) => Ok(build_delta_index(
                &spec.input,
                key,
                schema,
                tracked_index_layout,
                blob_tables,
                encryption_key,
            )
            .await
            .with_context(|| {
                format!(
                    "delta index build failed (kind={:?}, source_id={}, table_id={}, index_id={})",
                    key.kind, key.source_id, key.table_id, key.index_id
                )
            })?
            .map(|built| {
                let entry = Arc::new(built);
                let _ = guard.insert(Arc::clone(&entry));
                entry
            })),
        }
    }
}

async fn build_delta_index(
    input: &FtsDeltaInput,
    key: CacheKey,
    schema: Schema,
    tracked_index_layout: &crate::table::fts::TrackedIndexLayout,
    blob_tables: Arc<HashMap<u64, BlobTable>>,
    encryption_key: Option<cloud_encryption::EncryptionKey>,
) -> Result<Option<FtsDeltaCacheEntry>> {
    let mut buffer = Vec::new();
    let mut builder = PackedFileBuilder::new(
        Cursor::new(&mut buffer),
        PackedFileBuilderOptions {
            checksum_type: ChecksumType::None,
            ..Default::default()
        },
    );

    let lp_key = compact::lp_key(key.table_id, key.index_id);
    let is_int_handle = !schema.is_common_handle();
    builder.start_lp(key.table_id, key.index_id, is_int_handle, &lp_key)?;

    let index_def = schema
        .fulltext_indexes
        .iter()
        .find(|index_def| index_def.get_index_id() == key.index_id)
        .ok_or_else(|| {
            anyhow::anyhow!(
                "missing fulltext index definition for table_id {}, index_id {}",
                key.table_id,
                key.index_id
            )
        })?;
    let tracked_index_fields = crate::table::fts::tracked_index_fields_from_layout_or_index_def(
        tracked_index_layout,
        index_def,
    )?;
    if tracked_index_fields.len() != tracked_index_layout.index_layout.fields().len() {
        bail!(
            "delta tracked field count {} does not match layout field count {}",
            tracked_index_fields.len(),
            tracked_index_layout.index_layout.fields().len()
        );
    }
    let col_id_to_pos = schema
        .columns
        .iter()
        .enumerate()
        .map(|(idx, col)| (col.get_column_id(), idx))
        .collect::<HashMap<_, _>>();
    let resolved_col_positions = tracked_index_fields
        .iter()
        .map(|field| {
            let col_id = field.get_col_id();
            col_id_to_pos.get(&col_id).copied().ok_or_else(|| {
                anyhow::anyhow!(
                    "delta index schema is missing tracked column {} for table_id {}, index_id {}",
                    col_id,
                    key.table_id,
                    key.index_id
                )
            })
        })
        .collect::<Result<Vec<_>>>()?;
    let value_converters = resolved_col_positions
        .iter()
        .zip(tracked_index_layout.index_layout.fields())
        .map(|(&col_pos, field_layout)| {
            build_fts_value_fn(&schema.columns[col_pos], field_layout.field_type())
        })
        .collect::<Vec<_>>();
    let mut index_writer = TantivyIndexWriter::new_in_memory_with_index_fields(
        &tracked_index_fields,
        &tracked_index_layout.index_layout,
    )?;

    let mut reader: Box<dyn ColumnarReader> = match input {
        FtsDeltaInput::Memtable { mem, .. } => {
            // FIXME: Currently in read node we can only get one virtual
            // memtable from the snapshot, so this kind is
            // currently unused.
            let iter = mem.get_cf(crate::WRITE_CF).new_iterator(false);
            Box::new(ColumnarRowTableReader::new(
                schema.clone(),
                iter,
                Some(blob_tables),
                false,
                encryption_key,
            )?)
        }
        FtsDeltaInput::RowL0 { l0, .. } => {
            let Some(l0_write) = l0.get_cf(crate::WRITE_CF) else {
                return Ok(None);
            };
            let iter = l0_write.new_iterator(false, true);
            Box::new(ColumnarRowTableReader::new(
                schema.clone(),
                iter,
                Some(blob_tables),
                false,
                encryption_key,
            )?)
        }
        FtsDeltaInput::ColumnarL0 { file, .. } => Box::new(ColumnarTableReader::new(
            file,
            schema.clone(),
            None,
            encryption_key,
        )?),
    };

    reader.seek(&[]).await?;

    let mut any_rows = false;
    let mut block = Block::new(&schema);

    loop {
        block.reset();
        let rows = reader.read(&mut block, 1024).await?;
        if rows == 0 {
            break;
        }
        any_rows = true;

        // Tantivy docs align with doc_id in the packed LP.
        let mut document_fields =
            Vec::with_capacity(tracked_index_layout.index_layout.fields().len());
        for row in 0..rows {
            document_fields.clear();
            for (field_i, (&col_pos, field_layout)) in resolved_col_positions
                .iter()
                .zip(tracked_index_layout.index_layout.fields().iter())
                .enumerate()
            {
                let Some(raw) = block.columns[col_pos].get_value(row) else {
                    continue;
                };
                let Some(value) = (value_converters[field_i])(raw)? else {
                    continue;
                };
                document_fields.push((field_layout.field(), value));
            }
            if document_fields.is_empty() {
                index_writer.add_null()?;
            } else {
                index_writer.add_document_values(document_fields.drain(..))?;
            }
        }

        if is_int_handle {
            for row in 0..rows {
                let version = block.versions.get_version(row);
                let is_deleted = block.versions.is_null(row);
                let pk = block.handles.get_int_handle_value(row);
                builder.add_pk_int(pk, version, is_deleted)?;
            }
        } else {
            for row in 0..rows {
                let version = block.versions.get_version(row);
                let is_deleted = block.versions.is_null(row);
                let pk = block.handles.get_not_null_value(row);
                builder.add_pk_common(pk, version, is_deleted)?;
            }
        }
    }

    if !any_rows {
        return Ok(None);
    }

    let dir = index_writer.finalize_as_dir()?;
    builder.finish_lp(&dir)?;
    builder.finish(input.snap_version())?;

    let bytes = Bytes::from(buffer);
    let file_id = key.source_id;
    let packed = PackedFile::new(
        Arc::new(InMemFile::new(file_id, bytes)),
        FtsCache::disabled(),
    )?;
    let lp = packed
        .cached_get_lp(&lp_key)
        .await?
        .expect("delta packed file must contain exactly one lp");
    let index_reader = lp.read_tantivy_index()?;
    Ok(Some(FtsDeltaCacheEntry { lp, index_reader }))
}

#[cfg(test)]
mod tests {
    use clara_fts::test_util::ngram_index_field;
    use kvenginepb::fts::{
        FtsFieldLayout, FullTextAnalyzerType, FullTextIndexDef, FullTextIndexFieldDef,
    };
    use tantivy::{DocSet, postings::Postings, schema::FieldType};
    use tidb_query_datatype::FieldTypeTp;

    use super::*;
    use crate::table::fts::{
        IntPk, new_tracked_index_meta,
        test_util::{SchemaBuilder, make_exact_match_expr_query, new_write_batch},
        tracked_index_layout_from_tracked_meta,
    };

    #[tokio::test]
    async fn test_delta_cache_builds_ngram_index_from_tracked_field_options() -> Result<()> {
        let table_id = 40;
        let index_id = 6;
        let source_id = 100;
        let index_fields = vec![ngram_index_field(7, 3, 3, "WORD", false, true)];
        let mut index_def = FullTextIndexDef {
            index_id,
            ..Default::default()
        };
        index_def.set_index_fields(index_fields.clone().into());
        let schema = SchemaBuilder::<IntPk>::new(table_id)
            .column(7, FieldTypeTp::String)
            .fts_index(index_def)
            .schema();
        let l0 = new_write_batch()
            .table::<IntPk>(&schema, |row| {
                row.v1(0, 1000, |datum| {
                    datum("Théllo Théllo");
                });
            })
            .finish_as_sst_l0(source_id);
        let tracked_meta = new_tracked_index_meta(
            table_id,
            index_id,
            index_fields,
            FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString,
        );
        let tracked_layout =
            tracked_index_layout_from_tracked_meta(&tracked_meta, &schema.columns)?;
        let spec = FtsDeltaBuildSpec {
            kind: FtsDeltaKind::RowL0,
            source_id,
            table_id,
            index_id,
            input: FtsDeltaInput::RowL0 {
                l0,
                snap_version: SnapVersion::zero(),
            },
        };
        let cache = FtsDeltaCache::new(64 << 20);

        let entry = cache
            .get_or_build(
                spec.clone(),
                schema.clone(),
                &tracked_layout,
                Arc::new(HashMap::new()),
                None,
            )
            .await?
            .unwrap();
        let cached = cache
            .get(FtsDeltaKind::RowL0, source_id, table_id, index_id)
            .unwrap();
        assert!(Arc::ptr_eq(&entry, &cached));
        let reused = cache
            .get_or_build(
                spec,
                schema,
                &tracked_layout,
                Arc::new(HashMap::new()),
                None,
            )
            .await?
            .unwrap();
        assert!(Arc::ptr_eq(&entry, &reused));

        let reader = entry.index_reader();
        let field = tracked_layout.index_layout.first_field().field();
        let tantivy_schema = reader.tantivy_index().schema();
        let FieldType::Str(text_options) = tantivy_schema.get_field_entry(field).field_type()
        else {
            panic!("expected text field");
        };
        assert_eq!(
            text_options.get_indexing_options().unwrap().tokenizer(),
            "NGRAM_V1_FIELD_0"
        );

        let searcher = reader.tantivy_reader().searcher();
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_text(field, "Thé"))?,
            1
        );
        assert_eq!(
            searcher.doc_freq(&tantivy::Term::from_field_text(field, "thé"))?,
            0
        );
        let inverted_index = searcher.segment_readers()[0].inverted_index(field)?;
        let mut postings = inverted_index
            .read_postings(
                &tantivy::Term::from_field_text(field, "Thé"),
                tantivy::schema::IndexRecordOption::WithFreqsAndPositions,
            )?
            .unwrap();
        assert_eq!(postings.doc(), 0);
        assert_eq!(postings.term_freq(), 2);
        let mut positions = Vec::new();
        postings.positions(&mut positions);
        assert_eq!(positions, [0, 7]);

        Ok(())
    }
    #[tokio::test]
    async fn test_delta_cache_builds_native_signed_exact_terms() -> Result<()> {
        let table_id = 41;
        let index_id = 7;
        let source_id = 101;
        let mut index_field = FullTextIndexFieldDef::new();
        index_field.set_col_id(7);
        index_field
            .mut_index_field_option()
            .set_analyzer_type(FullTextAnalyzerType::FtsAnalyzerExactValueV1);
        index_field.mut_index_field_option().set_enable_bm25(false);
        let index_fields = vec![index_field];
        let mut index_def = FullTextIndexDef {
            index_id,
            ..Default::default()
        };
        index_def.set_index_fields(index_fields.clone().into());
        let schema = SchemaBuilder::<IntPk>::new(table_id)
            .column(7, FieldTypeTp::LongLong)
            .fts_index(index_def)
            .schema();
        let l0 = new_write_batch()
            .table::<IntPk>(&schema, |row| {
                row.v1(0, 1000, |datum| datum(7i64));
                row.v1(1, 1000, |datum| datum(-9i64));
            })
            .finish_as_sst_l0(source_id);
        let tracked_meta = new_tracked_index_meta(
            table_id,
            index_id,
            index_fields,
            FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString,
        );
        let tracked_layout =
            tracked_index_layout_from_tracked_meta(&tracked_meta, &schema.columns)?;
        let spec = FtsDeltaBuildSpec {
            kind: FtsDeltaKind::RowL0,
            source_id,
            table_id,
            index_id,
            input: FtsDeltaInput::RowL0 {
                l0,
                snap_version: SnapVersion::zero(),
            },
        };

        let entry = FtsDeltaCache::new(64 << 20)
            .get_or_build(
                spec,
                schema,
                &tracked_layout,
                Arc::new(HashMap::new()),
                None,
            )
            .await?
            .unwrap();
        let field = tracked_layout.index_layout.first_field().field();
        assert!(matches!(
            entry
                .index_reader()
                .tantivy_index()
                .schema()
                .get_field_entry(field)
                .field_type(),
            FieldType::I64(_)
        ));
        assert_eq!(
            entry
                .index_reader()
                .tantivy_reader()
                .searcher()
                .doc_freq(&tantivy::Term::from_field_i64(field, 7))?,
            1
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_delta_cache_matches_empty_exact_text() -> Result<()> {
        let table_id = 42;
        let index_id = 8;
        let source_id = 102;
        let mut index_field = FullTextIndexFieldDef::new();
        index_field.set_col_id(7);
        index_field
            .mut_index_field_option()
            .set_analyzer_type(FullTextAnalyzerType::FtsAnalyzerExactValueV1);
        index_field.mut_index_field_option().set_enable_bm25(false);
        let index_fields = vec![index_field];
        let mut index_def = FullTextIndexDef {
            index_id,
            ..Default::default()
        };
        index_def.set_index_fields(index_fields.clone().into());
        let schema = SchemaBuilder::<IntPk>::new(table_id)
            .column(7, FieldTypeTp::String)
            .fts_index(index_def)
            .schema();
        let l0 = new_write_batch()
            .table::<IntPk>(&schema, |row| {
                row.v1(0, 1000, |datum| datum(""));
                row.v1(1, 1000, |datum| datum("non-empty"));
            })
            .finish_as_sst_l0(source_id);
        let tracked_meta = new_tracked_index_meta(
            table_id,
            index_id,
            index_fields.clone(),
            FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString,
        );
        let tracked_layout =
            tracked_index_layout_from_tracked_meta(&tracked_meta, &schema.columns)?;
        let spec = FtsDeltaBuildSpec {
            kind: FtsDeltaKind::RowL0,
            source_id,
            table_id,
            index_id,
            input: FtsDeltaInput::RowL0 {
                l0,
                snap_version: SnapVersion::zero(),
            },
        };

        let entry = FtsDeltaCache::new(64 << 20)
            .get_or_build(
                spec,
                schema,
                &tracked_layout,
                Arc::new(HashMap::new()),
                None,
            )
            .await?
            .unwrap();
        let query = make_exact_match_expr_query(
            "",
            7,
            index_id,
            &index_fields,
            &tracked_layout.index_layout,
        )?;
        let mut results = Vec::new();
        entry.index_reader().search(&query, &mut results)?;
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].doc_id, 0);
        Ok(())
    }
}