// Copyright 2023 PingCAP, Inc.

use std::{
    collections::HashMap,
    net::SocketAddr,
    str::FromStr,
    sync::{Arc, Mutex},
    time::Duration,
};

use futures::{FutureExt, TryFutureExt};
use grpcio::{ChannelBuilder, EnvBuilder, Environment, RpcStatus, RpcStatusCode, ServerBuilder};
use kvproto::{
    coprocessor::{DelegateRequest, Request, Response},
    resource_usage_agent::create_resource_metering_pub_sub,
    tikvpb::{Tikv, TikvClient, create_tikv},
};
use resource_metering::FutureExt as _;
use tikv::coprocessor::{
    REQ_TYPE_ANALYZE, REQ_TYPE_CHECKSUM, REQ_TYPE_DAG, REQ_TYPE_DB9_DAG, RemoteCopRequestKind,
    RemoteCopRequestType, RemoteCopScanStats, parse_request_and_handle_remote_cop,
    remote_cop_request_type,
};
use tikv_util::{defer, info, sys::thread::GRPC_SERVER_THREAD, thd_name, time::Instant, warn};
use tracker::{
    GLOBAL_TRACKERS, RequestInfo, RequestType, TrackedFuture, clear_tls_tracker_token,
    set_tls_tracker_token,
};

use crate::{
    metrics::*,
    server::{Context, get_cop_req_tag},
};

pub(crate) fn record_remote_cop_outcome(
    request_type: RemoteCopRequestType,
    outcome: RemoteCopRequestOutcome,
) {
    let requests = match request_type {
        RemoteCopRequestType::Dag => &REMOTE_COPR_REQUESTS_COUNTER_VEC.dag,
        RemoteCopRequestType::Db9 => &REMOTE_COPR_REQUESTS_COUNTER_VEC.db9,
        RemoteCopRequestType::Analyze => &REMOTE_COPR_REQUESTS_COUNTER_VEC.analyze,
        RemoteCopRequestType::Checksum => &REMOTE_COPR_REQUESTS_COUNTER_VEC.checksum,
        RemoteCopRequestType::Unknown => &REMOTE_COPR_REQUESTS_COUNTER_VEC.unknown,
    };
    requests.get(outcome).inc();
}

pub(crate) fn record_remote_cop_scan_stats(scan_stats: RemoteCopScanStats) {
    let scan_keys = match scan_stats.request_kind {
        RemoteCopRequestKind::Select => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.select,
        RemoteCopRequestKind::Index => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.index,
        RemoteCopRequestKind::Db9Select => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.db9_select,
        RemoteCopRequestKind::Db9Index => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.db9_index,
        RemoteCopRequestKind::AnalyzeTable => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.analyze_table,
        RemoteCopRequestKind::AnalyzeIndex => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.analyze_index,
        RemoteCopRequestKind::AnalyzeFullSampling => {
            &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.analyze_full_sampling
        }
        RemoteCopRequestKind::ChecksumTable => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.checksum_table,
        RemoteCopRequestKind::ChecksumIndex => &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.checksum_index,
    };
    scan_keys.processed.inc_by(scan_stats.processed_keys);
    scan_keys.total.inc_by(scan_stats.total_keys);
    REMOTE_COPR_SCAN_BYTES_COUNTER_VEC
        .processed
        .inc_by(scan_stats.processed_bytes);
    REMOTE_COPR_SCAN_BYTES_COUNTER_VEC
        .total
        .inc_by(scan_stats.total_bytes);
}

#[derive(Clone, Default, Serialize, Deserialize, PartialEq, Debug)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct Config {
    pub addr: String,
    pub max_handle_duration: Duration,
}

pub struct RemoteCopServer {
    grpc_server: grpcio::Server,
}

impl RemoteCopServer {
    pub(crate) fn new(ctx: Arc<Context>, cfg: Config) -> RemoteCopServer {
        let env = Arc::new(
            EnvBuilder::new()
                .cq_count(16)
                .name_prefix(thd_name!(GRPC_SERVER_THREAD))
                .build(),
        );
        let channel_args = ChannelBuilder::new(env.clone())
            .stream_initial_window_size(2 * 1024 * 1024)
            .max_concurrent_stream(1024)
            .max_receive_message_len(-1)
            .max_send_message_len(-1)
            .http2_max_ping_strikes(i32::MAX) // For pings without data from clients.
            .keepalive_time(Duration::from_secs(10))
            .keepalive_timeout(Duration::from_secs(3))
            .build_args();
        let addr = SocketAddr::from_str(&cfg.addr).unwrap();
        let cop_service = CopService::new(ctx, env.clone(), cfg);
        let security_mgr = cop_service.ctx.pd.get_security_mgr();
        let sb = ServerBuilder::new(env)
            .channel_args(channel_args)
            .register_service(create_tikv(cop_service.clone()))
            .register_service(create_resource_metering_pub_sub(
                cop_service.ctx.rsmeter_pubsub_service.clone(),
            ));
        let sb = security_mgr.bind(sb, &addr.ip().to_string(), addr.port());
        let grpc_server = sb.build().unwrap();
        Self { grpc_server }
    }

