// Copyright 2026 PingCAP, Inc.

use std::sync::Arc;

use async_trait::async_trait;
use bytes::{Buf, Bytes};
use cloud_encryption::EncryptionKey;
use futures::{StreamExt, TryStreamExt, stream};
use tikv_util::sys::SysQuota;

use crate::table::{
    self, Error, Result,
    blobtable::blobtable::BlobTable,
    columnar::{Block, ColumnarReader},
    fts::{
        CommonPk, FtsDeltaCache, FtsIndexReader, IntPk, PkType,
        delta_cache::{FtsDeltaBuildSpec, FtsDeltaSource},
        level::FtsLevels,
    },
    schema_file::Schema,
};

/// Columnar reader that serves FTS index results. It runs the FTS index search
/// once, materializes handle/version (and optional score) into an in-memory
/// block, and then responds to `seek`/`read` from that cached block without
/// performing further I/O.
///
/// # Schema (`output_schema`)
///
/// In production this is `read_row_schema` from
/// [`super::FtsBruteForceCondReader::build_inner_schema`]: the merge/MVCC input
/// schema shared with unindexed fallback readers. It contains:
/// - `handle_column` / `version_column`
/// - one physical column per query source column (from `document_column_ids`),
///   with original table nullability; TableScan project-only columns are
///   excluded
/// - a trailing nullable internal score column (`VIRTUAL_SCORE_COLUMN_ID`)
///
/// Indexed hits fill source columns with ignored placeholders (`push_null`:
/// zero/empty data, plus a null bit only when the column is nullable) and the
/// score column with a non-NULL index score. Unindexed fallback rows use real
/// source values and a NULL score.
pub struct FtsIndexColumnarReader {
    /// Merge/MVCC input schema (`read_row_schema`), not the final TableScan
    /// projection. See the struct-level doc for column layout.
    output_schema: Schema,
    query: clara_fts::Query,
    tracked_index_layout: crate::table::fts::TrackedIndexLayout,
    delta_index_schema: Schema,
    table_id: i64,
    read_ts: u64,
    fts_levels: Arc<FtsLevels>,
    delta_sources: Vec<FtsDeltaSource>,
    delta_build_specs: Vec<FtsDeltaBuildSpec>,
    fts_delta_cache: FtsDeltaCache,
    blob_tables: Arc<std::collections::HashMap<u64, BlobTable>>,
    encryption_key: Option<EncryptionKey>,
    start_handle: Option<Bytes>, // Inclusive
    end_handle: Option<Bytes>,   // Exclusive

    // Cached results
    loaded: bool,
    block: Block,
    idx: usize,
}

impl FtsIndexColumnarReader {
    pub fn new(
        output_schema: Schema,
        query: clara_fts::Query,
        table_id: i64,
        read_ts: u64,
        fts_levels: Arc<FtsLevels>,
        delta_sources: Vec<FtsDeltaSource>,
        delta_build_specs: Vec<FtsDeltaBuildSpec>,
        fts_delta_cache: FtsDeltaCache,
        blob_tables: Arc<std::collections::HashMap<u64, BlobTable>>,
        encryption_key: Option<EncryptionKey>,
        start_handle: Option<Bytes>, // Inclusive
        end_handle: Option<Bytes>,   // Exclusive
    ) -> Result<Self> {
        let delta_index_schema = output_schema.clone();
        Self::new_with_layout(
            output_schema,
            query,
            crate::table::fts::TrackedIndexLayout {
                index_fields: Vec::new(),
                field_layout: kvenginepb::fts::FtsFieldLayout::FtsFieldLayoutLegacyBody,
                index_layout: clara_fts::IndexLayout::legacy_body(),
            },
            delta_index_schema,
            table_id,
            read_ts,
            fts_levels,
            delta_sources,
            delta_build_specs,
            fts_delta_cache,
            blob_tables,
            encryption_key,
            start_handle,
            end_handle,
        )
    }

