use std::sync::Arc;

use anyhow::{Context, Result, anyhow, ensure};
use tantivy::{
    DocAddress, Index, Order, Searcher, SegmentOrdinal, SegmentReader, TantivyDocument,
    collector::{Collector, Count, TopDocs},
    query::EnableScoring,
    schema::{Schema, Type},
};
use tici_grpc::FragMeta;

use crate::{
    reader::{
        doc_cache::{DocCache, DocCacheKey},
        frag_data_cache::{FragDataCache, FragLease},
        query_info::QueryInfo,
        query_result::{ColumnValue, TopKValue},
        seg_query_iter::{SegQueryIter, SegmentDocProvider},
        shard_reader::ShardReaderStat,
    },
    utils::{
        self,
        dedup_bitset::FragAliveBitSetBytes,
        key::{ShardingKey, ShardingKeyRange},
        s3_frag,
        tantivy_doc::DocFieldAccess,
    },
};

mod estimate;
mod stream;

pub use stream::FragRowsStream;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ReadRangeScope {
    Disjoint,
    Subset,
    Overlap,
}

impl ReadRangeScope {
    fn is_subset(self) -> bool {
        matches!(self, ReadRangeScope::Subset)
    }
}

#[derive(Clone)]
pub struct FragReaderMeta {
    pub frag_id: u64,
    pub local_frag_path: String,
    pub min_handle: ShardingKey,
    pub max_handle: ShardingKey,

    _frag_lease: Option<FragLease>,
}

impl FragReaderMeta {
    #[cfg(any(test, feature = "bench", feature = "testexport"))]
    pub fn new(frag_path: &str, min_handle: Vec<u8>, max_handle: Vec<u8>) -> Result<Self> {
        Ok(Self {
            frag_id: s3_frag::parse_frag_id(frag_path)?,
            local_frag_path: frag_path.to_string(),
            min_handle: ShardingKey::from_raw(min_handle),
            max_handle: ShardingKey::from_raw(max_handle),
            _frag_lease: None,
        })
    }

    // The lifecycle of local data is managed by the caller.
    pub fn from_local_path(frag_meta: &FragMeta, local_frag_path: String) -> Result<Self> {
        let property = frag_meta
            .property
            .as_ref()
            .ok_or_else(|| anyhow!("Fragment property is required, meta={:?}", frag_meta))?;
        let frag_id = s3_frag::parse_frag_id(&frag_meta.frag_path)
            .with_context(|| format!("parse frag_id failed, frag_path={}", frag_meta.frag_path))?;
        Ok(Self {
            frag_id,
            local_frag_path,
            min_handle: ShardingKey::from_raw(property.min_handle.clone()),
            max_handle: ShardingKey::from_raw(property.max_handle.clone()),
            _frag_lease: None,
        })
    }

    // The lifecycle of local data is managed by frag_data_cache.
    pub async fn from_frag_data_cache(
        frag_meta: &FragMeta,
        frag_data_cache: Arc<FragDataCache>,
        s3_bucket: &str,
    ) -> Result<Self> {
        let frag_lease = frag_data_cache.acquire_frag(s3_bucket, frag_meta).await?;
        let property = frag_meta
            .property
            .as_ref()
            .ok_or_else(|| anyhow!("Fragment property is required, meta={:?}", frag_meta))?;
        let frag_id = s3_frag::parse_frag_id(&frag_meta.frag_path)
            .with_context(|| format!("parse frag_id failed, frag_path={}", frag_meta.frag_path))?;
        Ok(Self {
            frag_id,
            local_frag_path: frag_lease.local_frag_path().to_string(),
            min_handle: ShardingKey::from_raw(property.min_handle.clone()),
            max_handle: ShardingKey::from_raw(property.max_handle.clone()),
            _frag_lease: Some(frag_lease),
        })
    }

    pub fn load_reader(
        self,
        doc_cache: Option<Arc<DocCache>>,
        valid_range: ShardingKeyRange,
    ) -> Result<FragReader> {
        self.load_reader_with_obsolete_doc_bitset(
            doc_cache,
            valid_range,
            FragAliveBitSetBytes::default(),
        )
    }

    pub fn load_reader_with_obsolete_doc_bitset(
        self,
        doc_cache: Option<Arc<DocCache>>,
        valid_range: ShardingKeyRange,
        obsolete_doc_bitset: FragAliveBitSetBytes,
    ) -> Result<FragReader> {
        // Cache a database block to avoid redundant reads of the same block when
        // performing a full table scan.
        const DOC_STORE_CACHE_NUM_BLOCKS: usize = 1;
        let index_path = utils::frag_index_path(&self.local_frag_path);
        let index = Index::open_in_dir(&index_path)
            .with_context(|| format!("Failed to open index at {}", index_path))?;
        let index_reader = index
            .reader_builder()
            .reload_policy(tantivy::ReloadPolicy::Manual)
            .doc_store_cache_num_blocks(DOC_STORE_CACHE_NUM_BLOCKS)
            .try_into()
            .with_context(|| {
                format!(
                    "Failed to create index reader for frag_path={}",
                    self.local_frag_path
                )
            })?;
        let searcher = index_reader.searcher();
        let custom_segment_readers = build_custom_segment_readers(
            &index,
            &searcher,
            &self.local_frag_path,
            obsolete_doc_bitset,
        )?;
        Ok(FragReader::new(
            self,
            searcher,
            valid_range,
            custom_segment_readers,
            doc_cache,
        ))
    }

