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

use anyhow::{Context, Result, anyhow, bail};
use backon::{ConstantBuilder, Retryable};
use log::{error, info, warn};
use tici_common::config::GrpcConfig;
use tici_grpc::GrpcMetaServiceInternalClient;
use tici_service_discovery::{EtcdClient, META_SERVICE_ELECTION_KEY, election::leader::LeaderTerm};
use tonic::transport::Channel;

use crate::utils::{
    meta_client::{GrpcMetaClient, MetaClient, MetaClientNotLeaderHandler, MockMetaClient},
    metrics::inc_meta_client_pool_failed_attempts,
};

const CONNECT_MAX_RETRIES: u64 = 5;
const TCP_KEEPALIVE: Duration = Duration::from_secs(15);
const HTTP2_KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(10);
const HTTP2_KEEP_ALIVE_TIMEOUT: Duration = Duration::from_secs(3);
const HTTP2_KEEP_ALIVE_IDLE: bool = true;

// =================================================================================
// MetaClientPool Trait Definition
// =================================================================================
#[async_trait::async_trait]
pub trait MetaClientPool: Send + Sync {
    // Returns a valid `MetaClient` instance (boxed trait object) and the current
    // leader address.
    async fn get_meta_client(&self) -> Result<Box<dyn MetaClient>>;
}

// =================================================================================
// GrpcMetaClientPool Implementation
// =================================================================================

#[derive(Clone)]
struct MetaChannel {
    addr: String,
    channel: Channel,
}

// Manages the gRPC connection to the Meta Service Leader.
//
// `GrpcMetaClientPool` implements a thread-safe singleton-like pattern that:
// 1. Caches the `Channel` (connection) to the current leader.
// 2. Automatically invalidates the cache and reconnects if the leader changes.
pub struct GrpcMetaClientPool {
    leader_term: Arc<LeaderTerm>,
    // We cache the Channel, not the Client. The Client is lightweight and
    // created on demand, allowing for flexible configuration (e.g., max message size).
    meta_channel: Arc<tokio::sync::RwLock<Option<MetaChannel>>>,
    grpc_config: GrpcConfig,
}

impl GrpcMetaClientPool {
    pub async fn new(
        pd_addresses: &[String],
        sec_config: &tici_common::security::SecurityConfig,
        grpc_config: GrpcConfig,
        wait_timeout: i32,
    ) -> Result<GrpcMetaClientPool> {
        let etcd_client = Arc::new(EtcdClient::new_tls(pd_addresses, sec_config)?);

        async fn start_leader_term_retry(etcd_client: Arc<EtcdClient>) -> Result<LeaderTerm> {
            let mut leader_term = LeaderTerm::new(META_SERVICE_ELECTION_KEY, etcd_client);
            leader_term.start().await.map_err(|e| anyhow!("{:?}", e))?;
            Ok(leader_term)
        }

        let delay: u64 = 2;
        let mut builder = ConstantBuilder::default().with_delay(Duration::from_secs(delay));
        if wait_timeout >= 0 {
            info!(
                "Waiting for meta service leader (timeout: {:?})",
                wait_timeout
            );
            builder = builder.with_max_times((wait_timeout as u32 / delay as u32) as usize);
        } else {
            info!("Waiting for meta service leader (timeout: forever)");
            builder = builder.without_max_times();
        }
        let retry_operation = (move || start_leader_term_retry(etcd_client.clone()))
            .retry(&builder)
            .notify(|_, dur| warn!("Retrying in {:?}...", dur));

        let leader_term = retry_operation
            .await
            .inspect_err(|e| error!("Failed to start leader term: {:?}", e))?;

        Ok(GrpcMetaClientPool {
            leader_term: Arc::new(leader_term),
            meta_channel: Arc::new(tokio::sync::RwLock::new(None)),
            grpc_config,
        })
    }

    fn not_leader_handler(&self) -> MetaClientNotLeaderHandler {
        let leader_term = self.leader_term.clone();
        let meta_channel = self.meta_channel.clone();
        Arc::new(move |stale_addr| {
            let leader_term = leader_term.clone();
            let meta_channel = meta_channel.clone();
            Box::pin(async move {
                let should_refresh = {
                    let mut guard = meta_channel.write().await;
                    match guard.as_ref() {
                        Some(channel) if channel.addr == stale_addr => {
                            guard.take();
                            true
                        }
                        Some(_) => false,
                        None => true,
                    }
                };

                if should_refresh {
                    let refreshed_addr = leader_term.refresh_leader_address().await?;
                    info!(
                        "refreshed meta leader after NOT_LEADER, stale_addr={}, refreshed_addr={}",
                        stale_addr, refreshed_addr
                    );
                }
                Ok(())
            })
        })
    }

    // Helper to get the channel with retry logic.
    async fn get_meta_channel_with_retry(&self, max_retries: u64) -> Result<MetaChannel> {
        for i in 0..max_retries {
            match self.get_meta_channel().await {
                Ok(meta_channel) => return Ok(meta_channel),
                Err(e) => {
                    inc_meta_client_pool_failed_attempts();
                    warn!(
                        "Attempt {}/{} => get meta channel failed, err={:?}",
                        i + 1,
                        max_retries,
                        e
                    );
                    // Linear backoff: 1s, 2s, 3s...
                    tokio::time::sleep(Duration::from_secs(i + 1)).await;
                }
            }
        }
        bail!("get meta channel failed after {} times", max_retries);
    }

