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

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

use crossbeam::atomic::AtomicCell;
use fail::fail_point;
use kvproto::{
    kvrpcpb::ExtraOp as TxnExtraOp,
    metapb,
    raft_cmdpb::{CmdType, RaftCmdRequest, RaftCmdResponse, ReadIndexResponse, Request, Response},
};
use pd_client::BucketMeta;
use raftstore::store::{
    TxnExt,
    util::{LeaseState, RemoteLease},
    worker_metrics::*,
};
use tikv_util::{
    debug, error,
    time::{ThreadReadId, Timespec, monotonic_raw_now},
};
use trace_event::types::TraceContext;
use tracker::GLOBAL_TRACKERS;

use crate::{
    Error, RaftRouter, Result,
    store::{
        Callback, Peer, PeerMsg, RaftCommand, ReadResponse, RegionSnapshot, RequestInspector,
        RequestPolicy, cf_name_to_num, cmd_resp, util,
    },
};

#[derive(Debug)]
pub enum ReadProgress {
    Term(u64),
    AppliedIndexTerm(u64),
    LeaderLease(RemoteLease),
    RegionBuckets(Arc<BucketMeta>),
    Region(Arc<metapb::Region>),
}

impl ReadProgress {
    pub fn term(term: u64) -> ReadProgress {
        ReadProgress::Term(term)
    }

    pub fn applied_index_term(applied_index_term: u64) -> ReadProgress {
        ReadProgress::AppliedIndexTerm(applied_index_term)
    }

    pub fn leader_lease(lease: RemoteLease) -> ReadProgress {
        ReadProgress::LeaderLease(lease)
    }

    pub fn region_buckets(bucket_meta: Arc<BucketMeta>) -> ReadProgress {
        ReadProgress::RegionBuckets(bucket_meta)
    }

    pub fn region(region: metapb::Region) -> ReadProgress {
        ReadProgress::Region(Arc::new(region))
    }
}

pub trait ReadExecutor {
    fn get_snapshot(&self, region_id: u64, region_ver: u64) -> Result<RegionSnapshot>;
    fn get_value(&self, req: &Request, region: &metapb::Region) -> Result<Response> {
        let key = req.get_get().get_key();
        // region key range has no data prefix, so we must use origin key to check.
        util::check_key_in_region(key, region)?;

        let region_snap =
            self.get_snapshot(region.get_id(), region.get_region_epoch().get_version())?;
        let cf_num = cf_name_to_num(req.get_get().get_cf());
        let item = region_snap.snap.get(cf_num, key, u64::MAX);
        let mut resp = Response::default();
        if item.value_len() > 0 {
            resp.mut_get().set_value(item.get_value().to_vec());
        }
        Ok(resp)
    }

    fn execute(
        &self,
        msg: &RaftCmdRequest,
        region: &Arc<metapb::Region>,
        read_index: Option<u64>,
        _ts: Option<ThreadReadId>,
    ) -> ReadResponse {
        let requests = msg.get_requests();
        let mut response = ReadResponse {
            response: RaftCmdResponse::default(),
            snapshot: None,
            txn_extra_op: TxnExtraOp::Noop,
        };
        let mut responses = Vec::with_capacity(requests.len());
        for req in requests {
            let cmd_type = req.get_cmd_type();
            let mut resp = match cmd_type {
                CmdType::Get => match self.get_value(req, region.as_ref()) {
                    Ok(resp) => resp,
                    Err(e) => {
                        error!(?e;
                            "failed to execute get command";
                            "region_id" => region.get_id(),
                        );
                        response.response = cmd_resp::new_error(e);
                        return response;
                    }
                },
                CmdType::Snap => {
                    match self
                        .get_snapshot(region.get_id(), region.get_region_epoch().get_version())
                    {
                        Ok(snapshot) => {
                            response.snapshot = Some(snapshot);
                            Response::default()
                        }
                        Err(e) => {
                            response.response = cmd_resp::new_error(e);
                            return response;
                        }
                    }
                }
                CmdType::ReadIndex => {
                    let mut resp = Response::default();
                    if let Some(read_index) = read_index {
                        let mut res = ReadIndexResponse::default();
                        res.set_read_index(read_index);
                        resp.set_read_index(res);
                    } else {
                        panic!("[region {}] can not get readindex", region.get_id());
                    }
                    resp
                }
                CmdType::Prewrite
                | CmdType::Put
                | CmdType::Delete
                | CmdType::DeleteRange
                | CmdType::IngestSst
                | CmdType::Invalid => unreachable!(),
            };
            resp.set_cmd_type(cmd_type);
            responses.push(resp);
        }
        response.response.set_responses(responses.into());
        response
    }
}