    pub fn frag_range(&self) -> Result<ShardingKeyRange> {
        ShardingKeyRange::half_open(
            self.min_handle.raw().to_vec(),
            ShardingKey::next_lexicographic(self.max_handle.raw()),
        )
    }

    pub fn load_index(&self) -> Result<Index> {
        let index_path = utils::frag_index_path(&self.local_frag_path);
        Index::open_in_dir(&index_path)
            .with_context(|| format!("Failed to open index at {}", index_path))
    }
}

pub struct FragReader {
    meta: FragReaderMeta,
    searcher: Searcher,
    valid_range: ShardingKeyRange, // valid_range is shrunk by shard_range
    custom_segment_readers: Vec<Arc<SegmentReader>>,
    doc_cache: Option<Arc<DocCache>>,
}

impl FragReader {
    fn new(
        meta: FragReaderMeta,
        searcher: Searcher,
        valid_range: ShardingKeyRange,
        custom_segment_readers: Vec<Arc<SegmentReader>>,
        doc_cache: Option<Arc<DocCache>>,
    ) -> Self {
        Self {
            meta,
            searcher,
            valid_range,
            custom_segment_readers,
            doc_cache,
        }
    }

    pub fn meta(&self) -> &FragReaderMeta {
        &self.meta
    }

    pub fn index(&self) -> &Index {
        self.searcher.index()
    }

    pub fn frag_path(&self) -> &str {
        &self.meta.local_frag_path
    }

    pub fn schema(&self) -> &Schema {
        self.searcher.schema()
    }

    pub fn live_doc_count(&self) -> u64 {
        self.segment_readers()
            .iter()
            .map(|segment_reader| segment_reader.num_docs() as u64)
            .sum()
    }

    pub fn fetch_doc(&self, addr: DocAddress) -> Result<Arc<TantivyDocument>> {
        let cache_key = self.doc_cache_key(addr)?;
        match &self.doc_cache {
            Some(doc_cache) => {
                doc_cache.get_or_insert_with(cache_key, || self.fetch_doc_uncached(addr))
            }
            None => self.fetch_doc_uncached(addr).map(Arc::new),
        }
    }

    fn fetch_doc_uncached(&self, addr: DocAddress) -> Result<TantivyDocument> {
        let doc = self.searcher.doc(addr).with_context(|| {
            format!(
                "Failed to get doc, addr={:?}, frag_path={}",
                addr,
                self.frag_path()
            )
        })?;
        Ok(doc)
    }

    fn doc_cache_key(&self, addr: DocAddress) -> Result<DocCacheKey> {
        if addr.segment_ord as usize >= self.segment_readers().len() {
            anyhow::bail!(
                "segment ordinal out of range, addr={:?}, frag_path={}",
                addr,
                self.frag_path()
            );
        }
        // Fragment is immutable on reader side, so segment ordinal is stable for the
        // fragment.
        Ok(DocCacheKey::new(self.meta.frag_id, addr))
    }

    fn search_with_collector<C>(
        &self,
        query: &dyn tantivy::query::Query,
        collector: &C,
    ) -> Result<C::Fruit>
    where
        C: Collector,
    {
        let segment_readers = &self.custom_segment_readers;

        let enable_scoring = if collector.requires_scoring() {
            EnableScoring::enabled_from_searcher(&self.searcher)
        } else {
            EnableScoring::disabled_from_searcher(&self.searcher)
        };
        let weight = query.weight(enable_scoring)?;
        let fruits = self.index().search_executor().map(
            |(segment_ord, segment_reader)| {
                collector.collect_segment(
                    weight.as_ref(),
                    segment_ord as u32,
                    segment_reader.as_ref(),
                )
            },
            segment_readers.iter().enumerate(),
        )?;
        Ok(collector.merge_fruits(fruits)?)
    }

    fn segment_readers(&self) -> &[Arc<SegmentReader>] {
        &self.custom_segment_readers
    }

    fn can_use_fast_fields(
        &self,
        query_info: &QueryInfo,
        subset_of_read_ranges: bool,
    ) -> Result<bool> {
        if !subset_of_read_ranges {
            return Ok(false);
        }

        let return_fields = query_info.return_fields()?;
        Ok(return_fields.iter().all(|field| {
            matches!(field.value_type, Type::U64 | Type::I64 | Type::F64)
                && (field.is_bm25_score() || self.schema().get_field_entry(field.field).is_fast())
        }))
    }

    fn passes_read_range(&self, doc: &TantivyDocument, query_info: &QueryInfo) -> Result<bool> {
        let raw_handle = doc
            .get_bytes(query_info.handle_field())?
            .context("handle is None")?;
        Ok(query_info.in_read_ranges(raw_handle))
    }

    pub fn search_with_count_collector(
        &self,
        query_info: &QueryInfo,
        stat: &mut ShardReaderStat,
    ) -> Result<u64> {
        let read_scope = self.classify_read_scope(query_info.read_ranges(), stat);
        if matches!(read_scope, ReadRangeScope::Disjoint) {
            return Ok(0);
        }
        let subset_of_read_ranges = read_scope.is_subset();
        let query = if subset_of_read_ranges {
            query_info.expr_query()
        } else {
            query_info.expr_query_with_read_ranges()?
        };
        let count = self.search_with_collector(query, &Count).with_context(|| {
            format!(
                "Failed to search index reader, frag_path={}",
                self.frag_path()
            )
        })? as u64;
        Ok(count)
    }

