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

use std::sync::{Arc, Mutex};

use crossbeam::channel::TrySendError;
use kvproto::{
    errorpb,
    metapb::{self, RegionEpoch},
    raft_cmdpb::{RaftCmdRequest, RaftCmdResponse},
    raft_serverpb::RaftMessage,
};
use raftstore::store::fsm::ChangeObserver;
use tikv_util::{
    deadline::Deadline,
    error,
    mpsc::{self, LooseBoundedSender, Receiver, Sender},
    time::ThreadReadId,
    warn,
};
use trace_event::types::TraceContext;
use txn_types::Key;

use crate::{
    Result as RaftStoreResult,
    store::{
        Callback, CasualMessage, CasualRouter, LocalReader, PeerMsg, ProposalRouter, RaftCommand,
        SignificantMsg, SplitSource, StoreMsg, StoreRouter,
        cmd_resp::message_error,
        metrics::{CHANNEL_FULL_COUNTER_VEC, RAFT_ROUTER_RAFT_MESSAGE_DROPPED_TOTAL},
    },
};

/// Routes messages to the raftstore.
pub trait RaftStoreRouter: StoreRouter + ProposalRouter + CasualRouter + Send + Clone {
    /// Sends RaftMessage to local store.
    fn send_raft_msg(&self, msg: RaftMessage);

    /// Sends a significant message. We should guarantee that the message can't
    /// be dropped.
    fn significant_send(&self, region_id: u64, msg: SignificantMsg);

    /// Send a casual message to the given region.
    fn send_casual_msg(&self, region_id: u64, msg: CasualMessage) {
        <Self as CasualRouter>::send(self, region_id, msg)
    }

    /// Send a store message to the backend raft batch system.
    fn send_store_msg(&self, msg: StoreMsg) {
        <Self as StoreRouter>::send(self, msg)
    }

    /// Sends RaftCmdRequest to local store.
    fn send_command(&self, req: RaftCmdRequest, trace_ctx: TraceContext, cb: Callback) {
        send_command_impl(self, req, trace_ctx, cb, None)
    }

    fn send_command_with_deadline(
        &self,
        req: RaftCmdRequest,
        trace_ctx: TraceContext,
        cb: Callback,
        deadline: Deadline,
    ) {
        send_command_impl(self, req, trace_ctx, cb, Some(deadline))
    }

    /// Reports the peer being unreachable to the Region.
    fn report_unreachable(&self, region_id: u64, store_id: u64) {
        let msg = SignificantMsg::StoreUnreachable { store_id };
        self.significant_send(region_id, msg)
    }

    /// Broadcast an `StoreUnreachable` event to all Raft groups.
    fn broadcast_unreachable(&self, store_id: u64) {
        self.send_store_msg(StoreMsg::StoreUnreachable { store_id });
    }
}

pub trait LocalReadRouter: Send + Clone {
    fn read(
        &self,
        read_id: Option<ThreadReadId>,
        req: RaftCmdRequest,
        cb: Callback,
    ) -> RaftStoreResult<()>;
}

impl StoreRouter for LocalReader {
    fn send(&self, msg: StoreMsg) {
        StoreRouter::send(&self.router, msg)
    }
}

impl ProposalRouter for LocalReader {
    fn send(&self, cmd: RaftCommand) {
        ProposalRouter::send(&self.router, cmd)
    }
}

impl CasualRouter for LocalReader {
    fn send(&self, region_id: u64, msg: CasualMessage) {
        CasualRouter::send(&self.router, region_id, msg)
    }
}

#[cfg(any(test, feature = "testexport"))]
impl RaftStoreRouter for LocalReader {
    fn send_raft_msg(&self, msg: RaftMessage) {
        RaftStoreRouter::send_raft_msg(&self.router, msg)
    }

    /// Sends a significant message. We should guarantee that the message can't
    /// be dropped.
    fn significant_send(&self, region_id: u64, msg: SignificantMsg) {
        RaftStoreRouter::significant_send(&self.router, region_id, msg)
    }

    fn broadcast_unreachable(&self, store_id: u64) {
        RaftStoreRouter::broadcast_unreachable(&self.router, store_id)
    }
}

#[derive(Clone)]
pub struct RaftStoreBlackHole;

