use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
};

use anyhow::{Context, Result, bail};
use futures::future::BoxFuture;
use log::warn;
use tici_common::{index::IndexKey, metrics::OperationMetricsGuard};
use tici_grpc::{
    AppendBaseFragMetaRequest, AppendDeltaFragMetaRequest, AppendFragMetaResponse, ErrorCode,
    FinishCompactFragRequest, FinishImportPartitionBuildRequest, FragMeta, GetIndexMetaRequest,
    GetIndexMetaResponse, GrpcMetaServiceInternalClient, KeyRange, ReaderNodeHeartbeatRequest,
    ReaderNodeHeartbeatResponse, WorkerNodeHeartbeatRequest, WorkerNodeHeartbeatResponse,
};
use tokio::sync::Mutex;

use crate::utils::{metrics::begin_meta_client_metrics, shard_tracing_id};

pub type MetaClientNotLeaderHandler =
    Arc<dyn Fn(String) -> BoxFuture<'static, Result<()>> + Send + Sync>;

fn mark_rpc_metrics_error_on_transport_failure<T, E>(
    rpc_metrics: &OperationMetricsGuard,
    rpc_result: &std::result::Result<T, E>,
) {
    if rpc_result.is_err() {
        rpc_metrics.mark_error();
    }
}

fn mark_rpc_metrics_error_if_non_success(rpc_metrics: &OperationMetricsGuard, status: i32) {
    if status != ErrorCode::Success as i32 {
        rpc_metrics.mark_error();
    }
}

#[async_trait::async_trait]
pub trait MetaClient: Send + Sync {
    fn get_addr(&self) -> &str;

    async fn append_base_frag(
        &mut self,
        req: AppendBaseFragMetaRequest,
    ) -> Result<AppendFragMetaResponse>;

    async fn append_delta_frag(
        &mut self,
        req: AppendDeltaFragMetaRequest,
    ) -> Result<AppendFragMetaResponse>;

    async fn get_frags_by_keyrange(
        &mut self,
        shard_id: u64,
        index_key: &IndexKey,
        key_range: KeyRange,
    ) -> Result<(Vec<FragMeta>, u64)>;

    async fn finish_compact_frags(&mut self, req: FinishCompactFragRequest) -> Result<()>;

    #[allow(clippy::too_many_arguments)]
    async fn finish_partition_build(
        &mut self,
        index_key: IndexKey,
        task_id: u64,
        owner_addr: String,
        owner_seq: u64,
        err_msg: Option<String>,
        rows_processed: u64,
        bytes_processed: u64,
        input_file_size: u64,
    ) -> Result<()>;

    async fn reader_node_heartbeat(
        &mut self,
        req: ReaderNodeHeartbeatRequest,
    ) -> Result<ReaderNodeHeartbeatResponse>;

    async fn worker_node_heartbeat(
        &mut self,
        req: WorkerNodeHeartbeatRequest,
    ) -> Result<Vec<WorkerNodeHeartbeatResponse>>;

    async fn get_shard_progress(&mut self, shard_id: u64) -> Result<Option<String>>;

    async fn reset_shard_progress(&mut self, _shard_id: u64) -> Result<()>;

    async fn get_index_meta(&mut self, index_key: &IndexKey) -> Result<GetIndexMetaResponse>;
}

#[derive(Clone)]
pub struct GrpcMetaClient {
    addr: String,
    client: GrpcMetaServiceInternalClient<tonic::transport::Channel>,
    not_leader_handler: Option<MetaClientNotLeaderHandler>,
}

impl GrpcMetaClient {
    pub fn new(
        addr: String,
        client: GrpcMetaServiceInternalClient<tonic::transport::Channel>,
    ) -> Self {
        GrpcMetaClient {
            addr,
            client,
            not_leader_handler: None,
        }
    }

    pub fn new_with_not_leader_handler(
        addr: String,
        client: GrpcMetaServiceInternalClient<tonic::transport::Channel>,
        not_leader_handler: MetaClientNotLeaderHandler,
    ) -> Self {
        GrpcMetaClient {
            addr,
            client,
            not_leader_handler: Some(not_leader_handler),
        }
    }

    async fn refresh_leader_if_not_leader(&self, status: i32) {
        if status != ErrorCode::NotLeader as i32 {
            return;
        }

        if let Some(handler) = self.not_leader_handler.as_ref() {
            if let Err(e) = handler(self.addr.clone()).await {
                warn!(
                    "failed to refresh meta leader after NOT_LEADER, stale_addr={}, err={:?}",
                    self.addr, e
                );
            }
        }
    }
}

#[async_trait::async_trait]
impl MetaClient for GrpcMetaClient {
    fn get_addr(&self) -> &str {
        &self.addr
    }