/// A read only delegate of `Peer`.
#[derive(Clone, Debug)]
pub struct ReadDelegate {
    pub region: Arc<metapb::Region>,
    pub peer_id: u64,
    pub store_id: u64,
    pub term: u64,
    pub applied_index_term: u64,
    pub leader_lease: Option<RemoteLease>,
    pub last_valid_ts: Timespec,

    pub txn_ext: Arc<TxnExt>,
    pub bucket_meta: Option<Arc<BucketMeta>>,
    pub txn_extra_op: Arc<AtomicCell<TxnExtraOp>>,
}

impl ReadDelegate {
    pub(crate) fn from_peer(peer: &Peer) -> Arc<ReadDelegate> {
        let region = peer.region().clone();
        let peer_id = peer.peer.get_id();
        let store_id = peer.peer.get_store_id();
        Arc::new(ReadDelegate {
            region: Arc::new(region),
            peer_id,
            store_id,
            term: peer.term(),
            applied_index_term: peer.get_store().applied_index_term(),
            leader_lease: None,
            last_valid_ts: Timespec::new(0, 0),
            txn_ext: peer.txn_ext.clone(),
            bucket_meta: peer.buckets.as_ref().map(|b| b.meta.clone()),
            txn_extra_op: peer.txn_extra_op.clone(),
        })
    }

    pub(crate) fn clone_for_update(&self) -> ReadDelegate {
        self.clone()
    }

    fn fresh_valid_ts(&mut self) {
        self.last_valid_ts = monotonic_raw_now();
    }

    pub fn update(&mut self, progress: ReadProgress) {
        self.fresh_valid_ts();
        match progress {
            ReadProgress::Term(term) => {
                self.term = term;
            }
            ReadProgress::AppliedIndexTerm(applied_index_term) => {
                self.applied_index_term = applied_index_term;
            }
            ReadProgress::LeaderLease(leader_lease) => {
                self.leader_lease = Some(leader_lease);
            }
            ReadProgress::RegionBuckets(bucket_meta) => {
                self.bucket_meta = Some(bucket_meta);
            }
            ReadProgress::Region(region) => {
                self.region = region;
            }
        }
    }

    fn is_in_leader_lease(&self, ts: Timespec) -> bool {
        if let Some(ref lease) = self.leader_lease {
            let term = lease.term();
            if term == self.term {
                if lease.inspect(Some(ts)) == LeaseState::Valid {
                    return true;
                } else {
                    TLS_LOCAL_READ_METRICS
                        .with(|m| m.borrow_mut().reject_reason.lease_expire.inc());
                }
            } else {
                TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.term_mismatch.inc());
            }
        }
        false
    }

    /// Used in some external tests.
    pub fn mock(region_id: u64) -> Self {
        let mut region: metapb::Region = Default::default();
        region.set_id(region_id);
        ReadDelegate {
            region: Arc::new(region),
            peer_id: 1,
            store_id: 1,
            term: 1,
            applied_index_term: 1,
            leader_lease: None,
            last_valid_ts: Timespec::new(0, 0),
            txn_extra_op: Default::default(),
            txn_ext: Default::default(),
            bucket_meta: None,
        }
    }
}

#[derive(Clone)]
pub struct LocalReader {
    pub inner: Arc<LocalReaderInner>,
    // A channel to raftstore.
    pub router: RaftRouter,
}

pub struct LocalReaderInner {
    store_readers: Arc<papaya::HashMap<u64, Arc<ReadDelegate>>>,
    pub kv_engine: kvengine::Engine,
}

impl ops::Deref for LocalReader {
    type Target = LocalReaderInner;

    fn deref(&self) -> &LocalReaderInner {
        &self.inner
    }
}

impl ReadExecutor for LocalReader {
    fn get_snapshot(&self, region_id: u64, region_ver: u64) -> Result<RegionSnapshot> {
        if let Some(snap) = self.kv_engine.try_get_snap_access(region_id, region_ver)? {
            // LocalReader bypasses raftstore and returns the snapshot directly.
            // Populate `term` from the cached delegate so storage can do term-based checks
            // (e.g. in-memory pessimistic lock table) deterministically.
            let (extra_op, term) = self
                .get_delegate(region_id)
                .as_ref()
                .map(|d| (d.txn_extra_op.load(), NonZeroU64::new(d.term)))
                .unwrap_or((TxnExtraOp::Noop, None));
            let mut snapshot =
                RegionSnapshot::from_snapshot(snap, self.kv_engine.get_value_cache(), None);
            snapshot.txn_extra_op = extra_op;
            snapshot.term = term;
            return Ok(snapshot);
        }
        Err(Error::StaleCommand)
    }
}