    pub fn start(&mut self) {
        self.grpc_server.start()
    }
}

#[derive(Clone)]
pub struct CopService {
    ctx: Arc<Context>,
    env: Arc<Environment>,
    store_addrs: Arc<Mutex<HashMap<u64, String>>>,
    channels: Arc<Mutex<HashMap<String, TikvClient>>>,
    cfg: Config,
}

impl CopService {
    fn new(ctx: Arc<Context>, env: Arc<Environment>, cfg: Config) -> Self {
        Self {
            ctx,
            env,
            store_addrs: Arc::new(Mutex::new(HashMap::new())),
            channels: Arc::new(Mutex::new(HashMap::new())),
            cfg,
        }
    }
}

impl Tikv for CopService {
    fn coprocessor(
        &mut self,
        ctx: grpcio::RpcContext<'_>,
        req: Request,
        sink: grpcio::UnarySink<Response>,
    ) {
        let handle_start = Instant::now_coarse();
        let tag = get_cop_req_tag(&req);
        let request_type = remote_cop_request_type(&req);
        REMOTE_COPR_RANGE_COUNT_HISTOGRAM.observe(req.get_ranges().len() as f64);
        let resource_tag = self.ctx.resource_tag_factory.new_tag(req.get_context());
        if self.ctx.config.black_list.rejects_request(&req) {
            warn!("request rejected by blacklist"; "tag" => &tag);
            REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.black_list.inc();
            record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
            ctx.spawn(
                sink.fail(RpcStatus::with_message(
                    RpcStatusCode::PERMISSION_DENIED,
                    "rejected".to_string(),
                ))
                .unwrap_or_else(|e| {
                    warn!("failed to send rpc status"; "err" => ?e);
                }),
            );
            return;
        }
        let store_id = req.get_context().get_peer().get_store_id();
        let client_res = self.get_client(store_id);
        if let Err(err) = client_res {
            warn!("get client failed"; "err" => ?err);
            REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.delegate_error.inc();
            record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
            ctx.spawn(
                sink.fail(RpcStatus::with_message(
                    RpcStatusCode::INTERNAL,
                    err.to_string(),
                ))
                .unwrap_or_else(|e| {
                    warn!("failed to send rpc status"; "err" => ?e);
                }),
            );
            return;
        }
        let client = client_res.unwrap();
        let mut delegate_req = DelegateRequest::new();
        let key_ranges = req.get_ranges().to_vec();
        delegate_req.set_context(req.get_context().clone());
        delegate_req.set_ranges(key_ranges.into());
        delegate_req.set_start_ts(req.get_start_ts());
        let max_handle_duration = self.cfg.max_handle_duration;
        let snap_ctx = self.ctx.get_snap_ctx(None);
        let mem_limiter = self.ctx.memory_limiter.clone();
        let peer = Some(ctx.peer());
        let max_resp_size = self.ctx.config.cop_max_resp_size.0;
        let future = async move {
            let delegate_start = Instant::now_coarse();
            let delegate_receiver =
                client
                    .delegate_coprocessor_async(&delegate_req)
                    .map_err(|e| {
                        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.delegate_error.inc();
                        record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                        tikv::coprocessor::Error::Other(format!("{:?}", e))
                    })?;
            let mut resp = delegate_receiver.await.map_err(|e| {
                REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.delegate_error.inc();
                record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                tikv::coprocessor::Error::Other(format!("{:?}", e))
            })?;

            let mem_table_data = resp.take_mem_table_data();
            let snapshot_data = resp.take_snapshot();

            let snap_start = Instant::now_coarse();
            let snap_access_res = kvengine::SnapAccess::construct_snapshot(
                &tag,
                &snap_ctx,
                &mem_table_data,
                &snapshot_data,
                mem_limiter,
            )
            .await;
            if let Err(err) = snap_access_res.as_ref() {
                if matches!(err, kvengine::Error::MemoryLimitExceeded(_)) {
                    REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.memory_limiter.inc();
                } else {
                    REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.snapshot_error.inc();
                }
                record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                return Err(tikv::coprocessor::Error::Other(format!("{:?}", err)));
            }
            let (snap_access, mem_limiter_guard) = snap_access_res.unwrap();
            let snapshot = rfstore::store::RegionSnapshot::from_snapshot(snap_access, None, None);

            let tracker = GLOBAL_TRACKERS.insert(tracker::Tracker::new(RequestInfo::new(
                req.get_context(),
                RequestType::Unknown,
                req.start_ts,
            )));
            set_tls_tracker_token(tracker);
            defer! {{
                clear_tls_tracker_token();
                GLOBAL_TRACKERS.remove(tracker);
            }};
            let process_start = Instant::now_coarse();
            let result = TrackedFuture::new(async move {
                parse_request_and_handle_remote_cop(
                    req,
                    peer,
                    max_handle_duration,
                    max_resp_size,
                    snapshot,
                )
                .await
            })
            .await;
            drop(mem_limiter_guard);

            match &result {
                Ok(result) => {
                    let finish_time = Instant::now_coarse();
                    let wait_dur = delegate_start.saturating_duration_since(handle_start);
                    let delegate_req_dur = snap_start.saturating_duration_since(delegate_start);
                    let snap_dur = process_start.saturating_duration_since(snap_start);
                    let process_dur = finish_time.saturating_duration_since(process_start);
                    let handle_dur = finish_time.saturating_duration_since(handle_start);
                    info!(
                        "finished remote cop service";
                        "tag" => tag,
                        "resp_size" => result.response.data.len(),
                        "max_handle_duration" => ?max_handle_duration,
                        "wait_dur" => ?wait_dur,
                        "delegate_req_dur" => ?delegate_req_dur,
                        "snap_dur" => ?snap_dur,
                        "process_dur" => ?process_dur,
                        "handle_dur" => ?handle_dur,
                    );
                }
                Err(err) => {
                    if matches!(err, tikv::coprocessor::Error::DeadlineExceeded) {
                        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
                            .process_timeout
                            .inc();
                    } else {
                        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC.process_error.inc();
                    }
                    record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::failure);
                }
            }
            result
        };
        let task = async move {
            match future.await {
                Ok(mut result) => {
                    let scan_stats = result.scan_stats;
                    record_remote_cop_scan_stats(scan_stats);
                    record_remote_cop_outcome(request_type, RemoteCopRequestOutcome::success);
                    if let Err(err) = sink.success(result.response.consume()).await {
                        REMOTE_COPR_FAILED_REQUESTS_COUNTER_VEC
                            .response_send_error
                            .inc();
                        return Err(err);
                    }
                }
                Err(err) => {
                    sink.fail(RpcStatus::with_message(
                        RpcStatusCode::INTERNAL,
                        err.to_string(),
                    ))
                    .await?
                }
            }
            Ok(())
        }
        .map_err(|e: grpcio::Error| {
            warn!("remote cop failed";
                "request" => "coprocessor",
                "err" => ?e
            );
        })
        .in_resource_metering_tag(resource_tag)
        .map(|_| ());
        ctx.spawn(task);
    }
}