    pub fn search_with_top_collector(
        &self,
        query_info: &QueryInfo,
        stat: &mut ShardReaderStat,
    ) -> Result<Vec<TopKValue>> {
        let read_scope = self.classify_read_scope(query_info.read_ranges(), stat);
        if matches!(read_scope, ReadRangeScope::Disjoint) {
            return Ok(Vec::new());
        }
        let subset_read_ranges = read_scope.is_subset();

        let query = if subset_read_ranges {
            query_info.expr_query()
        } else {
            query_info.expr_query_with_read_ranges()?
        };

        let topk_params = query_info.topk_params().context("topk_params is not set")?;
        ensure!(
            topk_params.sort_fields.len() == 1,
            "only single sort field is supported"
        );
        let sf = &topk_params.sort_fields[0];
        let field_entry = self.searcher.schema().get_field_entry(sf.field.field);
        ensure!(field_entry.is_fast(), "sort field is not a fast field");
        let field_name = sf.field.name.clone();
        let order = if sf.is_asc { Order::Asc } else { Order::Desc };

        // $sort_fn:  sorting method name
        // $val_type: return value type (String, f64, i64, u64, etc.)
        // $mapper:   closure to set value to ColumnValue
        macro_rules! search_and_collect {
            ($sort_fn:ident, $val_type:ty, $mapper:expr) => {{
                let collector =
                    TopDocs::with_limit(topk_params.limit).$sort_fn(field_name.clone(), order);
                let top_docs: Vec<($val_type, DocAddress)> = self
                    .search_with_collector(query, &collector)
                    .with_context(|| {
                        format!(
                            "Failed to search index reader, frag_path={}",
                            self.frag_path()
                        )
                    })?;

                let mut topk_vals: Vec<TopKValue> = Vec::with_capacity(top_docs.len());
                for (val, addr) in top_docs {
                    let mut col = ColumnValue::new(field_name.clone());
                    $mapper(&mut col, val);
                    topk_vals.push(TopKValue {
                        cols: vec![col],
                        doc_addr: addr,
                    });
                }
                stat.docs.tantivy_matched += topk_vals.len() as u64;
                Ok(topk_vals)
            }};
        }

        match sf.field.value_type {
            Type::Str => search_and_collect!(
                order_by_string_fast_field,
                String,
                |col: &mut ColumnValue, v: String| col.string_value = Some(v.into_bytes())
            ),
            Type::F64 => {
                search_and_collect!(order_by_fast_field, f64, |col: &mut ColumnValue, v: f64| {
                    col.float_value = Some(v)
                })
            }
            Type::I64 => {
                search_and_collect!(order_by_fast_field, i64, |col: &mut ColumnValue, v: i64| {
                    col.int_value = Some(v)
                })
            }
            Type::U64 => {
                search_and_collect!(order_by_fast_field, u64, |col: &mut ColumnValue, v: u64| {
                    col.int_value = Some(v as i64)
                })
            }
            _ => anyhow::bail!(
                "Unsupported sort field type for top collector: {:?}",
                sf.field.value_type
            ),
        }
    }

    /// Opens a batched row stream over this fragment.
    ///
    /// Rows are produced in `DocAddress` order: ascending `segment_ord`, then
    /// ascending `doc_id` within each segment. Deleted documents are skipped by
    /// the segment scanners before materialization.
    ///
    /// For the returned stream, `next_batch(batch_size, stat)` returns at most
    /// `batch_size` rows. Documents filtered out by read ranges do not count
    /// toward that limit. It returns `None` when the fragment is disjoint from
    /// the read ranges or when all segment scanners are exhausted.
    pub fn open_rows_stream(
        self: Arc<Self>,
        query_info: Arc<QueryInfo>,
        stat: &mut ShardReaderStat,
    ) -> Result<FragRowsStream> {
        FragRowsStream::open(self, query_info, stat)
    }

    pub fn segment_scanners(
        self: &Arc<Self>,
        query_info: &QueryInfo,
        stat: &mut ShardReaderStat,
    ) -> Result<Vec<SegQueryIter>> {
        let read_scope = self.classify_read_scope(query_info.read_ranges(), stat);
        if matches!(read_scope, ReadRangeScope::Disjoint) {
            return Ok(Vec::new());
        }
        let query = query_info.expr_query();
        let mut seg_scanners = Vec::with_capacity(self.segment_readers().len());
        let doc_provider: Arc<dyn SegmentDocProvider> = self.clone();
        for (i, seg_reader) in self.segment_readers().iter().enumerate() {
            let scanner = SegQueryIter::new(
                i as SegmentOrdinal,
                doc_provider.clone(),
                seg_reader.clone(),
                query,
            )?;
            seg_scanners.push(scanner);
        }
        Ok(seg_scanners)
    }

    fn classify_read_scope(
        &self,
        read_ranges: &[ShardingKeyRange],
        stat: &mut ShardReaderStat,
    ) -> ReadRangeScope {
        if self.valid_range.is_disjoint_with_all(read_ranges) {
            stat.frags.disjoint += 1;
            return ReadRangeScope::Disjoint;
        }
        if self.valid_range.is_subset_of_any(read_ranges) {
            stat.frags.subset += 1;
            ReadRangeScope::Subset
        } else {
            stat.frags.overlap += 1;
            ReadRangeScope::Overlap
        }
    }
}