impl CasualRouter for RaftStoreBlackHole {
    fn send(&self, _: u64, _: CasualMessage) {}
}

impl ProposalRouter for RaftStoreBlackHole {
    fn send(&self, _: RaftCommand) {}
}

impl StoreRouter for RaftStoreBlackHole {
    fn send(&self, _: StoreMsg) {}
}

impl RaftStoreRouter for RaftStoreBlackHole {
    /// Sends RaftMessage to local store.
    fn send_raft_msg(&self, _: RaftMessage) {}

    /// Sends a significant message. We should guarantee that the message can't
    /// be dropped.
    fn significant_send(&self, _: u64, _: SignificantMsg) {}
}

/// Sender type for the peer message channel used by `RaftRouter`.
pub type RaftRouterSender = LooseBoundedSender<(u64, Box<PeerMsg>)>;

/// Receiver type for the peer message channel used by `RaftRouter`.
pub type RaftRouterReceiver = Receiver<(u64, Box<PeerMsg>)>;

#[inline]
pub fn raft_router_peer_channel(cap: usize) -> (RaftRouterSender, RaftRouterReceiver) {
    mpsc::loose_bounded(cap)
}

#[derive(Clone)]
pub struct RaftRouter {
    pub(crate) store_sender: Sender<StoreMsg>,
    pub(crate) peer_sender: RaftRouterSender,
}

impl RaftRouter {
    pub fn new(peer_sender: RaftRouterSender, store_sender: Sender<StoreMsg>) -> Self {
        Self {
            store_sender,
            peer_sender,
        }
    }

    pub(crate) fn try_send(&self, id: u64, msg: PeerMsg) -> bool {
        if let Err(err) = self.peer_sender.force_send((id, Box::new(msg))) {
            warn!("send failed {:?}", err);
            return false;
        }
        true
    }

    pub(crate) fn send(&self, id: u64, msg: PeerMsg) {
        self.try_send(id, msg);
    }

    /// Sends a peer message directly for integration tests.
    #[cfg(any(test, feature = "testexport"))]
    pub fn send_peer_msg_for_test(&self, id: u64, msg: PeerMsg) {
        self.send(id, msg);
    }

    pub(crate) fn send_command_checked(
        &self,
        req: RaftCmdRequest,
        trace_ctx: TraceContext,
        cb: Callback,
    ) -> bool {
        let region_id = req.get_header().get_region_id();
        let cmd = RaftCommand::new(req, trace_ctx, cb);
        self.try_send(region_id, PeerMsg::RaftCommand(cmd))
    }

    pub(crate) fn send_engine_change_set(&self, id: u64, cs: kvenginepb::ChangeSet, cb: Callback) {
        let msg = PeerMsg::GenerateEngineChangeSet(cs, cb);
        if let Err(mpsc::SendError((region_id, msg))) =
            self.peer_sender.force_send((id, Box::new(msg)))
        {
            error!("send engine change set failed"; "region_id" => region_id);
            if let PeerMsg::GenerateEngineChangeSet(_, cb) = *msg {
                cb.invoke_with_response(region_missing_response(
                    region_id,
                    "peer channel is disconnected",
                ));
            }
        }
    }

    pub(crate) fn send_store_engine_change_set(&self, cs: kvenginepb::ChangeSet, cb: Callback) {
        let region_id = cs.get_shard_id();
        let msg = StoreMsg::GenerateEngineChangeSet(cs, cb);
        if let Err(mpsc::SendError(msg)) = self.store_sender.send(msg) {
            error!("send store engine change set failed"; "region_id" => region_id);
            if let StoreMsg::GenerateEngineChangeSet(_, cb) = msg {
                cb.invoke_with_response(message_error(format!(
                    "store channel is disconnected while sending engine change set for region {}",
                    region_id
                )));
            }
        }
    }

    pub(crate) fn send_store(&self, msg: StoreMsg) {
        if let Err(err) = self.store_sender.send(msg) {
            warn!("send store failed {:?}", err)
        }
    }

    pub(crate) fn force_send_raft_msg(&self, msg: RaftMessage) {
        let region_id = msg.get_region_id();
        let raft_msg = PeerMsg::RaftMessage(msg);
        self.send(region_id, raft_msg);
    }

