// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
#![feature(type_alias_impl_trait)]

#[macro_use]
extern crate log;

mod errors;
pub mod region;

pub use crate::errors::{HasError, HasRegionError};
#[doc(inline)]
pub use crate::region::{Region, RegionId, RegionVerId, StoreId};
pub use kvproto::tikvpb::TikvClient;
pub use tikv_client_common::{security::SecurityManager, Error, ErrorKind, Key, Result};

use async_trait::async_trait;
use derive_new::new;
use grpcio::{CallOption, Environment};
use std::{sync::Arc, time::Duration};
use tikv_client_common::stats::tikv_stats;

/// A trait for connecting to TiKV stores.
pub trait KvConnect: Sized + Send + Sync + 'static {
    type KvClient: KvClient + Clone + Send + Sync + 'static;

    fn connect(&self, address: &str) -> Result<Self::KvClient>;

    fn connect_to_store(&self, region: Region, address: String) -> Result<Store<Self::KvClient>> {
        info!("connect to tikv endpoint: {:?}", &address);
        let client = self.connect(address.as_str())?;
        Ok(Store::new(region, client))
    }
}

pub type RpcFnType<Req, Resp> =
    for<'a, 'b> fn(
        &'a TikvClient,
        &'b Req,
        CallOption,
    )
        -> std::result::Result<::grpcio::ClientUnaryReceiver<Resp>, ::grpcio::Error>;

#[derive(new, Clone)]
pub struct TikvConnect {
    env: Arc<Environment>,
    security_mgr: Arc<SecurityManager>,
    timeout: Duration,
}

impl KvConnect for TikvConnect {
    type KvClient = impl KvClient + Clone + Send + Sync + 'static;

    fn connect(&self, address: &str) -> Result<Self::KvClient> {
        self.security_mgr
            .connect(self.env.clone(), address, TikvClient::new)
            .map(|c| KvRpcClient::new(Arc::new(c), self.timeout))
    }
}

#[async_trait]
pub trait KvClient {
    async fn dispatch<Req, Resp>(&self, fun: RpcFnType<Req, Resp>, request: Req) -> Result<Resp>
    where
        Req: Send + Sync + 'static,
        Resp: HasError + Sized + Clone + Send + 'static;
}

/// This client handles requests for a single TiKV node. It converts the data
/// types and abstractions of the client program into the grpc data types.
#[derive(new, Clone)]
struct KvRpcClient {
    rpc_client: Arc<TikvClient>,
    timeout: Duration,
}

#[async_trait]
impl KvClient for KvRpcClient {
    async fn dispatch<Req, Resp>(&self, fun: RpcFnType<Req, Resp>, request: Req) -> Result<Resp>
    where
        Req: Send + Sync + 'static,
        Resp: HasError + Sized + Clone + Send + 'static,
    {
        fun(
            &self.rpc_client,
            &request,
            CallOption::default().timeout(self.timeout),
        )?
        .await
        .map_err(|e| ErrorKind::Grpc(e).into())
    }
}

#[derive(new)]
pub struct Store<Client: KvClient> {
    pub region: Region,
    pub client: Client,
}

impl<Client: KvClient> Store<Client> {
    pub async fn dispatch<Req, Resp>(
        &self,
        request_name: &'static str,
        fun: RpcFnType<Req, Resp>,
        request: Req,
    ) -> Result<Resp>
    where
        Req: Send + Sync + 'static,
        Resp: HasError + Sized + Clone + Send + 'static,
    {
        let result = self.client.dispatch(fun, request).await;

        tikv_stats(request_name).done(result)
    }
}