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

use std::{sync::Arc, thread, time::Duration};

use fail::fail_point;
use futures::{SinkExt, TryFutureExt, TryStreamExt, future};
use grpcio::{
    ClientStreamingSink, DuplexSink, EnvBuilder, RequestStream, RpcContext, RpcStatus,
    RpcStatusCode, Server as GrpcServer, ServerBuilder, ServerStreamingSink, UnarySink, WriteFlags,
};
use kvproto::pdpb::*;
use pd_client::Error as PdError;
use security::*;

use super::mocker::*;

pub struct Server<C: PdMocker> {
    server: Option<GrpcServer>,
    mocker: PdMock<C>,
}

impl Server<Service> {
    pub fn new(eps_count: usize) -> Server<Service> {
        let mgr = SecurityManager::new(&SecurityConfig::default()).unwrap();
        let eps = vec![("127.0.0.1".to_owned(), 0); eps_count];
        let case = Option::None::<Arc<Service>>;
        Self::with_configuration(&mgr, eps, case)
    }

    pub fn default_handler(&self) -> &Service {
        &self.mocker.default_handler
    }
}

impl<C: PdMocker + Send + Sync + 'static> Server<C> {
    pub fn with_case(eps_count: usize, case: Arc<C>) -> Server<C> {
        let mgr = SecurityManager::new(&SecurityConfig::default()).unwrap();
        let eps = vec![("127.0.0.1".to_owned(), 0); eps_count];
        Server::with_configuration(&mgr, eps, Some(case))
    }

    pub fn with_configuration(
        mgr: &SecurityManager,
        eps: Vec<(String, u16)>,
        case: Option<Arc<C>>,
    ) -> Server<C> {
        let cluster_id = case
            .as_ref()
            .map_or(DEFAULT_CLUSTER_ID, |c| c.get_cluster_id());
        let handler = Arc::new(Service::new(cluster_id));
        let default_handler = Arc::clone(&handler);
        let mocker = PdMock {
            default_handler,
            case,
        };
        let mut server = Server {
            server: None,
            mocker,
        };
        server.start(mgr, eps);
        server
    }

    pub fn start(&mut self, mgr: &SecurityManager, eps: Vec<(String, u16)>) {
        let service = create_pd(self.mocker.clone());
        let env = Arc::new(
            EnvBuilder::new()
                .cq_count(1)
                .name_prefix(thd_name!("mock-server"))
                .build(),
        );
        let mut sb = ServerBuilder::new(env).register_service(service);
        for (host, port) in eps {
            sb = mgr.bind(sb, &host, port);
        }

        let mut server = sb.build().unwrap();
        {
            let addrs: Vec<String> = server
                .bind_addrs()
                .map(|(host, port)| format!("{}:{}", host, port))
                .collect();
            self.mocker.default_handler.set_endpoints(addrs.clone());
            if let Some(case) = self.mocker.case.as_ref() {
                case.set_endpoints(addrs);
            }
        }

        server.start();
        self.server = Some(server);
        // Ensure that server is ready.
        thread::sleep(Duration::from_secs(1));
    }

    pub fn stop(&mut self) {
        self.server
            .take()
            .expect("Server is not started")
            .shutdown();
    }

    pub fn bind_addrs(&self) -> Vec<(String, u16)> {
        self.server
            .as_ref()
            .unwrap()
            .bind_addrs()
            .map(|(host, port)| (host.clone(), port))
            .collect()
    }
}

