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

use anyhow::{Result, bail};
use tantivy::schema::Schema;

use crate::{
    reader::{
        FragReaders,
        query_info::{QueryInfo, SearchExecutionPlan, TopKRequest},
        query_result::SearchRows,
    },
    utils::{key::ShardingKeyRange, metrics::begin_shard_reader_query_metrics, schema::IndexMeta},
};

mod estimate;
mod stream;

pub use estimate::{
    DEFAULT_FRAG_SAMPLE_COUNT, DEFAULT_SHARD_SAMPLE_COUNT, EstimateSampleCounts,
    ShardEstimatePopulation, ShardEstimateSample,
};
pub use stream::{DEFAULT_SEARCH_STREAM_DRAIN_BATCH_SIZE, SearchRowsBatch, ShardSearchStream};

#[derive(Clone, Debug, Default)]
pub struct FragScopeStat {
    /// Frags whose `valid_range` is disjoint from all read ranges.
    pub disjoint: usize,
    /// Frags whose `valid_range` is fully covered by a read range.
    pub subset: usize,
    /// Frags that intersect a read range but are not fully covered by any
    /// single range.
    pub overlap: usize,
}

#[derive(Clone, Debug, Default)]
pub struct DocFlowStat {
    /// Docs entering the output pipeline before caller-visible rows.
    ///
    /// Plan-specific semantics (not a comparable global Tantivy hit count):
    /// - `FragScanRows`: every doc returned by a segment scanner (true query
    ///   match count before post-scan read-range filtering).
    /// - `StoreSortTopKRows`: every doc pulled from the k-way merge iterator
    ///   while collecting top-k (candidate count, capped by the limit).
    /// - `FastFieldTopKRows`: sum of per-frag `TopDocs::with_limit` results
    ///   (partial top-k candidate count; can reach `num_frags * limit`, not the
    ///   query's total match population).
    /// - `FragScanCount`: final count returned by the count collector.
    pub tantivy_matched: u64,
    /// Docs dropped by handle read-range filtering after `tantivy_matched`.
    ///
    /// Only incremented on paths that filter read ranges after Tantivy scan
    /// (`FragScanRows`, `StoreSortTopKRows`). On `FastFieldTopKRows`, read
    /// ranges are usually folded into the Tantivy query, so this stays zero.
    pub read_range_filtered: u64,
    /// Rows returned to the caller; only incremented for row plans (see struct
    /// docs).
    pub returned_rows: u64,
}

#[derive(Clone, Debug, Default)]
pub struct QueryTimingBreakdown {
    pub open_us: u64,
    pub next_batch_us: u64,
    /// Total time spent in frag scan `next_batch` on the fast-field path. It is
    /// a subset of next_batch_us.
    pub next_ff_batch_us: u64,
}

#[derive(Clone, Debug, Default)]
pub struct StreamStat {
    pub batch_count: u64,
}

#[derive(Clone, Debug)]
pub struct ShardReaderStat {
    pub plan: SearchExecutionPlan,
    pub frags: FragScopeStat,
    pub docs: DocFlowStat,
    pub timing: QueryTimingBreakdown,
    pub stream: StreamStat,
    pub start: Instant,
}

impl Default for ShardReaderStat {
    fn default() -> Self {
        Self {
            plan: SearchExecutionPlan::FragScanRows,
            frags: FragScopeStat::default(),
            docs: DocFlowStat::default(),
            timing: QueryTimingBreakdown::default(),
            stream: StreamStat::default(),
            start: Instant::now(),
        }
    }
}

// ShardReader serves as a read-only snapshot of a Shard. The current
// implementation simply merges results from multiple fragments without
// deduplication support.
pub struct ShardReader {
    #[allow(dead_code)]
    shard_range: ShardingKeyRange,
    index_meta: IndexMeta,
    _seq: u64,
    estimate_sample_counts: EstimateSampleCounts,
    frag_readers: Arc<FragReaders>,
}

impl ShardReader {
    pub fn new(
        shard_range: ShardingKeyRange,
        index_meta: IndexMeta,
        seq: u64,
        estimate_sample_counts: EstimateSampleCounts,
        frag_readers: Arc<FragReaders>,
    ) -> Self {
        Self {
            shard_range,
            index_meta,
            _seq: seq,
            estimate_sample_counts,
            frag_readers,
        }
    }

    pub fn estimate_sample_counts(&self) -> EstimateSampleCounts {
        self.estimate_sample_counts
    }

    pub fn search_with_query_info(
        &self,
        query_info: impl Into<Arc<QueryInfo>>,
    ) -> Result<(SearchRows, ShardReaderStat)> {
        let query_info: Arc<QueryInfo> = query_info.into();
        query_info.ensure_is_rows_query()?;
        let mut stream = self.open_search_stream(query_info.clone())?;
        let mut merged_rows = query_info.new_search_rows()?;
        loop {
            let SearchRowsBatch { rows, is_eof } =
                stream.next_batch(stream::DEFAULT_SEARCH_STREAM_DRAIN_BATCH_SIZE)?;
            merged_rows.append(rows)?;
            if is_eof {
                break;
            }
        }
        Ok((merged_rows, stream.stat_snapshot()))
    }

    pub fn count_with_query_info(
        &self,
        query_info: impl Into<Arc<QueryInfo>>,
    ) -> Result<(u64, ShardReaderStat)> {
        let query_info: Arc<QueryInfo> = query_info.into();
        query_info.ensure_is_count_query()?;
        let metrics_guard = begin_shard_reader_query_metrics(query_info.execution_plan().as_str());
        let result = self.count_with_query_info_inner(query_info.as_ref());
        metrics_guard.mark_result(&result);
        result
    }

    fn count_with_query_info_inner(
        &self,
        query_info: &QueryInfo,
    ) -> Result<(u64, ShardReaderStat)> {
        let mut stat = ShardReaderStat {
            plan: SearchExecutionPlan::FragScanCount,
            ..Default::default()
        };
        let mut count = 0_u64;
        for frag_reader in self.frag_readers.iter() {
            count += frag_reader.search_with_count_collector(query_info, &mut stat)?;
        }
        stat.docs.tantivy_matched = count;
        Ok((count, stat))
    }

    pub fn parse_expr(
        &self,
        expr: &tipb::Expr,
        return_fields: &[String],
        read_ranges: &[ShardingKeyRange],
        limit: usize,
        sort_field_names: Vec<String>,
        is_asc: Vec<bool>,
    ) -> Result<QueryInfo> {
        if is_asc.len() != sort_field_names.len() {
            bail!(
                "topk sort params length mismatch: sort_field_names.len()={}, is_asc.len()={}",
                sort_field_names.len(),
                is_asc.len()
            );
        }
        let schema = self.get_schema()?;
        QueryInfo::parse_expr(
            &self.index_meta,
            &schema,
            expr,
            read_ranges,
            return_fields,
            TopKRequest {
                limit,
                sort_field_names,
                is_asc,
            },
        )
    }

    pub fn get_schema(&self) -> Result<Schema> {
        Ok(self.index_meta.get_tantivy_schema().clone())
    }
}