    // Core logic to resolve leader address and establish/reuse connection.
    async fn get_meta_channel(&self) -> Result<MetaChannel> {
        let meta_leader_addr = self
            .leader_term
            .get_leader_address()
            .await
            .with_context(|| "get meta leader address")?;

        // Fast path: reuse existing channel if the leader hasn't changed.
        if let Some(meta_channel) = self.meta_channel.read().await.clone() {
            if meta_channel.addr == meta_leader_addr {
                return Ok(meta_channel);
            }
        }

        // meta_channel is None or addr changed, try to connect to the new leader.
        // Acquire write lock to prevent multiple tasks from connecting simultaneously.
        let mut guard = self.meta_channel.write().await;

        // Double check inside the write lock.
        if let Some(meta_channel) = guard.clone() {
            if meta_channel.addr == meta_leader_addr {
                // Another task has already connected to the new leader.
                return Ok(meta_channel);
            }
        }

        info!("get_meta_channel, connecting to addr={}", meta_leader_addr);

        let endpoint = Channel::from_shared(format!("http://{}", meta_leader_addr))
            .with_context(|| format!("invalid uri: {}", meta_leader_addr))?
            .connect_timeout(self.grpc_config.connect_timeout.into())
            .tcp_keepalive(Some(TCP_KEEPALIVE))
            .http2_keep_alive_interval(HTTP2_KEEP_ALIVE_INTERVAL)
            .keep_alive_timeout(HTTP2_KEEP_ALIVE_TIMEOUT)
            .keep_alive_while_idle(HTTP2_KEEP_ALIVE_IDLE)
            .timeout(self.grpc_config.request_timeout.into());

        // Explicitly connect to ensure the network is reachable before returning.
        let channel = endpoint
            .connect()
            .await
            .with_context(|| format!("connect to meta leader failed, addr={}", meta_leader_addr))?;

        info!(
            "get_meta_channel, connected to addr={}, grpc_config={:?}",
            meta_leader_addr, self.grpc_config
        );

        let meta_channel = MetaChannel {
            addr: meta_leader_addr,
            channel,
        };

        guard.replace(meta_channel.clone());
        Ok(meta_channel)
    }
}

#[async_trait::async_trait]
impl MetaClientPool for GrpcMetaClientPool {
    async fn get_meta_client(&self) -> Result<Box<dyn MetaClient>> {
        let meta_channel = self
            .get_meta_channel_with_retry(CONNECT_MAX_RETRIES)
            .await?;

        // Create the internal tonic client using the cached channel.
        let internal_client = GrpcMetaServiceInternalClient::new(meta_channel.channel)
            .max_decoding_message_size(self.grpc_config.max_decoding_message_size.0 as usize)
            .max_encoding_message_size(self.grpc_config.max_encoding_message_size.0 as usize);

        Ok(Box::new(GrpcMetaClient::new_with_not_leader_handler(
            meta_channel.addr,
            internal_client,
            self.not_leader_handler(),
        )))
    }
}

// =================================================================================
// MockMetaClientPool Implementation
// =================================================================================
pub struct MockMetaClientPool {
    mock_client: MockMetaClient,
}

impl MockMetaClientPool {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            mock_client: MockMetaClient::new(),
        }
    }
}

#[async_trait::async_trait]
impl MetaClientPool for MockMetaClientPool {
    async fn get_meta_client(&self) -> Result<Box<dyn MetaClient>> {
        // MockMetaClient is internally thread-safe (using Arc/Mutex),
        // so we can clone it and box it.
        Ok(Box::new(self.mock_client.clone()))
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use backon::{BackoffBuilder, ConstantBuilder};

    // Regression test: verifies that wait_timeout correctly controls the retry
    // count in GrpcMetaClientPool::new(). Previously, the return value of
    // ConstantBuilder::with_max_times() was discarded (ConstantBuilder is Copy),
    // causing retry count to always be 3 regardless of wait_timeout.
    #[test]
    fn test_wait_timeout_retry_count() {
        let delay: u64 = 2;
        let wait_timeout: i32 = 10;

        // Mirrors the pattern in GrpcMetaClientPool::new()
        let mut builder = ConstantBuilder::default().with_delay(Duration::from_secs(delay));
        if wait_timeout >= 0 {
            builder = builder.with_max_times((wait_timeout as u32 / delay as u32) as usize);
        } else {
            builder = builder.without_max_times();
        }

        let count = builder.build().count();
        assert_eq!(
            count,
            5,
            "retry count should be wait_timeout/delay={}",
            wait_timeout as u64 / delay,
        );
    }

    #[test]
    fn test_wait_timeout_unlimited() {
        let delay: u64 = 2;
        let wait_timeout: i32 = -1;

        let mut builder = ConstantBuilder::default().with_delay(Duration::from_secs(delay));
        if wait_timeout >= 0 {
            builder = builder.with_max_times((wait_timeout as u32 / delay as u32) as usize);
        } else {
            builder = builder.without_max_times();
        }

        // When wait_timeout < 0, retry should be unlimited
        let count = builder.build().take(10_000).count();
        assert_eq!(count, 10_000, "unlimited retry should not stop");
    }
}