// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

mod storage_impl;

use std::sync::Arc;

use api_version::KvFormat;
use async_trait::async_trait;
use kvproto::{
    coprocessor::{KeyRange, Response},
    kvrpcpb,
    metapb::Region,
};
use protobuf::Message;
use tidb_query_common::{
    execute_stats::ExecSummary,
    storage,
    storage::{
        FindRegionResult, IntervalRange, RegionStorageAccessor, scanner::ScanRescheduleOptions,
    },
};
use tidb_query_executors::interface::CoprRuntimeBudgetRef;
use tikv_alloc::trace::MemoryTraceGuard;
use tikv_util::quota_limiter::QuotaLimiter;
use tipb::{DagRequest, SelectResponse, StreamResponse};

pub use self::storage_impl::TikvStorage;
use crate::{
    coprocessor::{Deadline, RequestHandler, Result, metrics::*, scan_limiter::ScanLimiterRef},
    storage::{Statistics, Store},
};

pub struct DagHandlerBuilder<R, S>
where
    R: RegionStorageAccessor<Storage = S> + 'static,
    S: Store + 'static,
{
    req: DagRequest,
    ranges: Vec<KeyRange>,
    store: S,
    extra_storage_accessor: Option<R>,
    data_version: Option<u64>,
    deadline: Deadline,
    batch_row_limit: usize,
    is_cache_enabled: bool,
    paging_size: Option<u64>,
    paging_size_bytes: Option<u64>,
    max_resp_size: u64,
    quota_limiter: Option<Arc<QuotaLimiter>>,
    reschedule: ScanRescheduleOptions,
    scan_limiter: Option<ScanLimiterRef>,
}

impl<R, S> DagHandlerBuilder<R, S>
where
    R: RegionStorageAccessor<Storage = S>,
    S: Store + 'static,
{
    pub fn new(
        req: DagRequest,
        ranges: Vec<KeyRange>,
        store: S,
        extra_storage_accessor: Option<R>,
        deadline: Deadline,
        batch_row_limit: usize,
        is_cache_enabled: bool,
        paging_size: Option<u64>,
        paging_size_bytes: Option<u64>,
        max_resp_size: u64,
        reschedule: ScanRescheduleOptions,
    ) -> Self {
        DagHandlerBuilder {
            req,
            ranges,
            store,
            extra_storage_accessor,
            data_version: None,
            deadline,
            batch_row_limit,
            is_cache_enabled,
            paging_size,
            paging_size_bytes,
            max_resp_size,
            quota_limiter: None,
            reschedule,
            scan_limiter: None,
        }
    }

    #[must_use]
    pub fn data_version(mut self, data_version: Option<u64>) -> Self {
        self.data_version = data_version;
        self
    }

    /// Sets the optional scan limiter for DAG scan batches.
    #[must_use]
    pub fn scan_limiter(mut self, scan_limiter: Option<ScanLimiterRef>) -> Self {
        self.scan_limiter = scan_limiter;
        self
    }

    /// Sets the optional pre-fairness global quota limiter.
    #[must_use]
    pub fn quota_limiter(mut self, quota_limiter: Option<Arc<QuotaLimiter>>) -> Self {
        self.quota_limiter = quota_limiter;
        self
    }

    pub fn build<F: KvFormat>(self) -> Result<Box<dyn RequestHandler>> {
        COPR_DAG_REQ_COUNT.with_label_values(&["batch"]).inc();
        Ok(BatchDagHandler::new_with_quota::<_, F>(
            self.req,
            self.ranges,
            self.scan_limiter,
            self.store,
            self.extra_storage_accessor,
            self.data_version,
            self.deadline,
            self.is_cache_enabled,
            self.batch_row_limit,
            self.paging_size,
            self.paging_size_bytes,
            self.max_resp_size,
            self.quota_limiter,
            self.reschedule,
        )?
        .into_boxed())
    }
}

/// Wraps the internal accessor to provide the accessor for the secondary
/// TikvStorage.
#[derive(Clone, Debug)]
pub struct ExtraStorageAccessor<R> {
    store_accessor: R,
}

impl<R> ExtraStorageAccessor<R> {
    pub fn from_store_accessor(store_accessor: R) -> Self {
        ExtraStorageAccessor { store_accessor }
    }
}

#[async_trait]
impl<R> RegionStorageAccessor for ExtraStorageAccessor<R>
where
    R: RegionStorageAccessor<Storage: Store>,
{
    type Storage = TikvStorage<R::Storage>;

    /// find the region by the specified key.
    /// The argument `key` should be the comparable format, you should use
    /// `Key::from_raw` encode the raw key.
    async fn find_region_by_key(&self, key: &[u8]) -> FindRegionResult {
        self.store_accessor.find_region_by_key(key).await
    }

    async fn get_local_region_storage(
        &self,
        region: &Region,
        key_ranges: &[KeyRange],
    ) -> storage::Result<Self::Storage> {
        let store = self
            .store_accessor
            .get_local_region_storage(region, key_ranges)
            .await?;
        let check_can_be_cached = store.is_check_has_newer_ts_data();
        Ok(Self::Storage::new(store, check_can_be_cached))
    }

    fn get_original_region_id(&self) -> Option<u64> {
        self.store_accessor.get_original_region_id()
    }
}

