// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// TODO: Maybe we can find a better place to put these interfaces, e.g. naming
// it as prelude?

//! Batch executor common structures.

use std::sync::Arc;

use async_trait::async_trait;
pub use tidb_query_common::execute_stats::{ExecSummaryCollector, ExecuteStats};
use tidb_query_common::{
    Result, execute_stats::ExecSummaryCollectorEnabled, storage::IntervalRange,
};
use tidb_query_datatype::{codec::batch::LazyBatchColumnVec, expr::EvalWarnings};
pub(crate) use tikv_util::cooperative_checkpoint;
use tipb::FieldType;

/// Combines an `ExecSummaryCollector` with another type. This inner type `T`
/// typically `Executor`/`BatchExecutor`, such that `WithSummaryCollector<C, T>`
/// would implement the same trait and collects the statistics into `C`.
pub struct WithSummaryCollector<C: ExecSummaryCollector, T> {
    pub summary_collector: C,
    pub inner: T,
    #[cfg(test)]
    pub inner_type: &'static str,
}

impl<C: ExecSummaryCollector, T> WithSummaryCollector<C, T> {
    #[cfg(test)]
    pub fn into_inner(self) -> T {
        self.inner
    }
}

/// Applies rate limiting for scan executors on a per-batch basis.
#[async_trait]
pub trait ScanLimiter: Send + Sync {
    /// Called after a batch is produced with the batch's raw KV byte size.
    async fn on_batch(&self, batch_bytes: usize) -> Result<()>;
}

/// Reports coprocessor runtime CPU usage at batch boundaries.
#[async_trait]
pub trait CoprRuntimeBudget: Send + Sync {
    /// Called after one DAG batch finishes execution.
    async fn on_batch(&self, cpu_delta_us: u64) -> Result<()>;
}

/// Shared coprocessor runtime budget handle.
pub type CoprRuntimeBudgetRef = Arc<dyn CoprRuntimeBudget>;

/// The interface for pull-based executors. It is similar to the Volcano
/// Iterator model, but pulls data in batch and stores data by column.
#[async_trait]
pub trait BatchExecutor: Send {
    type StorageStats;

    /// Gets the schema of the output.
    fn schema(&self) -> &[FieldType];

    /// Gets the schema of intermediate result
    fn intermediate_schema(&self, index: usize) -> Result<&[FieldType]>;

    /// Consumes the inner intermediate results and fills the arguments vector.
    fn consume_and_fill_intermediate_results(
        &mut self,
        results: &mut [Vec<BatchExecuteResult>],
    ) -> Result<()>;

    /// Pulls next several rows of data (stored by column).
    ///
    /// This function might return zero rows, which doesn't mean that there is
    /// no more result. See `is_drained` in `BatchExecuteResult`.
    async fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult;

    /// Collects execution statistics (including but not limited to metrics and
    /// execution summaries) accumulated during execution and prepares for
    /// next collection.
    ///
    /// The executor implementation must invoke this function for each children
    /// executor. However the invocation order of children executors is not
    /// stipulated.
    ///
    /// This function may be invoked several times during execution. For each
    /// invocation, it should not contain accumulated meta data in last
    /// invocation. Normally the invocation frequency of this function is
    /// less than `next_batch()`.
    fn collect_exec_stats(&mut self, dest: &mut ExecuteStats);

    /// Returns the cumulative bytes scanned by the underlying storage without
    /// draining execution or storage statistics.
    fn peek_scanned_bytes_sum(&self) -> usize;

    /// Collects underlying storage statistics accumulated during execution and
    /// prepares for next collection.
    ///
    /// Similar to `collect_exec_stats()`, the implementation must invoke this
    /// function for each children executor and this function may be invoked
    /// several times during execution.
    fn collect_storage_stats(&mut self, dest: &mut Self::StorageStats);

    fn take_scanned_range(&mut self) -> IntervalRange;

    fn can_be_cached(&self) -> bool;

    fn collect_summary(
        self,
        output_index: usize,
    ) -> WithSummaryCollector<ExecSummaryCollectorEnabled, Self>
    where
        Self: Sized,
    {
        WithSummaryCollector {
            summary_collector: ExecSummaryCollectorEnabled::new(output_index),
            inner: self,
            #[cfg(test)]
            inner_type: std::any::type_name::<Self>(),
        }
    }
}

#[async_trait]
impl<T: BatchExecutor + ?Sized> BatchExecutor for Box<T> {
    type StorageStats = T::StorageStats;

    fn schema(&self) -> &[FieldType] {
        (**self).schema()
    }

    #[inline]
    fn intermediate_schema(&self, index: usize) -> Result<&[FieldType]> {
        (**self).intermediate_schema(index)
    }

    #[inline]
    fn consume_and_fill_intermediate_results(
        &mut self,
        results: &mut [Vec<BatchExecuteResult>],
    ) -> Result<()> {
        (**self).consume_and_fill_intermediate_results(results)
    }

    async fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult {
        let result = (**self).next_batch(scan_rows).await;
        // Production executor trees connect every parent and child through a
        // boxed executor, making this the common post-next_batch checkpoint.
        cooperative_checkpoint().await;
        result
    }

