use std::{sync::Arc, time::Instant};

use anyhow::{Result, ensure};
use tantivy::{DocAddress, SegmentOrdinal};

use crate::reader::{
    frag_reader::{FragReader, ReadRangeScope},
    query_info::{DocField, QueryInfo},
    query_result::{FastFieldColumnCache, SearchRows},
    seg_query_iter::{SegQueryIter, SegmentDocProvider},
    shard_reader::ShardReaderStat,
};

enum FragRowsStreamState {
    Done,
    Scanning {
        scanners: Vec<SegQueryIter>,
        scanner_idx: usize,
    },
}

pub struct FragRowsStream {
    frag_reader: Arc<FragReader>,
    query_info: Arc<QueryInfo>,
    state: FragRowsStreamState,
    subset_of_read_ranges: bool,
    use_fast_fields: bool,
    return_fields: Vec<DocField>,
    fast_field_columns: FastFieldColumnCache,
}

impl FragRowsStream {
    pub(super) fn open(
        frag_reader: Arc<FragReader>,
        query_info: Arc<QueryInfo>,
        stat: &mut ShardReaderStat,
    ) -> Result<Self> {
        let read_scope = frag_reader.classify_read_scope(query_info.read_ranges(), stat);
        if matches!(read_scope, ReadRangeScope::Disjoint) {
            return Ok(Self {
                frag_reader,
                query_info,
                state: FragRowsStreamState::Done,
                subset_of_read_ranges: false,
                use_fast_fields: false,
                return_fields: Vec::new(),
                fast_field_columns: FastFieldColumnCache::default(),
            });
        }

        let subset_of_read_ranges = read_scope.is_subset();
        let use_fast_fields =
            frag_reader.can_use_fast_fields(&query_info, subset_of_read_ranges)?;
        let return_fields = query_info.return_fields()?.to_vec();
        let mut scanners = Vec::with_capacity(frag_reader.segment_readers().len());
        let doc_provider: Arc<dyn SegmentDocProvider> = frag_reader.clone();
        for (i, seg_reader) in frag_reader.segment_readers().iter().enumerate() {
            scanners.push(SegQueryIter::new(
                i as SegmentOrdinal,
                doc_provider.clone(),
                seg_reader.clone(),
                query_info.expr_query(),
            )?);
        }

        Ok(Self {
            frag_reader,
            query_info,
            state: FragRowsStreamState::Scanning {
                scanners,
                scanner_idx: 0,
            },
            subset_of_read_ranges,
            use_fast_fields,
            return_fields,
            fast_field_columns: FastFieldColumnCache::default(),
        })
    }

    pub fn next_batch(
        &mut self,
        batch_size: usize,
        stat: &mut ShardReaderStat,
    ) -> Result<Option<SearchRows>> {
        ensure!(batch_size > 0, "batch_size must be greater than 0");
        if matches!(self.state, FragRowsStreamState::Done) {
            return Ok(None);
        }
        if self.use_fast_fields {
            self.next_fast_field_batch(batch_size, stat)
        } else {
            self.next_document_batch(batch_size, stat)
        }
    }

    fn next_fast_field_batch(
        &mut self,
        batch_size: usize,
        stat: &mut ShardReaderStat,
    ) -> Result<Option<SearchRows>> {
        let batch_start = Instant::now();
        let mut rows = SearchRows::with_capacity(&self.return_fields, batch_size)?;

        while rows.len() < batch_size {
            let Some(doc_addr) = Self::next_scanned_doc_addr(&mut self.state, stat) else {
                break;
            };
            self.fast_field_columns.ensure_segment(
                &self.frag_reader.searcher,
                doc_addr.segment_ord,
                &self.return_fields,
            )?;
            rows.push_fast_fields(doc_addr, &self.return_fields, &self.fast_field_columns)?;
        }
        stat.timing.next_ff_batch_us += batch_start.elapsed().as_micros() as u64;
        if rows.is_empty() {
            return Ok(None);
        }
        Ok(Some(rows))
    }