    /// Routes a load-split candidate through the current local peer so its
    /// leadership and identity are validated at execution time. Returns false
    /// if the peer channel is disconnected.
    pub fn ask_load_split(
        &self,
        region_id: u64,
        region_epoch: RegionEpoch,
        split_key: Key,
    ) -> bool {
        self.try_send(
            region_id,
            PeerMsg::CasualMessage(CasualMessage::SplitRegion {
                region_epoch,
                split_keys: vec![split_key.into_encoded()],
                callback: Callback::None,
                source: SplitSource::Load,
            }),
        )
    }
}

pub(crate) fn region_missing_response(region_id: u64, reason: &str) -> RaftCmdResponse {
    let mut resp = RaftCmdResponse::default();
    let mut err = errorpb::Error::default();
    err.set_message(format!("region {} is missing: {}", region_id, reason));
    err.mut_region_not_found().set_region_id(region_id);
    resp.mut_header().set_error(err);
    resp
}

impl RaftStoreRouter for RaftRouter {
    fn send_raft_msg(&self, msg: RaftMessage) {
        let region_id = msg.get_region_id();
        let raft_msg = PeerMsg::RaftMessage(msg);
        match self.peer_sender.try_send((region_id, Box::new(raft_msg))) {
            Ok(()) => {}
            Err(TrySendError::Full(_)) => {
                CHANNEL_FULL_COUNTER_VEC
                    .with_label_values(&["normal"])
                    .inc();
                RAFT_ROUTER_RAFT_MESSAGE_DROPPED_TOTAL.full.inc();
            }
            Err(TrySendError::Disconnected(_)) => {
                RAFT_ROUTER_RAFT_MESSAGE_DROPPED_TOTAL.disconnected.inc();
            }
        }
    }

    fn significant_send(&self, region_id: u64, msg: SignificantMsg) {
        let msg = PeerMsg::SignificantMsg(msg);
        self.send(region_id, msg);
    }

    fn broadcast_unreachable(&self, store_id: u64) {
        self.send_store_msg(StoreMsg::StoreUnreachable { store_id });
    }
}

fn send_command_impl(
    router: &impl ProposalRouter,
    req: RaftCmdRequest,
    trace_ctx: TraceContext,
    cb: Callback,
    _deadline: Option<Deadline>,
) {
    let cmd = RaftCommand::new(req, trace_ctx, cb);
    // TODO(x) handle deadline
    router.send(cmd)
}

/// A handle for cdc and pitr to schedule some command back to raftstore.
pub trait CdcHandle: Clone + Send {
    fn capture_change(
        &self,
        region_id: u64,
        region_epoch: metapb::RegionEpoch,
        change_observer: ChangeObserver,
        callback: Callback,
    ) -> RaftStoreResult<()>;

    fn check_leadership(&self, region_id: u64, callback: Callback) -> RaftStoreResult<()>;
}

impl<T: CdcHandle> CdcHandle for Arc<Mutex<T>> {
    fn capture_change(
        &self,
        region_id: u64,
        region_epoch: metapb::RegionEpoch,
        change_observer: ChangeObserver,
        callback: Callback,
    ) -> RaftStoreResult<()> {
        Mutex::lock(self).unwrap().capture_change(
            region_id,
            region_epoch,
            change_observer,
            callback,
        )
    }

    fn check_leadership(&self, region_id: u64, callback: Callback) -> RaftStoreResult<()> {
        Mutex::lock(self)
            .unwrap()
            .check_leadership(region_id, callback)
    }
}

/// A wrapper of SignificantRouter that is specialized for implementing
/// CdcHandle.
#[derive(Clone)]
pub struct CdcRaftRouter<T>(pub T);

impl<T> std::ops::Deref for CdcRaftRouter<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T: RaftStoreRouter> CdcHandle for CdcRaftRouter<T> {
    fn capture_change(
        &self,
        region_id: u64,
        region_epoch: metapb::RegionEpoch,
        change_observer: ChangeObserver,
        callback: Callback,
    ) -> RaftStoreResult<()> {
        self.0.significant_send(
            region_id,
            SignificantMsg::CaptureChange {
                cmd: change_observer,
                region_epoch,
                callback,
                can_apply: false,
            },
        );
        Ok(())
    }