pub struct BatchDagHandler {
    runner: tidb_query_executors::runner::BatchExecutorsRunner<Statistics>,
    data_version: Option<u64>,
}

impl BatchDagHandler {
    pub fn new<S: Store + 'static, F: KvFormat>(
        req: DagRequest,
        ranges: Vec<KeyRange>,
        scan_limiter: Option<ScanLimiterRef>,
        store: S,
        extra_storage_accessor: Option<impl RegionStorageAccessor<Storage = S> + 'static>,
        data_version: Option<u64>,
        deadline: Deadline,
        is_cache_enabled: bool,
        streaming_batch_limit: usize,
        paging_size: Option<u64>,
        paging_size_bytes: Option<u64>,
        max_resp_size: u64,
        reschedule: ScanRescheduleOptions,
    ) -> Result<Self> {
        Self::new_with_quota::<S, F>(
            req,
            ranges,
            scan_limiter,
            store,
            extra_storage_accessor,
            data_version,
            deadline,
            is_cache_enabled,
            streaming_batch_limit,
            paging_size,
            paging_size_bytes,
            max_resp_size,
            None,
            reschedule,
        )
    }

    fn new_with_quota<S: Store + 'static, F: KvFormat>(
        req: DagRequest,
        ranges: Vec<KeyRange>,
        scan_limiter: Option<ScanLimiterRef>,
        store: S,
        extra_storage_accessor: Option<impl RegionStorageAccessor<Storage = S> + 'static>,
        data_version: Option<u64>,
        deadline: Deadline,
        is_cache_enabled: bool,
        streaming_batch_limit: usize,
        paging_size: Option<u64>,
        paging_size_bytes: Option<u64>,
        max_resp_size: u64,
        quota_limiter: Option<Arc<QuotaLimiter>>,
        reschedule: ScanRescheduleOptions,
    ) -> Result<Self> {
        let snap = store.get_kvengine_snap();
        let extra_storage_accessor =
            extra_storage_accessor.map(ExtraStorageAccessor::from_store_accessor);
        Ok(Self {
            runner: tidb_query_executors::runner::BatchExecutorsRunner::from_request_with_quota::<
                _,
                F,
            >(
                req,
                ranges,
                scan_limiter,
                TikvStorage::new(store, is_cache_enabled),
                extra_storage_accessor,
                deadline,
                streaming_batch_limit,
                paging_size,
                paging_size_bytes,
                max_resp_size,
                quota_limiter,
                snap,
                reschedule,
            )?,
            data_version,
        })
    }
}

#[async_trait]
impl RequestHandler for BatchDagHandler {
    async fn handle_request(&mut self) -> Result<MemoryTraceGuard<Response>> {
        let result = self.runner.handle_request().await;
        handle_qe_response(result, self.runner.can_be_cached(), self.data_version).map(|x| x.into())
    }

    async fn handle_streaming_request(&mut self) -> Result<(Option<Response>, bool)> {
        handle_qe_stream_response(self.runner.handle_streaming_request().await)
    }

    fn collect_scan_statistics(&mut self, dest: &mut Statistics) {
        self.runner.collect_storage_stats(dest);
    }

    fn collect_scan_summary(&mut self, dest: &mut ExecSummary) {
        self.runner.collect_scan_summary(dest);
    }

    fn set_copr_runtime_budget(&mut self, budget: Option<CoprRuntimeBudgetRef>) -> bool {
        self.runner.set_copr_runtime_budget(budget);
        true
    }

    fn collect_executor_inputs(&self, dest: &mut kvproto::kvrpcpb::ExecutorInputs) {
        self.runner.collect_executor_inputs(dest);
    }

    fn collect_executor_iterations(&self) -> u64 {
        self.runner.collect_executor_iterations()
    }

    fn set_ru_v2_metrics(&mut self, ru_v2: &mut kvrpcpb::Ruv2, resp_bytes: u64) {
        if !self.runner.take_streaming_ru_v2_metrics(ru_v2, resp_bytes) {
            let iterations = self.runner.collect_executor_iterations();
            let mut inputs = kvrpcpb::ExecutorInputs::new();
            self.runner.collect_executor_inputs(&mut inputs);
            ru_v2.set_coprocessor_response_bytes(resp_bytes);
            ru_v2.set_coprocessor_executor_iterations(iterations);
            ru_v2.set_executor_inputs(inputs);
        }
    }
}