impl LocalReader {
    pub fn new(
        kv_engine: kvengine::Engine,
        store_readers: Arc<papaya::HashMap<u64, Arc<ReadDelegate>>>,
        router: RaftRouter,
    ) -> Self {
        LocalReader {
            inner: Arc::new(LocalReaderInner {
                store_readers,
                kv_engine,
            }),
            router,
        }
    }

    fn redirect(&self, cmd: RaftCommand) {
        debug!("localreader redirects command"; "command" => ?cmd);
        let region_id = cmd.request.get_header().get_region_id();
        self.router.send(region_id, PeerMsg::RaftCommand(cmd));
    }

    // Ideally `get_delegate` should return `Option<&ReadDelegate>`, but if so the
    // lifetime of the returned `&ReadDelegate` will bind to `self`, and make it
    // impossible to use `&mut self` while the `&ReadDelegate` is alive, a
    // better choice is use `Rc` but `LocalReader: Send` will be violated, which
    // is required by `LocalReadRouter: Send`, use `Arc` will introduce extra cost
    // but make the logic clear
    fn get_delegate(&self, region_id: u64) -> Option<Arc<ReadDelegate>> {
        let readers = self.store_readers.pin();
        readers.get(&region_id).cloned()
    }

    pub fn get_txn_ext(&self, region_id: u64) -> Option<Arc<TxnExt>> {
        self.get_delegate(region_id)
            .map(|delegate| delegate.txn_ext.clone())
    }

    fn pre_propose_raft_command(
        &self,
        req: &RaftCmdRequest,
    ) -> Result<Option<(Arc<ReadDelegate>, RequestPolicy)>> {
        // Check region id.
        let region_id = req.get_header().get_region_id();
        let delegate = match self.get_delegate(region_id) {
            Some(d) => d,
            None => {
                TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.no_region.inc());
                debug!("rejected by no region"; "region_id" => region_id);
                return Ok(None);
            }
        };