    pub fn new_with_layout(
        output_schema: Schema,
        query: clara_fts::Query,
        tracked_index_layout: crate::table::fts::TrackedIndexLayout,
        delta_index_schema: Schema,
        table_id: i64,
        read_ts: u64,
        fts_levels: Arc<FtsLevels>,
        delta_sources: Vec<FtsDeltaSource>,
        delta_build_specs: Vec<FtsDeltaBuildSpec>,
        fts_delta_cache: FtsDeltaCache,
        blob_tables: Arc<std::collections::HashMap<u64, BlobTable>>,
        encryption_key: Option<EncryptionKey>,
        // Handle bounds for FTS index search.
        //
        // NOTE: int handles are encoded in *TiDB row-key* format (i64 memcomparable,
        // i.e. `IntPk::encode`), NOT the columnar in-block format (i64 little-endian).
        // Common handles are raw bytes.
        start_handle: Option<Bytes>, // Inclusive
        end_handle: Option<Bytes>,   // Exclusive
    ) -> Result<Self> {
        super::validate_fts_source_score_schema(&output_schema, &query, "FtsIndexColumnarReader")?;
        Ok(Self {
            block: Block::new(&output_schema),
            output_schema,
            query,
            tracked_index_layout,
            delta_index_schema,
            table_id,
            read_ts,
            fts_levels,
            delta_sources,
            delta_build_specs,
            fts_delta_cache,
            blob_tables,
            encryption_key,
            start_handle,
            end_handle,
            loaded: false,
            idx: 0,
        })
    }

    async fn ensure_loaded(&mut self) -> Result<()> {
        if self.loaded {
            return Ok(());
        }

        // Lazily build delta indexes for immutable unindexed sources so the
        // first query pays the build cost only once (singleflight).
        if !self.delta_build_specs.is_empty() {
            let schema = self.delta_index_schema.clone();
            let blob_tables = Arc::clone(&self.blob_tables);
            let encryption_key = self.encryption_key.clone();
            let tracked_index_layout = self.tracked_index_layout.clone();

            let cache = self.fts_delta_cache.clone();
            let specs = std::mem::take(&mut self.delta_build_specs);
            let concurrency = (SysQuota::cpu_cores_quota() as usize).clamp(1, 8);
            let mut built: Vec<FtsDeltaSource> = stream::iter(specs)
                .map(|spec| {
                    let schema = schema.clone();
                    let blob_tables = Arc::clone(&blob_tables);
                    let encryption_key = encryption_key.clone();
                    let cache = cache.clone();
                    let tracked_index_layout = tracked_index_layout.clone();
                    async move {
                        let seq = spec.input.snap_version().into_inner();
                        let entry = cache
                            .get_or_build(
                                spec,
                                schema,
                                &tracked_index_layout,
                                blob_tables,
                                encryption_key,
                            )
                            .await
                            .map_err(|e| Error::Other(e.to_string()))?;
                        Ok::<Option<FtsDeltaSource>, Error>(
                            entry.map(|entry| FtsDeltaSource { entry, seq }),
                        )
                    }
                })
                .buffer_unordered(concurrency)
                .try_filter_map(|x| async move { Ok(x) })
                .try_collect()
                .await?;
            self.delta_sources.append(&mut built);
        }

        let mut hits = if self.output_schema.is_common_handle() {
            FtsIndexReader::<CommonPk>::new_with_layout(
                self.fts_levels.clone(),
                self.table_id,
                self.read_ts,
                self.start_handle.clone(),
                self.end_handle.clone(),
                self.query.clone(),
                self.tracked_index_layout.index_layout.clone(),
            )
            .search(&self.delta_sources)
            .await
            .map_err(|e| Error::Other(e.to_string()))?
        } else {
            FtsIndexReader::<IntPk>::new_with_layout(
                self.fts_levels.clone(),
                self.table_id,
                self.read_ts,
                self.start_handle.clone(),
                self.end_handle.clone(),
                self.query.clone(),
                self.tracked_index_layout.index_layout.clone(),
            )
            .search(&self.delta_sources)
            .await
            .map_err(|e| Error::Other(e.to_string()))?
        };

        self.block.reset();

        if !hits.is_empty() {
            let score_column_idx = self.output_schema.columns.len() - 1;
            // Both common and int handles are encoded in memory-comparable order
            hits.sort_unstable_by(|a, b| a.pk.as_ref().cmp(b.pk.as_ref()));

            // Indexed hits must use the same block layout as unindexed fallback
            // rows so `ColumnarMergeReader` and `ColumnarMvccReader` can merge
            // them. Source columns get placeholders only: `push_null` writes a
            // zero/empty value and, when the column is nullable, a null bit.
            // Placeholders are never evaluated because the non-NULL score marks
            // an indexed row. `FtsBruteForceCondReader` consumes these columns;
            // final projection is loaded by `FtsJoinReader`.
            if self.output_schema.is_common_handle() {
                for hit in hits.iter() {
                    self.block.handles.push_value(&hit.pk);
                    self.block.versions.push_version(hit.version, false);
                    for column in &mut self.block.columns[..score_column_idx] {
                        column.push_null();
                    }
                    let score_bytes = (hit.score as f64).to_le_bytes();
                    self.block.columns[score_column_idx].push_value(&score_bytes);
                }
            } else {
                for hit in hits.iter() {
                    let pk_int: i64 =
                        IntPk::decode(&hit.pk).map_err(|e| Error::Other(e.to_string()))?;
                    self.block.handles.push_value(&pk_int.to_le_bytes());
                    self.block.versions.push_version(hit.version, false);
                    for column in &mut self.block.columns[..score_column_idx] {
                        column.push_null();
                    }
                    let score_bytes = (hit.score as f64).to_le_bytes();
                    self.block.columns[score_column_idx].push_value(&score_bytes);
                }
            }
        }

        self.idx = 0;
        self.loaded = true;
        Ok(())
    }
}