    fn next_document_batch(
        &mut self,
        batch_size: usize,
        stat: &mut ShardReaderStat,
    ) -> Result<Option<SearchRows>> {
        let return_fields = &self.return_fields;
        let mut rows = SearchRows::with_capacity(return_fields, batch_size)?;

        while rows.len() < batch_size {
            let Some(doc_addr) = Self::next_scanned_doc_addr(&mut self.state, stat) else {
                break;
            };
            let doc = self.frag_reader.fetch_doc(doc_addr)?;
            if !self.subset_of_read_ranges
                && !self
                    .frag_reader
                    .passes_read_range(doc.as_ref(), &self.query_info)?
            {
                stat.docs.read_range_filtered += 1;
                continue;
            }
            rows.push_document(doc.as_ref(), return_fields)?;
        }
        if rows.is_empty() {
            Ok(None)
        } else {
            Ok(Some(rows))
        }
    }

    fn next_scanned_doc_addr(
        state: &mut FragRowsStreamState,
        stat: &mut ShardReaderStat,
    ) -> Option<DocAddress> {
        let FragRowsStreamState::Scanning {
            scanners,
            scanner_idx,
        } = state
        else {
            return None;
        };
        while *scanner_idx < scanners.len() {
            if let Some(doc_addr) = scanners[*scanner_idx].next_doc_addr() {
                stat.docs.tantivy_matched += 1;
                return Some(doc_addr);
            }
            *scanner_idx += 1;
        }
        *state = FragRowsStreamState::Done;
        None
    }
}

#[cfg(test)]
mod tests {
    use anyhow::{Context, Result};

    use crate::{
        reader::{query_info::QueryInfo, query_result::SearchRows},
        test_utils::{
            TEST_HDFSLOGS_INDEX_META, TEST_INDEX_META_NORMAL, TestDoc, build_test_frag,
            collect_versions, deleted_across_segments, load_frag_reader, open_frag_rows_stream,
            row_versions, test_get_inf_range,
        },
        utils::{key::ShardingKeyRange, schema::VERSION_STORE_NAME},
    };

    fn id_tenant_pairs(rows: &SearchRows) -> Result<Vec<(i64, i64)>> {
        let id_column = rows.i64_columns.first().context("id column is missing")?;
        let tenant_id_column = rows
            .i64_columns
            .get(1)
            .context("tenant_id column is missing")?;
        let mut pairs = Vec::with_capacity(rows.len());
        for row_idx in 0..rows.len() {
            pairs.push((
                id_column.value_at(row_idx).context("id is missing")?,
                tenant_id_column
                    .value_at(row_idx)
                    .context("tenant_id is missing")?,
            ));
        }
        Ok(pairs)
    }

    #[test]
    fn test_next_batch_excludes_deleted_docs_across_segments() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_test_frag(
            index_meta,
            vec![
                vec![
                    TestDoc::normal(b"doc_00", 10, "alpha testword"),
                    TestDoc::normal(b"doc_01", 20, "beta testword"),
                    TestDoc::normal(b"doc_02", 30, "gamma testword"),
                ],
                vec![
                    TestDoc::normal(b"doc_10", 40, "delta testword"),
                    TestDoc::normal(b"doc_11", 50, "epsilon testword"),
                ],
            ],
        )?;
        let frag_reader = load_frag_reader(meta, &deleted_across_segments())?;
        let query_info = QueryInfo::query_cols_for_test(
            index_meta,
            frag_reader.schema(),
            vec![test_get_inf_range()],
            vec!["name", VERSION_STORE_NAME],
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;
        let first = stream
            .next_batch(2, &mut stat)?
            .context("first batch is missing")?;
        let second = stream
            .next_batch(2, &mut stat)?
            .context("second batch is missing")?;

        assert_eq!(first.len(), 2);
        assert_eq!(second.len(), 1);
        let mut versions = row_versions(&first)?;
        versions.extend(row_versions(&second)?);
        versions.sort_unstable();
        assert_eq!(versions, vec![10, 30, 50]);
        assert!(stream.next_batch(2, &mut stat)?.is_none());
        assert!(stream.next_batch(2, &mut stat)?.is_none());
        assert_eq!(stat.docs.tantivy_matched, 3);
        Ok(())
    }