impl SegmentDocProvider for FragReader {
    fn fetch_doc(&self, addr: DocAddress) -> Result<Arc<TantivyDocument>> {
        self.fetch_doc(addr)
    }
}

pub(crate) fn build_custom_segment_readers(
    index: &Index,
    searcher: &Searcher,
    frag_path: &str,
    obsolete_doc_bitset: FragAliveBitSetBytes,
) -> Result<Vec<Arc<SegmentReader>>> {
    if obsolete_doc_bitset.segments.is_empty() {
        return Ok(shared_searcher_segment_readers(searcher));
    }

    let segments = index.searchable_segments().with_context(|| {
        format!(
            "Failed to load searchable segments for frag_path={}",
            frag_path
        )
    })?;
    ensure!(
        segments.len() == obsolete_doc_bitset.segments.len(),
        "segment count mismatch for custom alive bitsets, frag_path={}, segments={}, bitsets={}",
        frag_path,
        segments.len(),
        obsolete_doc_bitset.segments.len()
    );
    obsolete_doc_bitset
        .validate_against_searchable_segments(&segments)
        .with_context(|| {
            format!(
                "custom alive bitset validation failed for frag_path={}",
                frag_path
            )
        })?;
    if !obsolete_doc_bitset.has_deleted_docs() {
        return Ok(shared_searcher_segment_readers(searcher));
    }

    let custom_alive_bitsets = obsolete_doc_bitset.open_alive_bitsets()?;

    let mut custom_segment_readers = Vec::with_capacity(segments.len());
    for (segment, custom_alive_bitset) in segments.iter().zip(custom_alive_bitsets) {
        let segment_reader =
            SegmentReader::open_with_custom_alive_set(segment, custom_alive_bitset).with_context(
                || {
                    format!(
                        "Failed to open custom segment reader, frag_path={}, segment_id={}",
                        frag_path,
                        segment.id()
                    )
                },
            )?;
        custom_segment_readers.push(Arc::new(segment_reader));
    }
    Ok(custom_segment_readers)
}

