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

use tidb_query_common::storage::{
    IntervalRange, OwnedKvPairEntry, PointRange, Result as QeResult, Storage,
};
use txn_types::Key;

use crate::{
    coprocessor::Error,
    storage::{Scanner, Statistics, Store, mvcc::NewerTsCheckState, txn::nextgen_scanned_bytes},
};

/// A `Storage` implementation over TiKV's storage.
pub struct TikvStorage<S: Store> {
    store: S,
    scanner: Option<S::Scanner>,
    cf_stats_backlog: Statistics,
    met_newer_ts_data_backlog: NewerTsCheckState,
}

impl<S: Store> TikvStorage<S> {
    pub fn new(store: S, check_can_be_cached: bool) -> Self {
        Self {
            store,
            scanner: None,
            cf_stats_backlog: Statistics::default(),
            met_newer_ts_data_backlog: if check_can_be_cached {
                NewerTsCheckState::NotMetYet
            } else {
                NewerTsCheckState::Unknown
            },
        }
    }

    pub fn get_kvengine_snap(&self) -> Option<kvengine::SnapAccess> {
        self.store.get_kvengine_snap()
    }
}

#[maybe_async::async_trait]
impl<S: Store> Storage for TikvStorage<S> {
    type Statistics = Statistics;

    #[maybe_async]
    async fn begin_scan(
        &mut self,
        is_backward_scan: bool,
        is_key_only: bool,
        load_commit_ts: bool,
        range: IntervalRange,
    ) -> QeResult<()> {
        let lower = Some(Key::from_raw(&range.lower_inclusive));
        let upper = Some(Key::from_raw(&range.upper_exclusive));
        if let Some(scanner) = &mut self.scanner {
            self.cf_stats_backlog.add(&scanner.take_statistics());
            if scanner.met_newer_ts_data() == NewerTsCheckState::Met {
                // always override if we met newer ts data
                self.met_newer_ts_data_backlog = NewerTsCheckState::Met;
            }
            if scanner
                .reset_range(
                    is_backward_scan,
                    load_commit_ts,
                    lower.clone(),
                    upper.clone(),
                )
                .map_err(Error::from)?
            {
                return Ok(());
            }
        }

        self.scanner = Some(
            self.store
                .scanner(
                    is_backward_scan,
                    is_key_only,
                    self.met_newer_ts_data_backlog == NewerTsCheckState::NotMetYet,
                    load_commit_ts,
                    lower,
                    upper,
                )
                .await
                .map_err(Error::from)?,
            // There is no transform from storage error to QE's StorageError,
            // so an intermediate error is needed.
        );
        Ok(())
    }

    #[maybe_async]
    async fn scan_next_entry(&mut self) -> QeResult<Option<OwnedKvPairEntry>> {
        // Unwrap is fine because `begin_scan` always initializes the scanner.
        let entry = self
            .scanner
            .as_mut()
            .unwrap()
            .next_entry()
            .await
            .map_err(Error::from)?;
        let entry = match entry {
            Some((key, entry)) => {
                let key = key
                    .into_raw()
                    .map_err(|err| Error::Other(format!("failed to decode scanned key: {err}")))?;
                Some(OwnedKvPairEntry {
                    key,
                    value: entry.value,
                    commit_ts: entry.commit_ts.map(|ts| ts.into_inner()),
                })
            }
            None => None,
        };
        if let Some(entry) = entry.as_ref() {
            let read_ts = self.store.get_read_ts();
            let snap_info = self
                .get_kvengine_snap()
                .map(|snap| (snap.get_keyspace_id(), snap.get_id()));
            txn_debug!(trace_event::types::Category::ReadDetails,
                "TikvStorage::scan_next_entry";
                "read_ts" => read_ts,
                "keyspace_id" => snap_info.map(|(keyspace, _)| keyspace),
                "region_id" => snap_info.map(|(_, region)| region),
                "key" => log_wrappers::Value::key(entry.key.as_slice()),
                "value" => log_wrappers::Value::value(entry.value.as_slice()),
                "commit_ts" => ?entry.commit_ts,
            );
        }
        Ok(entry)
    }

    #[maybe_async]
    async fn get_entry(
        &mut self,
        _is_key_only: bool,
        load_commit_ts: bool,
        range: PointRange,
    ) -> QeResult<Option<OwnedKvPairEntry>> {
        let key = range.0;
        let entry = self
            .store
            .incremental_get_entry(&Key::from_raw(&key), load_commit_ts)
            .await
            .map_err(Error::from)?;
        Ok(entry.map(move |entry| OwnedKvPairEntry {
            key,
            value: entry.value,
            commit_ts: entry.commit_ts.map(|ts| ts.into_inner()),
        }))
    }

    #[inline]
    fn met_uncacheable_data(&self) -> Option<bool> {
        if let Some(scanner) = &self.scanner {
            if scanner.met_newer_ts_data() == NewerTsCheckState::Met {
                return Some(true);
            }
        }
        if self.store.incremental_get_met_newer_ts_data() == NewerTsCheckState::Met {
            return Some(true);
        }
        match self.met_newer_ts_data_backlog {
            NewerTsCheckState::Unknown => None,
            NewerTsCheckState::Met => Some(true),
            NewerTsCheckState::NotMetYet => Some(false),
        }
    }

    fn collect_statistics(&mut self, dest: &mut Statistics) {
        self.cf_stats_backlog
            .add(&self.store.incremental_get_take_statistics());
        if let Some(scanner) = &mut self.scanner {
            self.cf_stats_backlog.add(&scanner.take_statistics());
        }
        dest.add(&self.cf_stats_backlog);
        self.cf_stats_backlog = Statistics::default();
    }

    fn scanned_bytes(&self) -> usize {
        let mut stats = self.cf_stats_backlog.clone();
        stats.add(&self.store.incremental_get_statistics());
        if let Some(scanner) = &self.scanner {
            stats.add(&scanner.statistics());
        }
        // CSE follows NextGen RU accounting here: total versions traversed, not
        // only visible entries returned to the executor.
        nextgen_scanned_bytes(&stats)
    }

    fn get_read_ts(&self) -> u64 {
        self.store.get_read_ts()
    }

    fn is_snapshot_sync(&self) -> bool {
        self.store.is_snapshot_sync()
    }
}