    #[test]
    fn test_next_batch_rejects_zero_batch_size() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_test_frag(
            index_meta,
            vec![vec![TestDoc::normal(b"doc_00", 10, "alpha testword")]],
        )?;
        let frag_reader = load_frag_reader(meta, &[])?;
        let query_info = QueryInfo::query_cols_for_test(
            index_meta,
            frag_reader.schema(),
            vec![test_get_inf_range()],
            vec!["name", VERSION_STORE_NAME],
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;

        assert!(stream.next_batch(0, &mut stat).is_err());
        Ok(())
    }

    #[test]
    fn test_next_batch_uses_fast_fields_for_numeric_return_columns() -> Result<()> {
        let index_meta = &*TEST_HDFSLOGS_INDEX_META;
        let (_td, meta) = build_test_frag(
            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 frag_reader = load_frag_reader(meta, &deleted_across_segments())?;
        let query_info = QueryInfo::query_cols_for_test(
            index_meta,
            frag_reader.schema(),
            vec![test_get_inf_range()],
            vec!["id", "tenant_id"],
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;
        let mut returned_pairs = Vec::new();
        let mut batch_lengths = Vec::new();
        while let Some(rows) = stream.next_batch(2, &mut stat)? {
            batch_lengths.push(rows.len());
            returned_pairs.extend(id_tenant_pairs(&rows)?);
        }

        returned_pairs.sort_unstable();
        assert_eq!(batch_lengths, vec![2, 1]);
        assert_eq!(returned_pairs, vec![(101, 1001), (103, 1003), (105, 1005)]);
        assert_eq!(stat.docs.tantivy_matched, 3);
        assert_eq!(stat.docs.read_range_filtered, 0);
        Ok(())
    }

    #[test]
    fn test_next_batch_uses_fast_fields_when_batch_exceeds_total_rows() -> Result<()> {
        let index_meta = &*TEST_HDFSLOGS_INDEX_META;
        let (_td, meta) = build_test_frag(
            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 frag_reader = load_frag_reader(meta, &deleted_across_segments())?;
        let query_info = QueryInfo::query_cols_for_test(
            index_meta,
            frag_reader.schema(),
            vec![test_get_inf_range()],
            vec!["id", "tenant_id"],
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;
        let rows = stream
            .next_batch(10, &mut stat)?
            .context("single batch is missing")?;
        let mut returned_pairs = id_tenant_pairs(&rows)?;

        returned_pairs.sort_unstable();
        assert_eq!(rows.len(), 3);
        assert_eq!(returned_pairs, vec![(101, 1001), (103, 1003), (105, 1005)]);
        assert!(stream.next_batch(10, &mut stat)?.is_none());
        assert_eq!(stat.docs.tantivy_matched, 3);
        assert_eq!(stat.docs.read_range_filtered, 0);
        Ok(())
    }

    #[test]
    fn test_next_batch_returns_none_for_disjoint_read_range() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_test_frag(
            index_meta,
            vec![vec![TestDoc::normal(b"doc_00", 10, "alpha testword")]],
        )?;
        let frag_reader = load_frag_reader(meta, &[])?;
        let read_range = ShardingKeyRange::half_open(b"zzz".to_vec(), b"zzzz".to_vec())?;
        let query_info = QueryInfo::query_cols_for_test(
            index_meta,
            frag_reader.schema(),
            vec![read_range],
            vec!["name", VERSION_STORE_NAME],
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;

        assert!(stream.next_batch(1, &mut stat)?.is_none());
        assert_eq!(stat.frags.disjoint, 1);
        assert_eq!(stat.frags.subset, 0);
        assert_eq!(stat.frags.overlap, 0);
        assert_eq!(stat.docs.tantivy_matched, 0);
        Ok(())
    }

    #[test]
    fn test_next_batch_returns_none_for_empty_subset_match() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_test_frag(
            index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
            ]],
        )?;
        let frag_reader = load_frag_reader(meta, &[])?;
        let read_ranges = vec![test_get_inf_range()];
        let query_info = QueryInfo::query_word(
            index_meta,
            frag_reader.schema(),
            &read_ranges,
            "column_1".to_string(),
            "missing_term".to_string(),
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;

        assert!(stream.next_batch(2, &mut stat)?.is_none());
        assert_eq!(stat.frags.disjoint, 0);
        assert_eq!(stat.frags.subset, 1);
        assert_eq!(stat.frags.overlap, 0);
        assert_eq!(stat.docs.tantivy_matched, 0);
        assert_eq!(stat.docs.read_range_filtered, 0);
        Ok(())
    }

    #[test]
    fn test_next_batch_returns_none_for_empty_overlap_match() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_test_frag(
            index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
            ]],
        )?;
        let frag_reader = load_frag_reader(meta, &[])?;
        let read_range = ShardingKeyRange::half_open(b"doc_01".to_vec(), b"doc_03".to_vec())?;
        let read_ranges = vec![read_range];
        let query_info = QueryInfo::query_word(
            index_meta,
            frag_reader.schema(),
            &read_ranges,
            "column_1".to_string(),
            "missing_term".to_string(),
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;

        assert!(stream.next_batch(2, &mut stat)?.is_none());
        assert_eq!(stat.frags.disjoint, 0);
        assert_eq!(stat.frags.subset, 0);
        assert_eq!(stat.frags.overlap, 1);
        assert_eq!(stat.docs.tantivy_matched, 0);
        assert_eq!(stat.docs.read_range_filtered, 0);
        Ok(())
    }