pub(crate) fn shared_searcher_segment_readers(searcher: &Searcher) -> Vec<Arc<SegmentReader>> {
    searcher
        .segment_readers()
        .iter()
        .cloned()
        .map(Arc::new)
        .collect()
}

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

    use anyhow::{Context, Result};
    use tantivy::{
        DocAddress, Index, doc,
        index::SegmentId,
        schema::{Field, STORED, STRING, Schema, Value},
    };
    use tempfile::TempDir;
    use tici_system::readable::ReadableSize;

    use super::*;
    use crate::{
        reader::{query_info::TantivySchemaExt, query_result::SearchRows},
        test_utils::{
            TEST_HDFSLOGS_INDEX_META, TEST_INDEX_META_NORMAL, TestDoc,
            build_test_frag as build_test_frag_from_docs, create_test_frag_path,
            test_get_inf_range,
        },
        utils::{
            create_index_writer,
            dedup_bitset::FragDocBitSet,
            schema::{HANDLE_STORE_NAME, IndexMeta, SortField, VERSION_STORE_NAME},
            tantivy_doc::DocFieldAccess,
        },
    };

    fn test_doc_cache() -> Arc<DocCache> {
        Arc::new(DocCache::new(ReadableSize::mb(16)))
    }

    fn build_test_index() -> Result<Index> {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_text_field("body", STRING | STORED);
        let schema = schema_builder.build();

        let index = Index::create_in_ram(schema);
        let mut writer = index.writer(15_000_000)?;
        writer.add_document(doc!(field => "a"))?;
        writer.add_document(doc!(field => "b"))?;
        writer.commit()?;
        writer.wait_merging_threads()?;
        Ok(index)
    }

    fn build_test_frag_with_bodies(
        first_body: &str,
        second_body: &str,
    ) -> Result<(TempDir, FragReaderMeta)> {
        let temp_dir = TempDir::new()?;
        let frag_path = create_test_frag_path(temp_dir.path())?;
        let index_path = utils::frag_index_path(&frag_path);
        fs::create_dir_all(&index_path)
            .with_context(|| format!("create_dir_all failed, path={}", index_path))?;

        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_text_field("body", STRING | STORED);
        let schema = schema_builder.build();
        let index = Index::create_in_dir(&index_path, schema)?;
        let mut writer = index.writer(15_000_000)?;
        writer.add_document(doc!(field => first_body))?;
        writer.add_document(doc!(field => second_body))?;
        writer.commit()?;
        writer.wait_merging_threads()?;

        Ok((temp_dir, FragReaderMeta::new(&frag_path, vec![0], vec![1])?))
    }

    fn build_test_frag() -> Result<(TempDir, FragReaderMeta)> {
        build_test_frag_with_bodies("a", "b")
    }

    fn load_frag_reader_with_cache(
        meta: FragReaderMeta,
        doc_cache: Arc<DocCache>,
    ) -> Result<FragReader> {
        let valid_range = meta.frag_range()?;
        meta.load_reader(Some(doc_cache), valid_range)
    }

    fn doc_text(doc: &TantivyDocument, field: Field) -> Result<&str> {
        doc.get_first(field)
            .and_then(|value| value.as_str())
            .context("text field is missing")
    }

    fn doc_version(reader: &FragReader, doc: &TantivyDocument) -> Result<u64> {
        let version_field = reader.schema().doc_field(VERSION_STORE_NAME)?;
        doc.get_u64(&version_field)?.context("version is missing")
    }

    #[test]
    fn test_load_reader_uses_empty_obsolete_doc_bitset_by_default() -> Result<()> {
        let (_temp_dir, frag_reader_meta) = build_test_frag()?;
        let valid_range = frag_reader_meta.frag_range()?;
        let frag_reader = frag_reader_meta.load_reader(Some(test_doc_cache()), valid_range)?;

        assert_eq!(frag_reader.segment_readers().len(), 1);
        assert!(!frag_reader.segment_readers()[0].is_deleted(0));
        assert!(!frag_reader.segment_readers()[0].is_deleted(1));
        Ok(())
    }

    #[test]
    fn test_load_reader_with_obsolete_doc_bitset_applies_deleted_docs() -> Result<()> {
        let (_temp_dir, frag_reader_meta) = build_test_frag()?;
        let index = frag_reader_meta.load_index()?;
        let mut obsolete_doc_bitset = FragDocBitSet::empty_for_index(&index)?;
        obsolete_doc_bitset.push_deleted_doc(DocAddress::new(0, 1))?;
        let valid_range = frag_reader_meta.frag_range()?;

        let frag_reader = frag_reader_meta.load_reader_with_obsolete_doc_bitset(
            Some(test_doc_cache()),
            valid_range,
            FragAliveBitSetBytes::from_doc_bitset(obsolete_doc_bitset)?,
        )?;

        assert_eq!(frag_reader.segment_readers().len(), 1);
        assert!(!frag_reader.segment_readers()[0].is_deleted(0));
        assert!(frag_reader.segment_readers()[0].is_deleted(1));
        Ok(())
    }

    #[test]
    fn test_build_custom_segment_readers_reuses_searcher_readers_without_deleted_docs() -> Result<()>
    {
        let index = build_test_index()?;
        let reader = index.reader()?;
        let searcher = reader.searcher();

        let custom_segment_readers = build_custom_segment_readers(
            &index,
            &searcher,
            "test_frag",
            FragAliveBitSetBytes::from_doc_bitset(FragDocBitSet::empty_for_index(&index)?)?,
        )?;

        assert_eq!(
            custom_segment_readers.len(),
            searcher.segment_readers().len()
        );
        for (custom_segment_reader, segment_reader) in custom_segment_readers
            .iter()
            .zip(searcher.segment_readers().iter())
        {
            assert_eq!(custom_segment_reader.max_doc(), segment_reader.max_doc());
            for doc_id in 0..segment_reader.max_doc() {
                assert_eq!(
                    custom_segment_reader.is_deleted(doc_id),
                    segment_reader.is_deleted(doc_id)
                );
            }
        }
        Ok(())
    }

    #[test]
    fn test_segment_scanners_reuse_frag_reader_segment_readers() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = Arc::new(load_frag_reader_with_cache(meta, test_doc_cache())?);

        assert_eq!(reader.custom_segment_readers.len(), 2);

        let schema = reader.schema().clone();
        let qi = QueryInfo::query_all_for_test(index_meta, &schema, vec![test_get_inf_range()])?;
        let mut stat = ShardReaderStat::default();
        let scanners = reader.segment_scanners(&qi, &mut stat)?;

        assert_eq!(scanners.len(), reader.custom_segment_readers.len());
        for (scanner, segment_reader) in scanners.iter().zip(&reader.custom_segment_readers) {
            assert!(Arc::ptr_eq(
                scanner.segment_reader_for_test(),
                segment_reader
            ));
        }
        Ok(())
    }

    #[test]
    fn test_build_custom_segment_readers_applies_deleted_docs() -> Result<()> {
        let index = build_test_index()?;
        let reader = index.reader()?;
        let searcher = reader.searcher();
        let mut obsolete_doc_bitset = FragDocBitSet::empty_for_index(&index)?;
        obsolete_doc_bitset.push_deleted_doc(DocAddress::new(0, 1))?;

        let custom_segment_readers = build_custom_segment_readers(
            &index,
            &searcher,
            "test_frag",
            FragAliveBitSetBytes::from_doc_bitset(obsolete_doc_bitset)?,
        )?;

        assert_eq!(custom_segment_readers.len(), 1);
        assert!(!searcher.segment_readers()[0].is_deleted(1));
        assert!(!custom_segment_readers[0].is_deleted(0));
        assert!(custom_segment_readers[0].is_deleted(1));
        Ok(())
    }

    #[test]
    fn test_build_custom_segment_readers_rejects_segment_id_mismatch() -> Result<()> {
        let index = build_test_index()?;
        let reader = index.reader()?;
        let searcher = reader.searcher();
        let mut obsolete_doc_bitset =
            FragAliveBitSetBytes::from_doc_bitset(FragDocBitSet::empty_for_index(&index)?)?;
        obsolete_doc_bitset.segments[0].segment_id = SegmentId::generate_random();

        let err = build_custom_segment_readers(&index, &searcher, "test_frag", obsolete_doc_bitset)
            .expect_err("segment id mismatch should be rejected");
        assert!(format!("{err:#}").contains("segment_id"));
        Ok(())
    }

    #[test]
    fn test_doc_cache_separates_segments_in_same_frag() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_temp_dir, frag_reader_meta) = build_multi_seg_frag(index_meta)?;
        let frag_reader = load_frag_reader_with_cache(frag_reader_meta, test_doc_cache())?;

        let seg0_doc = frag_reader.fetch_doc(DocAddress::new(0, 0))?;
        let seg1_doc = frag_reader.fetch_doc(DocAddress::new(1, 0))?;

        assert_eq!(doc_version(&frag_reader, &seg0_doc)?, 10);
        assert_eq!(doc_version(&frag_reader, &seg1_doc)?, 40);
        assert!(!Arc::ptr_eq(&seg0_doc, &seg1_doc));

        assert!(Arc::ptr_eq(
            &seg0_doc,
            &frag_reader.fetch_doc(DocAddress::new(0, 0))?
        ));
        assert!(Arc::ptr_eq(
            &seg1_doc,
            &frag_reader.fetch_doc(DocAddress::new(1, 0))?
        ));
        Ok(())
    }

    #[test]
    fn test_doc_cache_separates_different_frags() -> Result<()> {
        let doc_cache = test_doc_cache();
        let (_temp_dir_1, first_meta) =
            build_test_frag_with_bodies("first_frag_doc_0", "first_frag_doc_1")?;
        let (_temp_dir_2, second_meta) =
            build_test_frag_with_bodies("second_frag_doc_0", "second_frag_doc_1")?;
        let first_reader = load_frag_reader_with_cache(first_meta, doc_cache.clone())?;
        let second_reader = load_frag_reader_with_cache(second_meta, doc_cache)?;
        let first_body_field = first_reader.schema().get_field("body")?;
        let second_body_field = second_reader.schema().get_field("body")?;

        let first_doc = first_reader.fetch_doc(DocAddress::new(0, 0))?;
        let second_doc = second_reader.fetch_doc(DocAddress::new(0, 0))?;

        assert_eq!(doc_text(&first_doc, first_body_field)?, "first_frag_doc_0");
        assert_eq!(
            doc_text(&second_doc, second_body_field)?,
            "second_frag_doc_0"
        );
        assert!(!Arc::ptr_eq(&first_doc, &second_doc));

        assert!(Arc::ptr_eq(
            &first_doc,
            &first_reader.fetch_doc(DocAddress::new(0, 0))?
        ));
        assert!(Arc::ptr_eq(
            &second_doc,
            &second_reader.fetch_doc(DocAddress::new(0, 0))?
        ));
        Ok(())
    }

    #[test]
    fn test_doc_cache_shared_by_readers_for_same_frag() -> Result<()> {
        let doc_cache = test_doc_cache();
        let (_temp_dir, frag_reader_meta) =
            build_test_frag_with_bodies("shared_frag_doc_0", "shared_frag_doc_1")?;
        let first_reader =
            load_frag_reader_with_cache(frag_reader_meta.clone(), doc_cache.clone())?;
        let second_reader = load_frag_reader_with_cache(frag_reader_meta, doc_cache)?;
        let body_field = first_reader.schema().get_field("body")?;

        let first_doc = first_reader.fetch_doc(DocAddress::new(0, 0))?;
        let second_doc = second_reader.fetch_doc(DocAddress::new(0, 0))?;

        assert_eq!(doc_text(&first_doc, body_field)?, "shared_frag_doc_0");
        assert!(Arc::ptr_eq(&first_doc, &second_doc));
        Ok(())
    }

    #[test]
    fn test_fetch_doc_without_doc_cache_returns_equivalent_doc() -> Result<()> {
        let (_temp_dir, frag_reader_meta) =
            build_test_frag_with_bodies("cached_doc_0", "cached_doc_1")?;
        let cached_reader = frag_reader_meta
            .clone()
            .load_reader(Some(test_doc_cache()), frag_reader_meta.frag_range()?)?;
        let uncached_reader =
            frag_reader_meta.load_reader(None, cached_reader.meta().frag_range()?)?;
        let body_field = uncached_reader.schema().get_field("body")?;
        let addr = DocAddress::new(0, 1);

        let cached_doc = cached_reader.fetch_doc(addr)?;
        let uncached_doc = uncached_reader.fetch_doc(addr)?;
        let uncached_doc_again = uncached_reader.fetch_doc(addr)?;

        assert_eq!(doc_text(&cached_doc, body_field)?, "cached_doc_1");
        assert_eq!(doc_text(&uncached_doc, body_field)?, "cached_doc_1");
        assert_eq!(uncached_doc.node_data, cached_doc.node_data);
        assert_eq!(uncached_doc_again.node_data, cached_doc.node_data);
        assert!(!Arc::ptr_eq(&uncached_doc, &uncached_doc_again));
        Ok(())
    }

    // --- Multi-segment query-level tests ---

    fn build_test_tantivy_doc(
        index_meta: &IndexMeta,
        handle: &[u8],
        version: u64,
        name: &str,
        desc: &str,
        pk1: &str,
        pk2: &str,
    ) -> Result<tantivy::TantivyDocument> {
        let schema = index_meta.get_tantivy_schema();
        let mut doc = tantivy::TantivyDocument::default();
        doc.add_bytes(schema.get_field(HANDLE_STORE_NAME)?, handle);
        doc.add_u64(schema.get_field(VERSION_STORE_NAME)?, version);
        doc.add_text(schema.get_field("column_1")?, name);
        doc.add_text(schema.get_field("column_2")?, desc);
        doc.add_text(schema.get_field("column_3")?, pk1);
        doc.add_text(schema.get_field("column_4")?, pk2);
        index_meta.build_fulltext_fields(&mut doc)?;
        Ok(doc)
    }

    /// Two-segment frag (NoMergePolicy):
    ///   Seg 0 — doc 0 (v=10), doc 1 (v=20), doc 2 (v=30)
    ///   Seg 1 — doc 0 (v=40), doc 1 (v=50)
    /// All docs contain "testword" in column_1 for FTS matching.
    fn build_multi_seg_frag(index_meta: &IndexMeta) -> Result<(TempDir, FragReaderMeta)> {
        let temp_dir = TempDir::new()?;
        let frag_path = create_test_frag_path(temp_dir.path())?;
        let index_path = utils::frag_index_path(&frag_path);
        fs::create_dir_all(&index_path)?;

        let schema = index_meta.get_tantivy_schema().clone();
        let mut index = Index::create_in_dir(&index_path, schema)?;
        index.set_tokenizers(index_meta.get_analyzer_manager().clone());
        let mut writer = create_index_writer(&index, 1, 15_000_000, &frag_path)?;

        // Segment 0: 3 docs
        writer.add_document(build_test_tantivy_doc(
            index_meta,
            b"doc_00",
            10,
            "alpha testword",
            "first",
            "p1_00",
            "p2_00",
        )?)?;
        writer.add_document(build_test_tantivy_doc(
            index_meta,
            b"doc_01",
            20,
            "beta testword",
            "second",
            "p1_01",
            "p2_01",
        )?)?;
        writer.add_document(build_test_tantivy_doc(
            index_meta,
            b"doc_02",
            30,
            "gamma testword",
            "third",
            "p1_02",
            "p2_02",
        )?)?;
        writer.commit()?;

        // Segment 1: 2 docs
        writer.add_document(build_test_tantivy_doc(
            index_meta,
            b"doc_10",
            40,
            "delta testword",
            "fourth",
            "p1_10",
            "p2_10",
        )?)?;
        writer.add_document(build_test_tantivy_doc(
            index_meta,
            b"doc_11",
            50,
            "epsilon testword",
            "fifth",
            "p1_11",
            "p2_11",
        )?)?;
        writer.commit()?;
        writer.wait_merging_threads()?;

        Ok((
            temp_dir,
            FragReaderMeta::new(&frag_path, b"doc_00".to_vec(), b"doc_11".to_vec())?,
        ))
    }

    fn load_frag_reader_with_deletions(
        meta: FragReaderMeta,
        deleted_docs: &[DocAddress],
    ) -> Result<FragReader> {
        let index = meta.load_index()?;
        let mut bitset = FragDocBitSet::empty_for_index(&index)?;
        for addr in deleted_docs {
            bitset.push_deleted_doc(*addr)?;
        }
        let valid_range = meta.frag_range()?;
        meta.load_reader_with_obsolete_doc_bitset(
            Some(test_doc_cache()),
            valid_range,
            FragAliveBitSetBytes::from_doc_bitset(bitset)?,
        )
    }

    // Delete seg-0 doc-1 (v=20) and seg-1 doc-0 (v=40).
    // Alive: v=10, v=30, v=50.
    fn deleted_across_segments() -> Vec<DocAddress> {
        vec![DocAddress::new(0, 1), DocAddress::new(1, 0)]
    }

    fn row_versions(rows: &SearchRows) -> Result<Vec<u64>> {
        let version_column = rows
            .u64_column(VERSION_STORE_NAME)
            .context("version column is missing from search result")?;
        (0..rows.len())
            .map(|row_idx| {
                version_column
                    .value_at(row_idx)
                    .context("version is missing")
            })
            .collect()
    }

    #[test]
    fn test_open_rows_stream_excludes_deleted_docs_across_segments() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = Arc::new(load_frag_reader_with_deletions(
            meta,
            &deleted_across_segments(),
        )?);

        let schema = reader.schema().clone();
        let qi = QueryInfo::query_cols_for_test(
            index_meta,
            &schema,
            vec![test_get_inf_range()],
            vec!["name", VERSION_STORE_NAME],
        )?;
        let mut stat = ShardReaderStat::default();
        let mut stream = reader.open_rows_stream(Arc::new(qi), &mut stat)?;
        let rows = stream
            .next_batch(10, &mut stat)?
            .context("row stream should return alive docs")?;

        assert_eq!(rows.len(), 3, "only alive docs should be returned");
        assert!(stream.next_batch(1, &mut stat)?.is_none());
        assert_eq!(stat.docs.tantivy_matched, 3);
        assert_eq!(stat.frags.subset, 1);
        assert_eq!(stat.docs.read_range_filtered, 0);
        let mut versions = row_versions(&rows)?;
        versions.sort_unstable();
        assert_eq!(versions, vec![10, 30, 50]);
        Ok(())
    }

    #[test]
    fn test_open_rows_stream_uses_fast_fields_for_numeric_return_columns() -> Result<()> {
        let index_meta = &*TEST_HDFSLOGS_INDEX_META;
        let (_td, meta) = build_test_frag_from_docs(
            index_meta,
            vec![
                vec![
                    TestDoc::numeric(b"doc_00", 10, 101, 1001),
                    TestDoc::numeric(b"doc_01", 20, 102, 1002),
                    TestDoc::numeric(b"doc_02", 30, 103, 1003),
                ],
                vec![
                    TestDoc::numeric(b"doc_10", 40, 104, 1004),
                    TestDoc::numeric(b"doc_11", 50, 105, 1005),
                ],
            ],
        )?;
        let reader = Arc::new(load_frag_reader_with_deletions(
            meta,
            &deleted_across_segments(),
        )?);

        let schema = reader.schema().clone();
        let qi = QueryInfo::query_cols_for_test(
            index_meta,
            &schema,
            vec![test_get_inf_range()],
            vec!["id", "tenant_id", VERSION_STORE_NAME],
        )?;
        let mut stat = ShardReaderStat::default();
        let mut stream = reader.open_rows_stream(Arc::new(qi), &mut stat)?;
        let rows = stream
            .next_batch(10, &mut stat)?
            .context("row stream should return alive docs")?;

        assert_eq!(rows.len(), 3, "only alive docs should be returned");
        assert!(stream.next_batch(1, &mut stat)?.is_none());
        assert_eq!(stat.docs.tantivy_matched, 3);
        assert_eq!(stat.frags.subset, 1);
        assert_eq!(stat.docs.read_range_filtered, 0);

        let mut versions = row_versions(&rows)?;
        versions.sort_unstable();
        assert_eq!(versions, vec![10, 30, 50]);

        let id_column = rows
            .i64_columns
            .first()
            .context("id column is missing from search result")?;
        let tenant_id_column = rows
            .i64_columns
            .get(1)
            .context("tenant_id column is missing from search result")?;
        let mut returned_pairs: Vec<(i64, i64)> = (0..rows.len())
            .map(|row_idx| -> Result<_> {
                Ok((
                    id_column.value_at(row_idx).context("id is missing")?,
                    tenant_id_column
                        .value_at(row_idx)
                        .context("tenant_id is missing")?,
                ))
            })
            .collect::<Result<_>>()?;
        returned_pairs.sort_unstable();
        assert_eq!(returned_pairs, vec![(101, 1001), (103, 1003), (105, 1005)]);
        Ok(())
    }

    #[test]
    fn test_search_with_count_collector_excludes_deleted_docs_across_segments() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = load_frag_reader_with_deletions(meta, &deleted_across_segments())?;

        let schema = reader.schema().clone();
        let qi = QueryInfo::query_word_count(
            index_meta,
            &schema,
            &[test_get_inf_range()],
            "column_1".to_string(),
            "testword".to_string(),
        )?;
        let mut stat = ShardReaderStat::default();
        let count = reader.search_with_count_collector(&qi, &mut stat)?;

        assert_eq!(count, 3, "deleted docs must not be counted");
        Ok(())
    }

    #[test]
    fn test_search_with_count_collector_updates_frag_scope_stat() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = load_frag_reader_with_deletions(meta, &[])?;

        let schema = reader.schema().clone();
        let qi = QueryInfo::query_all_for_test(index_meta, &schema, vec![test_get_inf_range()])?;
        let mut stat = ShardReaderStat::default();
        let count = reader.search_with_count_collector(&qi, &mut stat)?;

        assert_eq!(count, 5, "multi-seg fixture has 5 alive docs");
        assert_eq!(stat.docs.read_range_filtered, 0);
        assert_eq!(stat.frags.subset, 1);
        Ok(())
    }

    #[test]
    fn test_live_doc_count_respects_custom_alive_bitset() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = load_frag_reader_with_deletions(meta, &deleted_across_segments())?;

        assert_eq!(reader.live_doc_count(), 3);
        Ok(())
    }

    #[test]
    fn test_estimated_live_doc_count_in_ranges_ignores_overlap_read_ranges() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = load_frag_reader_with_deletions(meta, &deleted_across_segments())?;
        let read_ranges = vec![ShardingKeyRange::half_open(
            b"doc_00".to_vec(),
            b"doc_10".to_vec(),
        )?];

        assert_eq!(reader.estimated_live_doc_count_in_ranges(&read_ranges)?, 3);
        Ok(())
    }

    #[test]
    fn test_search_with_top_collector_excludes_deleted_docs_across_segments() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = load_frag_reader_with_deletions(meta, &deleted_across_segments())?;

        let schema = reader.schema().clone();
        let base_qi =
            QueryInfo::query_all_for_test(index_meta, &schema, vec![test_get_inf_range()])?;
        let sort_field = SortField {
            field: schema.doc_field(VERSION_STORE_NAME)?,
            is_asc: false,
        };
        let qi = base_qi.add_topk(&schema, 10, vec![sort_field], false)?;

        let mut stat = ShardReaderStat::default();
        let topk = reader.search_with_top_collector(&qi, &mut stat)?;

        assert_eq!(topk.len(), 3, "deleted docs must be excluded from top-k");
        let versions: Vec<i64> = topk.iter().map(|v| v.cols[0].int_value.unwrap()).collect();
        assert_eq!(versions, vec![50, 30, 10], "descending by version");
        Ok(())
    }

    #[test]
    fn test_segment_scanners_skip_deleted_docs_across_segments() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_multi_seg_frag(index_meta)?;
        let reader = Arc::new(load_frag_reader_with_deletions(
            meta,
            &deleted_across_segments(),
        )?);

        let schema = reader.schema().clone();
        let qi = QueryInfo::query_all_for_test(index_meta, &schema, vec![test_get_inf_range()])?;
        let mut stat = ShardReaderStat::default();
        let mut scanners = reader.segment_scanners(&qi, &mut stat)?;

        assert_eq!(scanners.len(), 2, "one scanner per segment");
        let mut doc_count = 0;
        for scanner in &mut scanners {
            while scanner.next_doc()?.is_some() {
                doc_count += 1;
            }
        }
        assert_eq!(doc_count, 3, "deleted docs must be skipped by scanners");
        Ok(())
    }
}