#[async_trait]
impl ColumnarReader for FtsIndexColumnarReader {
    fn schema(&self) -> &Schema {
        &self.output_schema
    }

    async fn seek(&mut self, handle: &[u8]) -> table::Result<()> {
        self.ensure_loaded().await?;
        if handle.is_empty() {
            self.idx = 0;
            return Ok(());
        }
        self.idx = if self.output_schema.is_common_handle() {
            crate::table::search(self.block.length(), |mid| {
                self.block.handles.get_not_null_value(mid) >= handle
            })
        } else {
            let int_handle = (&handle[..]).get_i64_le();
            crate::table::search(self.block.length(), |mid| {
                self.block.handles.get_int_handle_value(mid) >= int_handle
            })
        };
        Ok(())
    }

    async fn read(&mut self, block: &mut Block, limit: usize) -> table::Result<usize> {
        self.ensure_loaded().await?;
        if self.idx >= self.block.length() {
            return Ok(0);
        }
        let read_end = std::cmp::min(self.idx + limit, self.block.length());
        block.append(&self.block, self.idx, read_end);
        let read_rows = read_end - self.idx;
        self.idx = read_end;
        Ok(read_rows)
    }

    fn reset(&mut self) -> table::Result<()> {
        self.idx = 0;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use clara_fts::{
        PhysicalFieldType,
        test_util::{PlainFtsQueryInfo, make_scored_query},
    };
    use kvenginepb::fts::{FtsFieldLayout, FullTextAnalyzerType, FullTextIndexFieldDef};
    use tidb_query_datatype::FieldTypeTp;

    use super::FtsIndexColumnarReader;
    use crate::table::{
        columnar::{Block, ColumnarReader},
        fts::{
            FtsDeltaCache, IntPk,
            level::FtsLevels,
            new_tracked_index_meta,
            reader::FtsIndexReader,
            test_util::{SchemaBuilder, make_exact_match_expr_query, new_block, new_packed},
            tracked_index_layout_from_tracked_meta,
        },
        schema_file::Schema,
    };

    const TABLE_ID: i64 = 808;
    const INDEX_ID: i64 = 9;
    const TEXT_COL_ID: i64 = 133;

    fn construction_error(schema: Schema) -> String {
        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "alpha".into(),
            column_id: TEXT_COL_ID,
            index_id: INDEX_ID,
            ..Default::default()
        });
        match FtsIndexColumnarReader::new(
            schema,
            query,
            TABLE_ID,
            20,
            Arc::new(FtsLevels::default()),
            vec![],
            vec![],
            FtsDeltaCache::disabled(),
            Arc::new(std::collections::HashMap::new()),
            None,
            None,
            None,
        ) {
            Ok(_) => panic!("expected schema validation error"),
            Err(error) => error.to_string(),
        }
    }

    #[test]
    fn rejects_non_string_source_column() {
        let schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column(TEXT_COL_ID, FieldTypeTp::LongLong)
            .fts_score_column(true)
            .schema();
        assert!(
            construction_error(schema).contains("has physical type I64"),
            "non-string FTS source column must be rejected"
        );
    }

    #[tokio::test]
    async fn index_columnar_accepts_not_null_source_and_fills_placeholder() {
        let mut levels = FtsLevels::default();
        levels.track_index(TABLE_ID, INDEX_ID);
        let file = new_packed(1, 5)
            .lp(TABLE_ID, INDEX_ID, |d| {
                d(5, 15, false, "alpha");
            })
            .finish_as_file();
        levels.mut_l0(|files| files.push(file));
        let levels = Arc::new(levels);

        let query = make_scored_query(&PlainFtsQueryInfo {
            query: "alpha".into(),
            column_id: TEXT_COL_ID,
            index_id: INDEX_ID,
            ..Default::default()
        });
        let schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column_not_null(TEXT_COL_ID, FieldTypeTp::String)
            .fts_score_column(true)
            .schema();

        let hits =
            FtsIndexReader::<IntPk>::new(levels.clone(), TABLE_ID, 20, None, None, query.clone())
                .search(&[])
                .await
                .unwrap();
        assert_eq!(hits.len(), 1);

        let mut reader = FtsIndexColumnarReader::new(
            schema.clone(),
            query,
            TABLE_ID,
            20,
            levels,
            vec![],
            vec![],
            FtsDeltaCache::disabled(),
            Arc::new(std::collections::HashMap::new()),
            None,
            None,
            None,
        )
        .unwrap();
        let mut block = Block::new(&schema);
        let rows = reader.read(&mut block, 8).await.unwrap();
        assert_eq!(rows, 1);

        let expected = new_block::<IntPk>(&schema, |row| {
            row(5, 15, false, |datum| {
                datum("");
                datum(hits[0].score as f64);
            });
        });
        assert!(block.eq(&expected));
    }
    #[test]
    fn accepts_exact_i64_source_column_and_uses_resolved_layout() {
        let output_schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column(TEXT_COL_ID, FieldTypeTp::LongLong)
            .fts_score_column(true)
            .schema();
        let delta_schema = SchemaBuilder::<IntPk>::new(TABLE_ID)
            .column(TEXT_COL_ID, FieldTypeTp::LongLong)
            .schema();
        let mut index_field = FullTextIndexFieldDef::new();
        index_field.set_col_id(TEXT_COL_ID);
        index_field
            .mut_index_field_option()
            .set_analyzer_type(FullTextAnalyzerType::FtsAnalyzerExactValueV1);
        index_field.mut_index_field_option().set_enable_bm25(false);
        let tracked_meta = new_tracked_index_meta(
            TABLE_ID,
            INDEX_ID,
            vec![index_field],
            FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString,
        );
        let tracked_layout =
            tracked_index_layout_from_tracked_meta(&tracked_meta, &delta_schema.columns).unwrap();
        let query = make_exact_match_expr_query(
            42,
            TEXT_COL_ID,
            INDEX_ID,
            &tracked_layout.index_fields,
            &tracked_layout.index_layout,
        )
        .unwrap();

        let reader = FtsIndexColumnarReader::new_with_layout(
            output_schema,
            query,
            tracked_layout,
            delta_schema,
            TABLE_ID,
            20,
            Arc::new(FtsLevels::default()),
            vec![],
            vec![],
            FtsDeltaCache::disabled(),
            Arc::new(std::collections::HashMap::new()),
            None,
            None,
            None,
        )
        .unwrap();

        assert_eq!(
            reader
                .tracked_index_layout
                .index_layout
                .first_field()
                .field_type(),
            PhysicalFieldType::I64
        );
    }
}