    async fn append_base_frag(
        &mut self,
        req: AppendBaseFragMetaRequest,
    ) -> Result<AppendFragMetaResponse> {
        let rpc_metrics = begin_meta_client_metrics("append_base_frag_meta");
        let rpc_result = self.client.append_base_frag_meta(req).await;
        mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
        let resp = rpc_result?.into_inner();
        mark_rpc_metrics_error_if_non_success(&rpc_metrics, resp.status);
        self.refresh_leader_if_not_leader(resp.status).await;
        Ok(resp)
    }

    async fn append_delta_frag(
        &mut self,
        req: AppendDeltaFragMetaRequest,
    ) -> Result<AppendFragMetaResponse> {
        let rpc_metrics = begin_meta_client_metrics("append_delta_frag_meta");
        let rpc_result = self.client.append_delta_frag_meta(req).await;
        mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
        let resp = rpc_result?.into_inner();
        mark_rpc_metrics_error_if_non_success(&rpc_metrics, resp.status);
        self.refresh_leader_if_not_leader(resp.status).await;
        Ok(resp)
    }

    async fn get_frags_by_keyrange(
        &mut self,
        shard_id: u64,
        index_key: &IndexKey,
        key_range: KeyRange,
    ) -> Result<(Vec<FragMeta>, u64)> {
        #[cfg(feature = "get_manifest_from_writer")]
        {
            let rpc_metrics = begin_meta_client_metrics("debug_get_shard_manifest");
            // debug_get_shard_manifest returns fragment information based on the start key.
            // If the shard splits or merges, the returned fragment may no longer belong to
            // the original shard_id, so verification is required.
            let rpc_result = self
                .client
                .debug_get_shard_manifest(tici_grpc::DebugGetShardManifestRequest {
                    keyspace_id: index_key.keyspace_id,
                    table_id: index_key.table_id,
                    index_id: index_key.index_id,
                    key_range: Some(key_range.clone()),
                })
                .await
                .with_context(|| shard_tracing_id(index_key, shard_id));
            mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
            let resp = rpc_result?.into_inner();
            if resp.status != ErrorCode::Success as i32 {
                rpc_metrics.mark_error();
                self.refresh_leader_if_not_leader(resp.status).await;
                bail!(
                    "{}, key_range: [{}, {}), invalid status={}.",
                    shard_tracing_id(index_key, shard_id),
                    hex::encode(key_range.start_key),
                    hex::encode(key_range.end_key),
                    resp.status
                );
            }
            if resp.shard_id != shard_id {
                rpc_metrics.mark_error();
                bail!(
                    "{}, key_range: [{}, {}), shard id not match, shard_id={}.",
                    shard_tracing_id(index_key, shard_id),
                    hex::encode(key_range.start_key),
                    hex::encode(key_range.end_key),
                    resp.shard_id
                );
            }
            Ok((resp.frag_metas, resp.seq))
        }
        #[cfg(not(feature = "get_manifest_from_writer"))]
        {
            let rpc_metrics = begin_meta_client_metrics("get_shard_manifest");
            let rpc_result = self
                .client
                .get_shard_manifest(tici_grpc::GetShardManifestRequest {
                    keyspace_id: index_key.keyspace_id,
                    table_id: index_key.table_id,
                    index_id: index_key.index_id,
                    shard_id,
                })
                .await
                .with_context(|| shard_tracing_id(index_key, shard_id));
            mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
            let resp = rpc_result?.into_inner();
            if resp.status != ErrorCode::Success as i32 {
                rpc_metrics.mark_error();
                self.refresh_leader_if_not_leader(resp.status).await;
                bail!(
                    "{}, key_range: [{}, {}), invalid status={}.",
                    shard_tracing_id(index_key, shard_id),
                    hex::encode(key_range.start_key),
                    hex::encode(key_range.end_key),
                    resp.status
                );
            }
            Ok((resp.frag_metas, resp.seq))
        }
    }

    async fn finish_compact_frags(&mut self, req: FinishCompactFragRequest) -> Result<()> {
        let rpc_metrics = begin_meta_client_metrics("finish_compact_fragments");
        let rpc_result = self
            .client
            .finish_compact_fragments(req.clone())
            .await
            .with_context(|| format!("FinishCompactFragRequest={:?}", req));
        mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
        let resp = rpc_result?.into_inner();
        if resp.status != ErrorCode::Success as i32 {
            // If status = 1, the region is splitted.
            rpc_metrics.mark_error();
            self.refresh_leader_if_not_leader(resp.status).await;
            Err(anyhow::anyhow!(
                "Finish compact frag failed, status={}. FinishCompactFragRequest={:?}",
                resp.status,
                req
            ))
        } else {
            Ok(())
        }
    }