fn handle_qe_response(
    result: tidb_query_common::Result<(SelectResponse, Option<IntervalRange>)>,
    can_be_cached: bool,
    data_version: Option<u64>,
) -> Result<Response> {
    use tidb_query_common::error::{ErrorInner, EvaluateError};

    use crate::coprocessor::Error;

    match result {
        Ok((sel_resp, range)) => {
            let mut resp = Response::default();
            if let Some(range) = range {
                resp.mut_range().set_start(range.lower_inclusive);
                resp.mut_range().set_end(range.upper_exclusive);
            }
            let mut data = Vec::with_capacity(sel_resp.compute_size() as usize);
            box_try!(sel_resp.write_to_vec(&mut data));
            resp.set_data(data);
            resp.set_can_be_cached(can_be_cached);
            resp.set_is_cache_hit(false);
            if let Some(v) = data_version {
                resp.set_cache_last_version(v);
            }
            Ok(resp)
        }
        Err(err) => match *err.0 {
            ErrorInner::Storage(err) => Err(err.into()),
            ErrorInner::Evaluate(EvaluateError::DeadlineExceeded) => Err(Error::DeadlineExceeded),
            ErrorInner::Evaluate(EvaluateError::KeyIsLocked(lock_info)) => {
                Err(Error::Locked(lock_info))
            }
            ErrorInner::Evaluate(err) => {
                let mut resp = Response::default();
                let mut sel_resp = SelectResponse::default();
                sel_resp.mut_error().set_code(err.code());
                sel_resp.mut_error().set_msg(err.to_string());
                resp.set_data(box_try!(sel_resp.write_to_bytes()));
                resp.set_can_be_cached(can_be_cached);
                resp.set_is_cache_hit(false);
                Ok(resp)
            }
        },
    }
}

fn handle_qe_stream_response(
    result: tidb_query_common::Result<(Option<(StreamResponse, IntervalRange)>, bool)>,
) -> Result<(Option<Response>, bool)> {
    use tidb_query_common::error::{ErrorInner, EvaluateError};

    use crate::coprocessor::Error;

    match result {
        Ok((Some((s_resp, range)), finished)) => {
            let mut resp = Response::default();
            resp.set_data(box_try!(s_resp.write_to_bytes()));
            resp.mut_range().set_start(range.lower_inclusive);
            resp.mut_range().set_end(range.upper_exclusive);
            Ok((Some(resp), finished))
        }
        Ok((None, finished)) => Ok((None, finished)),
        Err(err) => match *err.0 {
            ErrorInner::Storage(err) => Err(err.into()),
            ErrorInner::Evaluate(EvaluateError::DeadlineExceeded) => Err(Error::DeadlineExceeded),
            ErrorInner::Evaluate(EvaluateError::KeyIsLocked(lock_info)) => {
                Err(Error::Locked(lock_info))
            }
            ErrorInner::Evaluate(err) => {
                let mut resp = Response::default();
                let mut s_resp = StreamResponse::default();
                s_resp.mut_error().set_code(err.code());
                s_resp.mut_error().set_msg(err.to_string());
                resp.set_data(box_try!(s_resp.write_to_bytes()));
                Ok((Some(resp), true))
            }
        },
    }
}

#[cfg(test)]
mod tests {
    use anyhow::anyhow;
    use protobuf::Message;
    use tidb_query_common::error::{Error as CommonError, EvaluateError, StorageError};

    use super::*;
    use crate::coprocessor::Error;

    #[test]
    fn test_handle_qe_response() {
        // Ok Response
        let ok_res = Ok((SelectResponse::default(), None));
        let res = handle_qe_response(ok_res, true, Some(1)).unwrap();
        assert!(res.can_be_cached);
        assert_eq!(res.get_cache_last_version(), 1);
        let mut select_res = SelectResponse::new();
        Message::merge_from_bytes(&mut select_res, res.get_data()).unwrap();
        assert!(!select_res.has_error());

        // Storage Error
        let storage_err = CommonError::from(StorageError(anyhow!("unknown")));
        let res = handle_qe_response(Err(storage_err), false, None);
        assert!(matches!(res, Err(Error::Other(_))));

        // Evaluate Error
        let err = CommonError::from(EvaluateError::DeadlineExceeded);
        let res = handle_qe_response(Err(err), false, None);
        assert!(matches!(res, Err(Error::DeadlineExceeded)));

        let err = CommonError::from(EvaluateError::InvalidCharacterString {
            charset: "test".into(),
        });
        let res = handle_qe_response(Err(err), false, None).unwrap();
        let mut select_res = SelectResponse::new();
        Message::merge_from_bytes(&mut select_res, res.get_data()).unwrap();
        assert_eq!(select_res.get_error().get_code(), 1300);
    }
}