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

mod raft_extension;

// #[PerformanceCriticalPath]
use std::{
    fmt::{self, Debug, Display, Formatter},
    sync::Arc,
};

use collections::HashMap;
use kvproto::kvrpcpb::Context;
pub use raft_extension::RaftRouterWrap;
use raftstore::{RegionInfoAccessor, router::RaftStoreRouter};
use tikv_kv::{BTreeEngine, BTreeEngineSnapshot, Modify, OnAppliedCb, TrackerToken};
use txn_types::{TxnExtra, TxnExtraScheduler};

use crate::storage::{
    kv,
    kv::{Engine, SnapContext, WriteData},
};

/// `RaftKv` is a storage engine base on `RaftStore`.
#[derive(Clone)]
pub struct RaftKv<S>
where
    S: RaftStoreRouter + 'static,
{
    router: RaftRouterWrap<S>,
    engine: BTreeEngine,
    txn_extra_scheduler: Option<Arc<dyn TxnExtraScheduler>>,
}

impl<S> RaftKv<S>
where
    S: RaftStoreRouter + 'static,
{
    /// Create a RaftKv using specified configuration.
    pub fn new(router: S, engine: BTreeEngine, _: RegionInfoAccessor) -> RaftKv<S> {
        RaftKv {
            router: RaftRouterWrap::new(router),
            engine,
            txn_extra_scheduler: None,
        }
    }

    pub fn set_txn_extra_scheduler(&mut self, txn_extra_scheduler: Arc<dyn TxnExtraScheduler>) {
        self.txn_extra_scheduler = Some(txn_extra_scheduler);
    }
}

impl<S> Display for RaftKv<S>
where
    S: RaftStoreRouter + 'static,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "RaftKv")
    }
}

impl<S> Debug for RaftKv<S>
where
    S: RaftStoreRouter + 'static,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "RaftKv")
    }
}

impl<S> Engine for RaftKv<S>
where
    S: RaftStoreRouter + 'static,
{
    type Snap = BTreeEngineSnapshot;
    type Local = BTreeEngine;

    fn kv_engine(&self) -> Option<Self::Local> {
        Some(self.engine.clone())
    }

    type RaftExtension = RaftRouterWrap<S>;
    #[inline]
    fn raft_extension(&self) -> &Self::RaftExtension {
        &self.router
    }

    fn modify_on_kv_engine(&self, region_modifies: HashMap<u64, Vec<Modify>>) -> kv::Result<()> {
        self.engine.modify_on_kv_engine(region_modifies)
    }

    fn precheck_write_with_ctx(&self, ctx: &Context) -> kv::Result<()> {
        self.engine.precheck_write_with_ctx(ctx)
    }

    type WriteRes = <BTreeEngine as Engine>::WriteRes;
    fn async_write(
        &self,
        ctx: &Context,
        batch: WriteData,
        subscribed: u8,
        on_applied: Option<OnAppliedCb>,
        tracker: Option<TrackerToken>,
    ) -> Self::WriteRes {
        self.engine
            .async_write(ctx, batch, subscribed, on_applied, tracker)
    }

    type SnapshotRes = <BTreeEngine as Engine>::SnapshotRes;
    fn async_snapshot(&mut self, ctx: SnapContext<'_>) -> Self::SnapshotRes {
        self.engine.async_snapshot(ctx)
    }

    fn schedule_txn_extra(&self, txn_extra: TxnExtra) {
        if let Some(tx) = self.txn_extra_scheduler.as_ref() {
            if !txn_extra.is_empty() {
                tx.schedule(txn_extra);
            }
        }
    }
}