    fn check_leadership(&self, region_id: u64, callback: Callback) -> RaftStoreResult<()> {
        self.0
            .significant_send(region_id, SignificantMsg::LeaderCallback(callback));
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use kvproto::raft_serverpb::RaftMessage;
    use raft_proto::eraftpb::{Entry, Message};

    use super::*;
    use crate::{
        RaftRouter, RaftStoreRouter,
        store::{
            Callback, PeerMsg,
            metrics::{CHANNEL_FULL_COUNTER_VEC, RAFT_ROUTER_RAFT_MESSAGE_DROPPED_TOTAL},
        },
    };

    fn new_raft_message(region_id: u64, payload_len: usize) -> RaftMessage {
        let mut entry = Entry::default();
        entry.set_data(vec![b'x'; payload_len].into());

        let mut message = Message::default();
        message.mut_entries().push(entry);

        let mut raft_message = RaftMessage::default();
        raft_message.set_region_id(region_id);
        raft_message.set_message(message);
        raft_message
    }

    #[test]
    fn test_send_store_engine_change_set_reports_send_failure() {
        let (peer_sender, _peer_receiver) = raft_router_peer_channel(0);
        let (store_sender, store_receiver) = mpsc::unbounded();
        drop(store_receiver);
        let router = RaftRouter::new(peer_sender, store_sender);
        let mut cs = kvenginepb::ChangeSet::default();
        cs.set_shard_id(42);
        let (tx, rx) = std::sync::mpsc::channel();
        let cb = Callback::write(Box::new(move |res| {
            tx.send(res.response).unwrap();
        }));

        router.send_store_engine_change_set(cs, cb);

        let resp = rx.recv_timeout(std::time::Duration::from_secs(1)).unwrap();
        assert!(resp.get_header().has_error());
        assert!(
            resp.get_header()
                .get_error()
                .get_message()
                .contains("store channel is disconnected")
        );
    }

    #[test]
    fn test_router_drops_only_bounded_raft_messages() {
        let (peer_sender, peer_receiver) = raft_router_peer_channel(1);
        let (store_sender, _store_receiver) = tikv_util::mpsc::unbounded();
        let router = RaftRouter::new(peer_sender, store_sender);
        let dropped_before = RAFT_ROUTER_RAFT_MESSAGE_DROPPED_TOTAL.full.get();
        let channel_full = CHANNEL_FULL_COUNTER_VEC.with_label_values(&["normal"]);
        let channel_full_before = channel_full.get();

        for _ in 0..10 {
            router.send_raft_msg(new_raft_message(1, 1));
        }
        router.send(1, PeerMsg::Tick);

        assert!(RAFT_ROUTER_RAFT_MESSAGE_DROPPED_TOTAL.full.get() > dropped_before);
        assert!(channel_full.get() > channel_full_before);
        while let Ok((_, msg)) = peer_receiver.try_recv() {
            if matches!(*msg, PeerMsg::Tick) {
                return;
            }
        }
        panic!("forced tick message should bypass the raft message bound");
    }

    #[test]
    fn test_ask_load_split_reports_send_failure() {
        let (peer_sender, peer_receiver) = raft_router_peer_channel(1);
        let (store_sender, _store_receiver) = tikv_util::mpsc::unbounded();
        let router = RaftRouter::new(peer_sender, store_sender);
        let mut epoch = RegionEpoch::default();
        epoch.set_version(7);

        assert!(router.ask_load_split(42, epoch.clone(), Key::from_encoded(b"k".to_vec())));
        let (region_id, msg) = peer_receiver.try_recv().unwrap();
        assert_eq!(region_id, 42);
        match *msg {
            PeerMsg::CasualMessage(CasualMessage::SplitRegion {
                region_epoch,
                split_keys,
                callback,
                source,
            }) => {
                assert_eq!(region_epoch, epoch);
                assert_eq!(split_keys, vec![b"k".to_vec()]);
                assert!(matches!(callback, Callback::None));
                assert_eq!(source, SplitSource::Load);
            }
            _ => panic!("unexpected peer message"),
        }

        drop(peer_receiver);
        assert!(!router.ask_load_split(42, epoch, Key::from_encoded(b"k".to_vec())));
    }
}