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

use std::{num::NonZeroU64, sync::Arc};

use engine_traits::{CfName, IterOptions, ReadOptions};
use kvengine::{SnapAccess, ValueCache};
use kvproto::kvrpcpb::ExtraOp as TxnExtraOp;
use pd_client::BucketMeta;
use raftstore::store::TxnExt;
use rfstore::Error as RaftServerError;
use txn_types::{Key, Value};

use crate::{self as kv, Error, ErrorInner, Iterator, Snapshot, SnapshotExt};

impl From<RaftServerError> for Error {
    fn from(e: RaftServerError) -> Error {
        Error(Box::new(ErrorInner::Request(e.into())))
    }
}

pub struct RegionSnapshotExt<'a> {
    snapshot: &'a rfstore::store::RegionSnapshot,
}

impl SnapshotExt for RegionSnapshotExt<'_> {
    #[inline]
    fn get_data_version(&self) -> Option<u64> {
        // TiDB coprocessor cache validates cached responses with this value.
        // The applied index advances after Raft applies a write to the region,
        // so equality means the snapshot still sees the same committed state.
        Some(self.snapshot.get_applied_index())
    }

    fn is_max_ts_synced(&self) -> bool {
        self.snapshot
            .txn_ext
            .as_ref()
            .map(|txn_ext| txn_ext.is_max_ts_synced())
            .unwrap_or(false)
    }

    fn get_term(&self) -> Option<NonZeroU64> {
        self.snapshot.term
    }

    fn get_txn_extra_op(&self) -> TxnExtraOp {
        // Check if old_value probe is enabled for testing
        if crate::PROBE_OLD_VALUES_IN_TEST.load(std::sync::atomic::Ordering::Relaxed) {
            TxnExtraOp::ReadOldValue
        } else {
            self.snapshot.txn_extra_op
        }
    }

    fn get_txn_ext(&self) -> Option<&Arc<TxnExt>> {
        self.snapshot.txn_ext.as_ref()
    }

    fn get_buckets(&self) -> Option<Arc<BucketMeta>> {
        self.snapshot.bucket_meta.clone()
    }
}

impl Snapshot for rfstore::store::RegionSnapshot {
    type Iter = rfstore::store::RegionSnapshotIterator;
    type Ext<'a> = RegionSnapshotExt<'a>;

    fn get(&self, _key: &Key) -> kv::Result<Option<Value>> {
        unreachable!()
    }

    fn get_cf(&self, _cf: CfName, _key: &Key) -> kv::Result<Option<Value>> {
        unreachable!()
    }

    fn get_cf_opt(&self, _opts: ReadOptions, _cf: CfName, _key: &Key) -> kv::Result<Option<Value>> {
        unreachable!()
    }

    fn iter(&self, _cf: CfName, _iter_opt: IterOptions) -> kv::Result<Self::Iter> {
        unreachable!()
    }

    #[inline]
    fn lower_bound(&self) -> Option<&[u8]> {
        unreachable!()
    }

    #[inline]
    fn upper_bound(&self) -> Option<&[u8]> {
        unreachable!()
    }

    fn ext(&self) -> RegionSnapshotExt<'_> {
        RegionSnapshotExt { snapshot: self }
    }

    fn get_kvengine_snap(&self) -> Option<&SnapAccess> {
        Some(&self.snap)
    }

    fn get_value_cache(&self) -> Option<&ValueCache> {
        self.value_cache.as_ref()
    }

    fn is_sync(&self) -> bool {
        self.snap.is_sync()
    }
}

#[allow(dead_code)]
impl Iterator for rfstore::store::RegionSnapshotIterator {
    fn next(&mut self) -> kv::Result<bool> {
        unreachable!()
    }

    fn prev(&mut self) -> kv::Result<bool> {
        unreachable!()
    }

    fn seek(&mut self, _key: &Key) -> kv::Result<bool> {
        unreachable!()
    }

    fn seek_for_prev(&mut self, _key: &Key) -> kv::Result<bool> {
        unreachable!()
    }

    fn seek_to_first(&mut self) -> kv::Result<bool> {
        unreachable!()
    }

    fn seek_to_last(&mut self) -> kv::Result<bool> {
        unreachable!()
    }

    fn valid(&self) -> kv::Result<bool> {
        unreachable!()
    }

    fn key(&self) -> &[u8] {
        unreachable!()
    }

    fn value(&self) -> &[u8] {
        unreachable!()
    }
}

#[cfg(test)]
mod tests {
    use kvengine::{WRITE_CF, WriteBatch};
    use kvproto::metapb::Region;
    use tikv_util::mpsc;

    use super::*;

    #[test]
    fn test_rfstore_snapshot_data_version_changes_after_write() {
        let (engine, applier_tx) = kvengine::test_engine::new_test_engine();
        let mut region = Region::default();
        region.set_id(1);

        let old_snap = rfstore::store::RegionSnapshot::from_raw(&engine, &region);
        let old_applied_index = old_snap.get_applied_index();
        assert_eq!(old_snap.ext().get_data_version(), Some(old_applied_index));

        let mut wb = WriteBatch::default();
        wb.reset(region.get_id());
        wb.put(WRITE_CF, b"key", b"value", 0, &[], 1);
        let write_sequence = write_data(wb, &applier_tx);

        let new_snap = rfstore::store::RegionSnapshot::from_raw(&engine, &region);
        assert!(new_snap.get_applied_index() > old_applied_index);
        assert_eq!(new_snap.get_applied_index(), write_sequence);
        assert_eq!(new_snap.ext().get_data_version(), Some(write_sequence));
    }

    fn write_data(
        wb: WriteBatch,
        applier_tx: &mpsc::Sender<kvengine::test_engine::ApplyTask>,
    ) -> u64 {
        let (result_tx, result_rx) = mpsc::bounded(1);
        applier_tx
            .send(kvengine::test_engine::ApplyTask::new_wb(wb, result_tx))
            .unwrap();
        result_rx.recv().unwrap().unwrap()
    }
}