    async fn finish_partition_build(
        &mut self,
        index_key: IndexKey,
        task_id: u64,
        owner_addr: String,
        owner_seq: u64,
        err_msg: Option<String>,
        rows_processed: u64,
        bytes_processed: u64,
        input_file_size: u64,
    ) -> Result<()> {
        let rpc_metrics = begin_meta_client_metrics("finish_import_partition_build");
        let req = if let Some(msg) = err_msg {
            FinishImportPartitionBuildRequest {
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_id: index_key.index_id,
                task_id,
                status: 1, // 1 means failed
                frag_metas: vec![],
                error_message: msg,
                rows_processed,
                bytes_processed,
                input_file_size,
                owner_addr,
                owner_seq,
            }
        } else {
            FinishImportPartitionBuildRequest {
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_id: index_key.index_id,
                task_id,
                status: 0, // 0 means success
                frag_metas: vec![],
                error_message: String::new(),
                rows_processed,
                bytes_processed,
                input_file_size,
                owner_addr,
                owner_seq,
            }
        };
        let rpc_result = self
            .client
            .finish_import_partition_build(req.clone())
            .await
            .with_context(|| format!("FinishPartitionRequest={:?}", req));
        mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
        let resp = rpc_result?.into_inner();
        if resp.status != ErrorCode::Success as i32 {
            rpc_metrics.mark_error();
            self.refresh_leader_if_not_leader(resp.status).await;
            Err(anyhow::anyhow!(
                "Finish partition build failed, status={}. FinishPartitionRequest={:?}",
                resp.status,
                req
            ))
        } else {
            Ok(())
        }
    }

    async fn reader_node_heartbeat(
        &mut self,
        req: ReaderNodeHeartbeatRequest,
    ) -> Result<ReaderNodeHeartbeatResponse> {
        let rpc_metrics = begin_meta_client_metrics("reader_node_heartbeat");
        let rpc_result = self.client.reader_node_heartbeat(req).await;
        mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
        let resp = rpc_result?.into_inner();
        mark_rpc_metrics_error_if_non_success(&rpc_metrics, resp.status);
        self.refresh_leader_if_not_leader(resp.status).await;
        Ok(resp)
    }

    async fn worker_node_heartbeat(
        &mut self,
        req: WorkerNodeHeartbeatRequest,
    ) -> Result<Vec<WorkerNodeHeartbeatResponse>> {
        let rpc_metrics = begin_meta_client_metrics("worker_node_heartbeat");
        let rpc_result = self.client.worker_node_heartbeat(req).await;
        mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
        let mut stream = rpc_result?.into_inner();
        let mut responses = Vec::new();
        while let Some(resp) = stream.message().await? {
            mark_rpc_metrics_error_if_non_success(&rpc_metrics, resp.status);
            self.refresh_leader_if_not_leader(resp.status).await;
            responses.push(resp);
        }
        if responses.is_empty() {
            rpc_metrics.mark_error();
            bail!("worker_node_heartbeat returned an empty stream");
        }
        Ok(responses)
    }

    async fn get_shard_progress(&mut self, _shard_id: u64) -> Result<Option<String>> {
        unimplemented!("GrpcMetaClient::get_shard_progress")
    }

    async fn reset_shard_progress(&mut self, _shard_id: u64) -> Result<()> {
        unimplemented!("GrpcMetaClient::reset_shard_progress")
    }

    async fn get_index_meta(&mut self, index_key: &IndexKey) -> Result<GetIndexMetaResponse> {
        let rpc_metrics = begin_meta_client_metrics("get_index_meta");
        let req = GetIndexMetaRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
        };
        let rpc_result = self
            .client
            .get_index_meta(req)
            .await
            .with_context(|| format!("get_index_meta index_key={:?}", index_key));
        mark_rpc_metrics_error_on_transport_failure(&rpc_metrics, &rpc_result);
        let resp = rpc_result?.into_inner();
        mark_rpc_metrics_error_if_non_success(&rpc_metrics, resp.status);
        self.refresh_leader_if_not_leader(resp.status).await;
        Ok(resp)
    }
}

#[derive(Default)]
struct MockShardMeta {
    frags: Vec<FragMeta>,
    seq: u64,
}

impl MockShardMeta {
    pub fn append_delta_frag(&mut self, frag_meta: FragMeta) {
        self.frags.push(frag_meta);
        self.seq += 1;
    }

    pub fn append_base_frag(&mut self, frag_meta: FragMeta) {
        self.frags.insert(0, frag_meta);
        self.seq += 1;
    }
}

#[derive(Clone)]
pub struct MockMetaClient {
    // TODO: Vec<FragMeta> to manifest.
    // ShardID -> MockShardMeta
    addr: String,
    shard_frags: Arc<Mutex<HashMap<u64, MockShardMeta>>>,
    shard_progress: Arc<Mutex<HashMap<u64, String>>>,
    finish_partition_build_failures_before_success: Arc<AtomicUsize>,
    finish_partition_build_calls: Arc<AtomicUsize>,
}