    fn collect_exec_stats(&mut self, dest: &mut ExecuteStats) {
        (**self).collect_exec_stats(dest);
    }

    fn peek_scanned_bytes_sum(&self) -> usize {
        (**self).peek_scanned_bytes_sum()
    }

    fn collect_storage_stats(&mut self, dest: &mut Self::StorageStats) {
        (**self).collect_storage_stats(dest);
    }

    fn take_scanned_range(&mut self) -> IntervalRange {
        (**self).take_scanned_range()
    }

    fn can_be_cached(&self) -> bool {
        (**self).can_be_cached()
    }
}

#[async_trait]
impl<C: ExecSummaryCollector + Send, T: BatchExecutor> BatchExecutor
    for WithSummaryCollector<C, T>
{
    type StorageStats = T::StorageStats;

    fn schema(&self) -> &[FieldType] {
        self.inner.schema()
    }

    #[inline]
    fn intermediate_schema(&self, index: usize) -> Result<&[FieldType]> {
        self.inner.intermediate_schema(index)
    }

    #[inline]
    fn consume_and_fill_intermediate_results(
        &mut self,
        results: &mut [Vec<BatchExecuteResult>],
    ) -> Result<()> {
        self.inner.consume_and_fill_intermediate_results(results)
    }

    async fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult {
        let timer = self.summary_collector.on_start_iterate();
        let result = self.inner.next_batch(scan_rows).await;
        self.summary_collector
            .on_finish_iterate(timer, result.logical_rows.len());
        result
    }

    fn collect_exec_stats(&mut self, dest: &mut ExecuteStats) {
        self.summary_collector
            .collect(&mut dest.summary_per_executor);
        self.inner.collect_exec_stats(dest);
    }

    fn peek_scanned_bytes_sum(&self) -> usize {
        self.inner.peek_scanned_bytes_sum()
    }

    fn collect_storage_stats(&mut self, dest: &mut Self::StorageStats) {
        self.inner.collect_storage_stats(dest);
    }

    fn take_scanned_range(&mut self) -> IntervalRange {
        self.inner.take_scanned_range()
    }

    fn can_be_cached(&self) -> bool {
        self.inner.can_be_cached()
    }
}

/// Data to be flowed between parent and child executors' single `next_batch()`
/// invocation.
///
/// Note: there are other data flow between executors, like metrics and output
/// statistics. However they are flowed at once, just before response, instead
/// of each step during execution. Hence they are not covered by this structure.
/// See `BatchExecuteMetaData`.
///
/// It is only `Send` but not `Sync` because executor returns its own data copy.
/// However `Send` enables executors to live in different threads.
///
/// It is designed to be used in new generation executors, i.e. executors
/// support batch execution. The old executors will not be refined to return
/// this kind of result.
pub struct BatchExecuteResult {
    /// The *physical* columns data generated during this invocation.
    ///
    /// Note 1: Empty column data doesn't mean that there is no more data. See
    /// `is_drained`.
    ///
    /// Note 2: This is only a *physical* store of data. The data may not be in
    /// desired order and there could be filtered out data stored inside. You
    /// should access the *logical* data via the `logical_rows` field. For the
    /// same reason, `rows_len() > 0` doesn't mean that there is logical data
    /// inside.
    pub physical_columns: LazyBatchColumnVec,

    /// Valid row offsets in `physical_columns`, placed in the logical order.
    pub logical_rows: Vec<usize>,

    /// The warnings generated during this invocation.
    // TODO: It can be more general, e.g. `ExecuteWarnings` instead of `EvalWarnings`.
    // TODO: Should be recorded by row.
    pub warnings: EvalWarnings,

    /// Whether or not there is no more data.
    ///
    /// This structure is a `Result`. When it is:
    /// - `Ok(batch_exec_is_drain)`: See the comment of `BatchExecIsDrain`.
    /// - `Err(_)`: Means that there is an error when trying to retrieve more
    ///   data. In this case, the error is returned and the executor is also
    ///   drained. Similar to `Ok(true)`, there could be some remaining data in
    ///   the `data` field which is valid data and should be processed. The
    ///   caller should NOT call `next_batch()` any more.
    pub is_drained: Result<BatchExecIsDrain>,
}

/// The result of batch execution.
/// - `Drain`: The executor is completely drained and no more data will be
///   returned in the given range.However there could be some (last) data in
///   `data` field this time. The caller should NOT call `next_batch()` any
///   more.
/// - `PagingDrain`: The executor output enough rows of the paging request,
///   there may be following data in the next paging request, the paging request
///   should be returned with scanned range in this case. Only used in paging
///   mode, Also check the last data in `data` field.
/// - `Remain`: The normal case, means that there could be more data. The caller
///   should continue calling `next_batch()` although for each call the returned
///   data may be empty.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BatchExecIsDrain {
    Remain,
    Drain,
    PagingDrain,
}

impl BatchExecIsDrain {
    #[inline]
    pub fn is_remain(&self) -> bool {
        *self == BatchExecIsDrain::Remain
    }

    /// the batch execution need to stop when the result status is Drain or
    /// PagingDrain, but only when we meet Drain, the resultset is really
    /// drained.
    #[inline]
    pub fn stop(&self) -> bool {
        !self.is_remain()
    }
}