        if let Err(e) = util::check_store_id(req, delegate.store_id) {
            TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.store_id_mismatch.inc());
            debug!("rejected by store id not match"; "err" => %e);
            return Err(e);
        }

        fail_point!("localreader_on_find_delegate");

        // Check peer id.
        if let Err(e) = util::check_peer_id(req, delegate.peer_id) {
            TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.peer_id_mismatch.inc());
            return Err(e);
        }

        // Check term.
        if let Err(e) = util::check_term(req, delegate.term) {
            debug!(
                "check term";
                "delegate_term" => delegate.term,
                "header_term" => req.get_header().get_term(),
            );
            TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.term_mismatch.inc());
            return Err(e);
        }

        // Check region epoch.
        if util::check_region_epoch(req, &delegate.region, false).is_err() {
            TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.epoch.inc());
            // Stale epoch, redirect it to raftstore to get the latest region.
            let delegate_ver = delegate.region.get_region_epoch().get_version();
            debug!("rejected by epoch not match req {:?}", req; "delegate_ver" => delegate_ver);
            return Ok(None);
        }

        let mut inspector = Inspector {
            delegate: &delegate,
        };
        match inspector.inspect(req) {
            RequestPolicy::ReadLocal => Ok(Some((delegate, RequestPolicy::ReadLocal))),
            // It can not handle other policies.
            _ => Ok(None),
        }
    }

    pub fn propose_raft_command(
        &self,
        mut read_id: Option<ThreadReadId>,
        req: RaftCmdRequest,
        trace_ctx: TraceContext,
        cb: Callback,
    ) {
        match self.pre_propose_raft_command(&req) {
            Ok(Some((delegate, policy))) => {
                let mut response = match policy {
                    // Leader can read local if and only if it is in lease.
                    RequestPolicy::ReadLocal => {
                        let snapshot_ts = match read_id.as_mut() {
                            // If this peer became Leader not long ago and just after the cached
                            // snapshot was created, this snapshot can not see all data of the peer.
                            Some(id) => {
                                if id.create_time <= delegate.last_valid_ts {
                                    id.create_time = monotonic_raw_now();
                                }
                                id.create_time
                            }
                            None => monotonic_raw_now(),
                        };
                        if !delegate.is_in_leader_lease(snapshot_ts) {
                            // Forward to raftstore.
                            self.redirect(RaftCommand::new(req, trace_ctx, cb));
                            return;
                        }
                        self.execute(&req, &delegate.region, None, read_id)
                    }
                    _ => unreachable!(),
                };
                cb.read_tracker().map(|tracker| {
                    GLOBAL_TRACKERS.with_tracker(tracker, |t| {
                        t.metrics.local_read = true;
                    })
                });

                response.txn_extra_op = delegate.txn_extra_op.load();
                cmd_resp::bind_term(&mut response.response, delegate.term);
                if let Some(snap) = response.snapshot.as_mut() {
                    snap.txn_ext = Some(delegate.txn_ext.clone());
                    snap.term = NonZeroU64::new(delegate.term);
                    snap.txn_extra_op = response.txn_extra_op;
                    snap.bucket_meta = delegate.bucket_meta.clone();
                }
                cb.invoke_read(response);
            }
            // Forward to raftstore.
            Ok(None) => self.redirect(RaftCommand::new(req, trace_ctx, cb)),
            Err(e) => {
                let mut response = cmd_resp::new_error(e);
                if let Some(delegate) = self.get_delegate(req.get_header().get_region_id()) {
                    cmd_resp::bind_term(&mut response, delegate.term);
                }
                cb.invoke_read(ReadResponse {
                    response,
                    snapshot: None,
                    txn_extra_op: TxnExtraOp::Noop,
                });
            }
        }
    }

    /// If read requests are received at the same RPC request, we can create one
    /// snapshot for all of them and check whether the time when the
    /// snapshot was created is in lease. We use ThreadReadId to figure out
    /// whether this RaftCommand comes from the same RPC request with
    /// the last RaftCommand which left a snapshot cached in LocalReader.
    /// ThreadReadId is composed by thread_id and a thread_local incremental
    /// sequence.
    #[inline]
    pub fn read(&self, read_id: Option<ThreadReadId>, req: RaftCmdRequest, cb: Callback) {
        self.propose_raft_command(read_id, req, TraceContext::default(), cb);
        maybe_tls_local_read_metrics_flush();
    }
}

#[cfg(all(test, feature = "testexport"))]
mod tests {
    use std::{num::NonZeroU64, sync::Arc};

    use tikv_util::mpsc;

    use super::{LocalReader, ReadDelegate, ReadExecutor};
    use crate::RaftRouter;

    #[test]
    fn test_localreader_direct_snapshot_sets_term() {
        let (test_engine, _applier_tx) = kvengine::test_engine::new_test_engine();

        let region_id = 1;
        let region_ver = test_engine
            .engine
            .get_snap_access(region_id)
            .expect("test shard must exist")
            .get_version();

        let store_readers = Arc::new(papaya::HashMap::new());
        let mut delegate = ReadDelegate::mock(region_id);
        delegate.term = 7;
        store_readers.pin().insert(region_id, Arc::new(delegate));

        let (peer_sender, _peer_receiver) = crate::raft_router_peer_channel(0);
        let (store_sender, _store_receiver) = mpsc::unbounded();
        let router = RaftRouter::new(peer_sender, store_sender);

        let reader = LocalReader::new(test_engine.engine.clone(), store_readers, router);

        let snap = reader
            .get_snapshot(region_id, region_ver)
            .expect("snapshot should be available");
        assert_eq!(snap.term, NonZeroU64::new(7));
    }
}

struct Inspector<'r> {
    delegate: &'r ReadDelegate,
}

impl RequestInspector for Inspector<'_> {
    fn has_applied_to_current_term(&mut self) -> bool {
        if self.delegate.applied_index_term == self.delegate.term {
            true
        } else {
            // only for metric.
            TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.applied_term.inc());
            false
        }
    }

    fn inspect_lease(&mut self) -> LeaseState {
        // TODO: disable localreader if we did not enable raft's check_quorum.
        if self.delegate.leader_lease.is_some() {
            // We skip lease check, because it is postponed until `handle_read`.
            LeaseState::Valid
        } else {
            TLS_LOCAL_READ_METRICS.with(|m| m.borrow_mut().reject_reason.no_lease.inc());
            LeaseState::Expired
        }
    }
}