impl MockMetaClient {
    pub fn new() -> Self {
        MockMetaClient {
            addr: "mock_meta_client".to_string(),
            shard_frags: Arc::new(Mutex::new(HashMap::new())),
            shard_progress: Arc::new(Mutex::new(HashMap::new())),
            finish_partition_build_failures_before_success: Arc::new(AtomicUsize::new(0)),
            finish_partition_build_calls: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub fn debug_set_finish_partition_build_failures_before_success(&self, failures: usize) {
        self.finish_partition_build_failures_before_success
            .store(failures, Ordering::SeqCst);
        self.finish_partition_build_calls.store(0, Ordering::SeqCst);
    }

    pub fn debug_finish_partition_build_call_count(&self) -> usize {
        self.finish_partition_build_calls.load(Ordering::SeqCst)
    }
}

impl Default for MockMetaClient {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait::async_trait]
impl MetaClient for MockMetaClient {
    fn get_addr(&self) -> &str {
        &self.addr
    }

    async fn append_base_frag(
        &mut self,
        req: AppendBaseFragMetaRequest,
    ) -> Result<AppendFragMetaResponse> {
        let mut shard_frags = self.shard_frags.lock().await;
        // insert into all shards.
        for meta in shard_frags.values_mut() {
            meta.append_base_frag(req.frag_metas[0].clone());
        }
        Ok(AppendFragMetaResponse {
            status: tici_grpc::ErrorCode::Success.into(),
        })
    }

    async fn append_delta_frag(
        &mut self,
        req: AppendDeltaFragMetaRequest,
    ) -> Result<AppendFragMetaResponse> {
        let shard_id = req.shard_id;
        if !req.frag_metas.is_empty() {
            let mut shard_frags = self.shard_frags.lock().await;
            let meta = shard_frags.entry(shard_id).or_default();
            meta.append_delta_frag(req.frag_metas[0].clone());
        }
        if let Some(progress) = req.progress {
            let mut shard_progress = self.shard_progress.lock().await;
            shard_progress.insert(shard_id, progress.last_cdc_file);
        }
        Ok(AppendFragMetaResponse {
            status: tici_grpc::ErrorCode::Success.into(),
        })
    }

    async fn get_frags_by_keyrange(
        &mut self,
        shard_id: u64,
        _index_key: &IndexKey,
        _key_range: KeyRange,
    ) -> Result<(Vec<FragMeta>, u64)> {
        let shard_frags = self.shard_frags.lock().await;
        if let Some(meta) = shard_frags.get(&shard_id) {
            Ok((meta.frags.clone(), meta.seq))
        } else {
            Ok((vec![], 0))
        }
    }

    async fn finish_compact_frags(&mut self, _req: FinishCompactFragRequest) -> Result<()> {
        unimplemented!()
    }

    async fn finish_partition_build(
        &mut self,
        _index_key: IndexKey,
        _task_id: u64,
        _owner_addr: String,
        _owner_seq: u64,
        _err_msg: Option<String>,
        _rows_processed: u64,
        _bytes_processed: u64,
        _input_file_size: u64,
    ) -> Result<()> {
        let call = self
            .finish_partition_build_calls
            .fetch_add(1, Ordering::SeqCst)
            + 1;
        if call
            <= self
                .finish_partition_build_failures_before_success
                .load(Ordering::SeqCst)
        {
            bail!("transient finish failure");
        }
        Ok(())
    }

    async fn reader_node_heartbeat(
        &mut self,
        _req: ReaderNodeHeartbeatRequest,
    ) -> Result<ReaderNodeHeartbeatResponse> {
        unimplemented!("MockMetaClient::reader_node_heartbeat")
    }

    async fn worker_node_heartbeat(
        &mut self,
        _req: WorkerNodeHeartbeatRequest,
    ) -> Result<Vec<WorkerNodeHeartbeatResponse>> {
        unimplemented!("MockMetaClient::worker_node_heartbeat")
    }

    async fn get_shard_progress(&mut self, shard_id: u64) -> Result<Option<String>> {
        Ok(self.shard_progress.lock().await.get(&shard_id).cloned())
    }

    async fn reset_shard_progress(&mut self, shard_id: u64) -> Result<()> {
        let mut shard_progress = self.shard_progress.lock().await;
        shard_progress.remove(&shard_id);
        Ok(())
    }

    async fn get_index_meta(&mut self, _index_key: &IndexKey) -> Result<GetIndexMetaResponse> {
        unimplemented!("MockMetaClient::get_index_meta")
    }
}