#[allow(unused_mut)]
fn hijack_unary<F, R, C: PdMocker>(
    mock: &mut PdMock<C>,
    ctx: RpcContext<'_>,
    sink: UnarySink<R>,
    f: F,
) where
    R: Send + 'static,
    F: Fn(&dyn PdMocker) -> Option<Result<R>>,
{
    let resp = mock
        .case
        .as_ref()
        .and_then(|case| f(case.as_ref()))
        .or_else(|| f(mock.default_handler.as_ref()));
    match resp {
        Some(Ok(resp)) => ctx.spawn(
            sink.success(resp)
                .unwrap_or_else(|e| error!("failed to reply: {:?}", e)),
        ),
        Some(Err(err)) => {
            let status = RpcStatus::with_message(RpcStatusCode::UNKNOWN, format!("{:?}", err));
            ctx.spawn(
                sink.fail(status)
                    .unwrap_or_else(|e| error!("failed to reply: {:?}", e)),
            );
        }
        None => {
            let mut status =
                RpcStatus::with_message(RpcStatusCode::UNIMPLEMENTED, "Unimplemented".to_owned());
            #[allow(clippy::redundant_closure_call)]
            (|| {
                fail_point!("connect_leader", |_| {
                    let details = ctx
                        .request_headers()
                        .iter()
                        .find(|(k, _)| *k == "pd-forwarded-host")
                        .map(|(_, v)| std::str::from_utf8(v).unwrap())
                        .unwrap_or("");
                    status =
                        RpcStatus::with_message(RpcStatusCode::UNAVAILABLE, details.to_string());
                })
            })();
            ctx.spawn(
                sink.fail(status)
                    .unwrap_or_else(|e| error!("failed to reply: {:?}", e)),
            );
        }
    }
}

#[derive(Debug)]
struct PdMock<C: PdMocker> {
    default_handler: Arc<Service>,
    case: Option<Arc<C>>,
}

impl<C: PdMocker> Clone for PdMock<C> {
    fn clone(&self) -> Self {
        PdMock {
            default_handler: Arc::clone(&self.default_handler),
            case: self.case.clone(),
        }
    }
}