    #[test]
    fn test_next_batch_filters_overlap_read_range() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_test_frag(
            index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
                TestDoc::normal(b"doc_02", 30, "gamma testword"),
            ]],
        )?;
        let frag_reader = load_frag_reader(meta, &[])?;
        let read_range = ShardingKeyRange::half_open(b"doc_01".to_vec(), b"doc_03".to_vec())?;
        let query_info = QueryInfo::query_cols_for_test(
            index_meta,
            frag_reader.schema(),
            vec![read_range],
            vec!["name", VERSION_STORE_NAME],
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;
        let mut versions = collect_versions(&mut stream, 1, &mut stat)?;

        versions.sort_unstable();
        assert_eq!(versions, vec![20, 30]);
        assert_eq!(stat.frags.overlap, 1);
        assert_eq!(stat.docs.tantivy_matched, 3);
        assert_eq!(stat.docs.read_range_filtered, 1);
        Ok(())
    }

    #[test]
    fn test_next_batch_filtered_doc_does_not_consume_batch_quota() -> Result<()> {
        let index_meta = &*TEST_INDEX_META_NORMAL;
        let (_td, meta) = build_test_frag(
            index_meta,
            vec![vec![
                TestDoc::normal(b"doc_00", 10, "alpha testword"),
                TestDoc::normal(b"doc_01", 20, "beta testword"),
                TestDoc::normal(b"doc_02", 30, "gamma testword"),
            ]],
        )?;
        let frag_reader = load_frag_reader(meta, &[])?;
        let read_range = ShardingKeyRange::half_open(b"doc_01".to_vec(), b"doc_03".to_vec())?;
        let query_info = QueryInfo::query_cols_for_test(
            index_meta,
            frag_reader.schema(),
            vec![read_range],
            vec!["name", VERSION_STORE_NAME],
        )?;

        let (mut stream, mut stat) = open_frag_rows_stream(frag_reader, query_info)?;
        let first = stream
            .next_batch(2, &mut stat)?
            .context("first batch is missing")?;

        assert_eq!(first.len(), 2);
        assert_eq!(row_versions(&first)?, vec![20, 30]);
        assert!(stream.next_batch(2, &mut stat)?.is_none());
        assert_eq!(stat.frags.overlap, 1);
        assert_eq!(stat.docs.tantivy_matched, 3);
        assert_eq!(stat.docs.read_range_filtered, 1);
        Ok(())
    }
}