impl CopService {
    fn resolve_store_addr(&mut self, store_id: u64) -> pd_client::Result<String> {
        let mut addrs = self.store_addrs.lock().unwrap();
        if let Some(addr) = addrs.get(&store_id) {
            return Ok(addr.clone());
        }
        let store = self.ctx.pd.get_store(store_id)?;
        addrs.insert(store_id, store.address.clone());
        Ok(store.address)
    }

    fn get_client(&mut self, store_id: u64) -> pd_client::Result<TikvClient> {
        let addr = self.resolve_store_addr(store_id)?;
        let mut channels = self.channels.lock().unwrap();
        if let Some(channel) = channels.get(&addr) {
            return Ok(channel.clone());
        }
        let cb = ChannelBuilder::new(self.env.clone());
        let security_mgr = self.ctx.pd.get_security_mgr();
        let channel = security_mgr.connect(cb, &addr);
        let client = TikvClient::new(channel);
        channels.insert(addr, client.clone());
        Ok(client)
    }
}

/// The blacklist entries for a single remote coprocessor request type.
#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, Debug, Default)]
#[serde(default, rename_all = "kebab-case")]
pub struct BlackListConfig {
    /// Keyspace IDs rejected for the associated request type.
    pub keyspace_ids: Vec<u32>,
}