impl<C: PdMocker + Send + Sync + 'static> Pd for PdMock<C> {
    fn load_global_config(
        &mut self,
        ctx: RpcContext<'_>,
        req: LoadGlobalConfigRequest,
        sink: UnarySink<LoadGlobalConfigResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.load_global_config(&req))
    }

    fn store_global_config(
        &mut self,
        _ctx: RpcContext<'_>,
        _req: StoreGlobalConfigRequest,
        _sink: UnarySink<StoreGlobalConfigResponse>,
    ) {
        unimplemented!()
    }

    fn watch_global_config(
        &mut self,
        ctx: RpcContext<'_>,
        _req: WatchGlobalConfigRequest,
        mut sink: ServerStreamingSink<WatchGlobalConfigResponse>,
    ) {
        ctx.spawn(async move {
            let mut name: usize = 0;
            loop {
                let mut change = GlobalConfigItem::new();
                change.set_name(format!("/global/config/{:?}", name).to_owned());
                change.set_value(format!("{:?}", name));
                let mut wc = WatchGlobalConfigResponse::default();
                wc.set_changes(vec![change].into());
                // simulate network delay
                std::thread::sleep(Duration::from_millis(10));
                name += 1;
                let _ = sink.send((wc, WriteFlags::default())).await;
                let _ = sink.flush().await;
            }
        })
    }

    fn get_members(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetMembersRequest,
        sink: UnarySink<GetMembersResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_members(&req))
    }

    fn tso(
        &mut self,
        ctx: RpcContext<'_>,
        req: RequestStream<TsoRequest>,
        mut resp: DuplexSink<TsoResponse>,
    ) {
        let mock = self.clone();
        let fut = async move {
            // Tolerate errors like RpcFinished(None).
            let _ = resp
                .send_all(&mut req.map_ok(move |r| {
                    let res = mock
                        .case
                        .as_ref()
                        .and_then(|case| case.tso(&r))
                        .or_else(|| mock.default_handler.as_ref().tso(&r));
                    let res = res.unwrap().unwrap(); // panic for simplicity.
                    (res, WriteFlags::default())
                }))
                .await;
            let _ = resp.close().await;
        };
        ctx.spawn(fut);
    }

    fn bootstrap(
        &mut self,
        ctx: RpcContext<'_>,
        req: BootstrapRequest,
        sink: UnarySink<BootstrapResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.bootstrap(&req))
    }

    fn is_bootstrapped(
        &mut self,
        ctx: RpcContext<'_>,
        req: IsBootstrappedRequest,
        sink: UnarySink<IsBootstrappedResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.is_bootstrapped(&req))
    }

    fn alloc_id(
        &mut self,
        ctx: RpcContext<'_>,
        req: AllocIdRequest,
        sink: UnarySink<AllocIdResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.alloc_id(&req))
    }

    fn get_store(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetStoreRequest,
        sink: UnarySink<GetStoreResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_store(&req))
    }

    fn put_store(
        &mut self,
        ctx: RpcContext<'_>,
        req: PutStoreRequest,
        sink: UnarySink<PutStoreResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.put_store(&req))
    }

    fn get_all_stores(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetAllStoresRequest,
        sink: UnarySink<GetAllStoresResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_all_stores(&req))
    }

    fn store_heartbeat(
        &mut self,
        ctx: RpcContext<'_>,
        req: StoreHeartbeatRequest,
        sink: UnarySink<StoreHeartbeatResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.store_heartbeat(&req))
    }

    fn region_heartbeat(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<RegionHeartbeatRequest>,
        sink: DuplexSink<RegionHeartbeatResponse>,
    ) {
        let mock = self.clone();
        ctx.spawn(async move {
            let mut stream = stream.map_err(PdError::from).try_filter_map(move |req| {
                let resp = mock
                    .case
                    .as_ref()
                    .and_then(|case| case.region_heartbeat(&req))
                    .or_else(|| mock.default_handler.region_heartbeat(&req));
                match resp {
                    None => future::ok(None),
                    Some(Ok(resp)) => future::ok(Some((resp, WriteFlags::default()))),
                    Some(Err(e)) => future::err(box_err!("{:?}", e)),
                }
            });
            let mut sink = sink.sink_map_err(PdError::from);
            if sink.send_all(&mut stream).await.is_ok() {
                let _ = sink.close().await;
            }
        });
    }

    fn get_region(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetRegionRequest,
        sink: UnarySink<GetRegionResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_region(&req))
    }

    fn get_region_by_id(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetRegionByIdRequest,
        sink: UnarySink<GetRegionResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_region_by_id(&req))
    }

    fn ask_split(
        &mut self,
        ctx: RpcContext<'_>,
        req: AskSplitRequest,
        sink: UnarySink<AskSplitResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.ask_split(&req))
    }

    fn report_split(
        &mut self,
        _: RpcContext<'_>,
        _: ReportSplitRequest,
        _: UnarySink<ReportSplitResponse>,
    ) {
        unimplemented!()
    }

    fn ask_batch_split(
        &mut self,
        ctx: RpcContext<'_>,
        req: AskBatchSplitRequest,
        sink: UnarySink<AskBatchSplitResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.ask_batch_split(&req))
    }

    fn report_batch_split(
        &mut self,
        ctx: RpcContext<'_>,
        req: ReportBatchSplitRequest,
        sink: UnarySink<ReportBatchSplitResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.report_batch_split(&req))
    }

    fn get_cluster_config(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetClusterConfigRequest,
        sink: UnarySink<GetClusterConfigResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_cluster_config(&req))
    }

    fn put_cluster_config(
        &mut self,
        ctx: RpcContext<'_>,
        req: PutClusterConfigRequest,
        sink: UnarySink<PutClusterConfigResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.put_cluster_config(&req))
    }

    fn scatter_region(
        &mut self,
        ctx: RpcContext<'_>,
        req: ScatterRegionRequest,
        sink: UnarySink<ScatterRegionResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.scatter_region(&req))
    }

    fn get_prev_region(
        &mut self,
        _: RpcContext<'_>,
        _: GetRegionRequest,
        _: UnarySink<GetRegionResponse>,
    ) {
        unimplemented!()
    }

    fn get_gc_safe_point(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetGcSafePointRequest,
        sink: UnarySink<GetGcSafePointResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_gc_safe_point(&req))
    }

    fn update_gc_safe_point(
        &mut self,
        ctx: RpcContext<'_>,
        req: UpdateGcSafePointRequest,
        sink: UnarySink<UpdateGcSafePointResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.update_gc_safe_point(&req))
    }

    fn sync_regions(
        &mut self,
        _ctx: RpcContext<'_>,
        _stream: RequestStream<SyncRegionRequest>,
        _sink: DuplexSink<SyncRegionResponse>,
    ) {
        unimplemented!()
    }

    fn get_operator(
        &mut self,
        ctx: RpcContext<'_>,
        req: GetOperatorRequest,
        sink: UnarySink<GetOperatorResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.get_operator(&req))
    }

    fn scan_regions(
        &mut self,
        _ctx: RpcContext<'_>,
        _req: ScanRegionsRequest,
        _sink: UnarySink<ScanRegionsResponse>,
    ) {
        unimplemented!()
    }

    fn update_service_gc_safe_point(
        &mut self,
        ctx: RpcContext<'_>,
        req: UpdateServiceGcSafePointRequest,
        sink: UnarySink<UpdateServiceGcSafePointResponse>,
    ) {
        hijack_unary(self, ctx, sink, |c| c.update_service_gc_safe_point(&req))
    }

    fn sync_max_ts(
        &mut self,
        _ctx: RpcContext<'_>,
        _req: SyncMaxTsRequest,
        _sink: UnarySink<SyncMaxTsResponse>,
    ) {
        unimplemented!()
    }

    fn split_regions(
        &mut self,
        _: RpcContext<'_>,
        _: kvproto::pdpb::SplitRegionsRequest,
        _: UnarySink<kvproto::pdpb::SplitRegionsResponse>,
    ) {
        unimplemented!()
    }

    fn get_dc_location_info(
        &mut self,
        _: RpcContext<'_>,
        _: kvproto::pdpb::GetDcLocationInfoRequest,
        _: UnarySink<kvproto::pdpb::GetDcLocationInfoResponse>,
    ) {
        unimplemented!()
    }

    fn split_and_scatter_regions(
        &mut self,
        _: RpcContext<'_>,
        _: kvproto::pdpb::SplitAndScatterRegionsRequest,
        _: UnarySink<kvproto::pdpb::SplitAndScatterRegionsResponse>,
    ) {
        unimplemented!()
    }

    fn report_buckets(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<kvproto::pdpb::ReportBucketsRequest>,
        sink: ClientStreamingSink<kvproto::pdpb::ReportBucketsResponse>,
    ) {
        // For tests that rely on a real PD server endpoint, TiKV may report
        // region buckets periodically. The mock PD doesn't need to handle it.
        //
        // Drain the stream to avoid breaking the client's streaming RPC.
        let fut = async move {
            let _ = stream.try_for_each(|_| future::ready(Ok(()))).await;
            let _ = sink
                .success(kvproto::pdpb::ReportBucketsResponse::default())
                .await;
        };
        ctx.spawn(fut);
    }

    fn report_min_resolved_ts(
        &mut self,
        ctx: RpcContext<'_>,
        _req: kvproto::pdpb::ReportMinResolvedTsRequest,
        sink: UnarySink<kvproto::pdpb::ReportMinResolvedTsResponse>,
    ) {
        // No-op for tests.
        let mut resp = kvproto::pdpb::ReportMinResolvedTsResponse::default();
        resp.set_header(self.default_handler.header());
        ctx.spawn(
            sink.success(resp)
                .unwrap_or_else(|e| error!("failed to reply: {:?}", e)),
        );
    }
}