/// The cloud-worker blacklist grouped by remote coprocessor request type.
#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, Debug, Default)]
#[serde(default, rename_all = "kebab-case")]
pub struct BlackListConfigWithReqTypes {
    /// Rejected keyspace IDs for DAG-style remote coprocessor requests,
    /// including DB9 DAG.
    pub remote_cop: BlackListConfig,
    /// Rejected keyspace IDs for remote analyze requests.
    pub analyze: BlackListConfig,
    /// Rejected keyspace IDs for remote checksum requests.
    pub checksum: BlackListConfig,
}

impl BlackListConfigWithReqTypes {
    /// Returns `true` when no request type rejects any keyspace.
    pub fn is_empty(&self) -> bool {
        self.remote_cop.keyspace_ids.is_empty()
            && self.analyze.keyspace_ids.is_empty()
            && self.checksum.keyspace_ids.is_empty()
    }

    /// Returns `true` when the given request type rejects the provided
    /// keyspace.
    pub fn contains_keyspace_id(&self, req_type: i64, keyspace_id: u32) -> bool {
        let keyspace_ids = match req_type {
            REQ_TYPE_DAG | REQ_TYPE_DB9_DAG => &self.remote_cop.keyspace_ids,
            REQ_TYPE_ANALYZE => &self.analyze.keyspace_ids,
            REQ_TYPE_CHECKSUM => &self.checksum.keyspace_ids,
            _ => return false,
        };
        keyspace_ids.contains(&keyspace_id)
    }

    /// Returns `true` when the blacklist rejects the supplied request.
    pub fn rejects_request(&self, req: &Request) -> bool {
        self.contains_keyspace_id(req.get_tp(), req.get_context().get_keyspace_id())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_record_remote_cop_scan_keys_by_request_kind() {
        let processed = &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.select.processed;
        let total = &REMOTE_COPR_SCAN_KEYS_COUNTER_VEC.select.total;
        let processed_before = processed.get();
        let total_before = total.get();

        record_remote_cop_scan_stats(RemoteCopScanStats {
            request_kind: RemoteCopRequestKind::Select,
            processed_keys: 3,
            total_keys: 5,
            processed_bytes: 7,
            total_bytes: 11,
        });

        assert_eq!(processed.get() - processed_before, 3);
        assert_eq!(total.get() - total_before, 5);
    }

    #[test]
    fn test_black_list_config_match_request_type() {
        let black_list = BlackListConfigWithReqTypes {
            remote_cop: BlackListConfig {
                keyspace_ids: vec![1],
            },
            analyze: BlackListConfig {
                keyspace_ids: vec![2],
            },
            checksum: BlackListConfig {
                keyspace_ids: vec![3],
            },
        };

        assert!(black_list.contains_keyspace_id(REQ_TYPE_DAG, 1));
        assert!(black_list.contains_keyspace_id(REQ_TYPE_DB9_DAG, 1));
        assert!(black_list.contains_keyspace_id(REQ_TYPE_ANALYZE, 2));
        assert!(black_list.contains_keyspace_id(REQ_TYPE_CHECKSUM, 3));

        assert!(!black_list.contains_keyspace_id(REQ_TYPE_DAG, 2));
        assert!(!black_list.contains_keyspace_id(REQ_TYPE_DB9_DAG, 2));
        assert!(!black_list.contains_keyspace_id(REQ_TYPE_ANALYZE, 3));
        assert!(!black_list.contains_keyspace_id(REQ_TYPE_CHECKSUM, 1));
        assert!(!black_list.contains_keyspace_id(0, 1));
    }

    #[test]
    fn test_black_list_config_rejects_request() {
        let black_list = BlackListConfigWithReqTypes {
            remote_cop: BlackListConfig {
                keyspace_ids: vec![7],
            },
            ..Default::default()
        };
        let mut req = Request::default();
        req.set_tp(REQ_TYPE_DAG);
        req.mut_context().set_keyspace_id(7);

        assert!(black_list.rejects_request(&req));

        req.set_tp(REQ_TYPE_DB9_DAG);
        assert!(black_list.rejects_request(&req));

        req.mut_context().set_keyspace_id(8);
        assert!(!black_list.rejects_request(&req));
    }

    #[test]
    fn test_worker_config_black_list_toml_layout() {
        let config: crate::Config = toml::from_str(
            r#"
                [black-list.remote-cop]
                keyspace-ids = [3, 1, 3]

                [black-list.analyze]
                keyspace-ids = [2]
            "#,
        )
        .unwrap();

        assert_eq!(
            config.black_list.remote_cop.keyspace_ids,
            vec![3_u32, 1_u32, 3_u32]
        );
        assert_eq!(config.black_list.analyze.keyspace_ids, vec![2_u32]);
        assert!(config.black_list.checksum.keyspace_ids.is_empty());
    }
}