use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicI32, Ordering};
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};

use futures::{Stream, future::join_all};
use pd_proto_grpc::{pdpb, tsopb};
use tokio::sync::RwLock;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Code, Request, Response, Status, Streaming};
use tracing::{Instrument, warn};

use crate::error::ServiceError;
use crate::forwarding::{
    ForwardChannelPool, LEGACY_FORWARD_FAILPOINT_RECV_FROM_TSO_TIMEOUT,
    LEGACY_FORWARD_FAILPOINT_SEND_TO_TSO_TIMEOUT, LEGACY_FORWARD_PROXY_TIMEOUT,
    LEGACY_PROXY_MAX_CONCURRENT_STREAMINGS_EXCEEDED_TOKEN, LEGACY_PROXY_NOT_FOUND_TSO_ADDR_TOKEN,
    LEGACY_PROXY_RECV_FROM_CLIENT_EOF_TOKEN, LEGACY_PROXY_RECV_FROM_CLIENT_TIMEOUT_TOKEN,
    LegacyForwardDispatchRequest, LegacyForwardDispatcherPool, LegacyForwardFailpoint,
    dispatch_host_key, fetch_min_ts_from_host, forwarded_channel_endpoint,
    legacy_forward_timeout_status, normalize_endpoint, stable_member_id,
};
use crate::grpc_codec::{
    from_pdpb_tso_request, from_tsopb_request_header, from_tsopb_tso_request, status_from_core,
    to_pdpb_error, to_pdpb_tso_response, to_tsopb_find_group_response,
    to_tsopb_get_min_ts_response, to_tsopb_tso_response,
};
use crate::metrics;
use crate::service::{TsoRequestRouteCache, TsoService};
use crate::trace_context::{TraceContext, rpc_span};

type StreamResult<T> = Result<T, Status>;
type TsoResponseStream = Pin<Box<dyn Stream<Item = StreamResult<tsopb::TsoResponse>> + Send>>;
type PdResponseStream<T> = Pin<Box<dyn Stream<Item = StreamResult<T>> + Send>>;
type LegacyTsoResponseStream = PdResponseStream<pdpb::TsoResponse>;

/// Metadata key used by strict compatibility tests to force one legacy forwarding watchdog path.
pub const LEGACY_FORWARD_FAILPOINT_METADATA_KEY: &str = "x-pd-legacy-forward-failpoint";
/// Metadata value that forces send-side watchdog timeout handling (`Canceled` surface).
pub const LEGACY_FORWARD_FAILPOINT_SEND_TIMEOUT_VALUE: &str =
    LEGACY_FORWARD_FAILPOINT_SEND_TO_TSO_TIMEOUT;
/// Metadata value that forces recv-side watchdog timeout handling (`Canceled` surface).
pub const LEGACY_FORWARD_FAILPOINT_RECV_TIMEOUT_VALUE: &str =
    LEGACY_FORWARD_FAILPOINT_RECV_FROM_TSO_TIMEOUT;

fn join_error_to_status(err: tokio::task::JoinError) -> Status {
    Status::internal(format!("blocking task failed: {err}"))
}

/// gRPC adapter exposing the core TSO service through real `tsopb` protobuf transport.
#[derive(Clone)]
pub struct TsoGrpcService {
    inner: Arc<RwLock<TsoService>>,
}

impl TsoGrpcService {
    /// Creates one `tsopb.TSO` gRPC adapter around shared service state.
    #[must_use]
    pub fn new(inner: Arc<RwLock<TsoService>>) -> Self {
        Self { inner }
    }
}

#[tonic::async_trait]
impl tsopb::tso_server::Tso for TsoGrpcService {
    type TsoStream = TsoResponseStream;

    async fn tso(
        &self,
        request: Request<Streaming<tsopb::TsoRequest>>,
    ) -> Result<Response<Self::TsoStream>, Status> {
        let trace_context = TraceContext::extract(request.metadata());
        let span = rpc_span("tsopb.TSO/Tso", &trace_context);
        let mut incoming = request.into_inner();
        let executor = {
            let guard = self.inner.read().await;
            guard.tso_executor()
        };
        let (tx, rx) = mpsc::channel::<StreamResult<tsopb::TsoResponse>>(32);

        tokio::spawn(
            async move {
                let mut route_cache = TsoRequestRouteCache::default();
                loop {
                    let next = incoming.message().await;
                    let request = match next {
                        Ok(Some(req)) => req,
                        Ok(None) => break,
                        Err(status) => {
                            let _ = tx.send(Err(status)).await;
                            break;
                        }
                    };

                    let core_request = from_tsopb_tso_request(request);
                    let response = executor.handle_tso(now_ms(), core_request, &mut route_cache);

                    match response {
                        Ok(resp) => {
                            if tx.send(Ok(to_tsopb_tso_response(resp))).await.is_err() {
                                break;
                            }
                        }
                        Err(status) => {
                            let _ = tx.send(Err(status_from_core(status))).await;
                            break;
                        }
                    }
                }
            }
            .instrument(span),
        );

        Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
    }

    async fn find_group_by_keyspace_id(
        &self,
        request: Request<tsopb::FindGroupByKeyspaceIdRequest>,
    ) -> Result<Response<tsopb::FindGroupByKeyspaceIdResponse>, Status> {
        let trace_context = TraceContext::extract(request.metadata());
        let span = rpc_span("tsopb.TSO/FindGroupByKeyspaceID", &trace_context);
        let request = request.into_inner();
        let service = Arc::clone(&self.inner);
        let response = tokio::task::spawn_blocking(move || {
            let guard = service.blocking_read();
            guard.find_group_by_keyspace_id(
                now_ms(),
                pd_proto_tso::FindGroupByKeyspaceIdRequest {
                    header: from_tsopb_request_header(request.header),
                    keyspace_id: request.keyspace_id,
                    mod_revision: request.mod_revision,
                },
            )
        })
        .instrument(span)
        .await
        .map_err(join_error_to_status)?;
        Ok(Response::new(to_tsopb_find_group_response(response)))
    }

    async fn get_min_ts(
        &self,
        request: Request<tsopb::GetMinTsRequest>,
    ) -> Result<Response<tsopb::GetMinTsResponse>, Status> {
        let trace_context = TraceContext::extract(request.metadata());
        let span = rpc_span("tsopb.TSO/GetMinTS", &trace_context);
        let request = request.into_inner();
        let service = Arc::clone(&self.inner);
        let response = tokio::task::spawn_blocking(move || {
            let guard = service.blocking_read();
            guard.get_min_ts(
                now_ms(),
                pd_proto_tso::GetMinTsRequest {
                    header: from_tsopb_request_header(request.header),
                    dc_location: request.dc_location,
                },
            )
        })
        .instrument(span)
        .await
        .map_err(join_error_to_status)?;
        Ok(Response::new(to_tsopb_get_min_ts_response(response)))
    }
}

/// Legacy `pdpb.PD` transport surface used by migration compatibility clients.
#[derive(Clone)]
pub struct LegacyPdGrpcService {
    inner: Arc<RwLock<TsoService>>,
    forward_pool: ForwardChannelPool,
    forward_dispatchers: LegacyForwardDispatcherPool,
    concurrent_legacy_proxy_streamings: Arc<AtomicI32>,
}

struct LegacyForwardStreamPermit {
    active_streams: Option<Arc<AtomicI32>>,
}

impl LegacyForwardStreamPermit {
    fn unbounded() -> Self {
        Self {
            active_streams: None,
        }
    }

    fn bounded(active_streams: Arc<AtomicI32>) -> Self {
        Self {
            active_streams: Some(active_streams),
        }
    }
}

impl Drop for LegacyForwardStreamPermit {
    fn drop(&mut self) {
        if let Some(active_streams) = &self.active_streams {
            active_streams.fetch_sub(1, Ordering::AcqRel);
        }
    }
}

impl LegacyPdGrpcService {
    /// Creates one legacy `pdpb.PD` compatibility adapter.
    #[must_use]
    pub fn new(inner: Arc<RwLock<TsoService>>) -> Self {
        Self {
            inner,
            forward_pool: ForwardChannelPool::default(),
            forward_dispatchers: LegacyForwardDispatcherPool::default(),
            concurrent_legacy_proxy_streamings: Arc::new(AtomicI32::new(0)),
        }
    }

    async fn should_forward_to_host(&self, forwarded_host: &str) -> bool {
        if forwarded_host.is_empty() {
            return false;
        }
        let mut local_urls = legacy_local_tso_urls();
        if local_urls.is_empty() {
            let guard = self.inner.read().await;
            local_urls = guard.manager().local_discovery_tso_urls();
        }
        !local_urls.iter().any(|local| local == forwarded_host)
    }

    async fn dispatch_legacy_forwarded_tso(
        &self,
        forwarded_host: &str,
        request: pdpb::TsoRequest,
        trace_context: TraceContext,
        failpoint: Option<LegacyForwardFailpoint>,
        idle_timeout: Duration,
    ) -> Result<pdpb::TsoResponse, Status> {
        if normalize_endpoint(forwarded_host).is_empty() {
            return Err(Status::not_found(LEGACY_PROXY_NOT_FOUND_TSO_ADDR_TOKEN));
        }
        let endpoint = forwarded_channel_endpoint(forwarded_host);
        let dispatch_host = dispatch_host_key(forwarded_host, failpoint);
        let dispatch_response_timeout =
            LEGACY_FORWARD_PROXY_TIMEOUT.saturating_add(Duration::from_millis(250));
        for _ in 0..2 {
            let queue = match self
                .forward_dispatchers
                .queue_for_host(&dispatch_host, &endpoint, &self.forward_pool, idle_timeout)
                .await
            {
                Ok(queue) => queue,
                Err(status) => {
                    if status.code() == Code::Internal
                        && status.message().starts_with("invalid endpoint:")
                    {
                        return Err(Status::not_found(LEGACY_PROXY_NOT_FOUND_TSO_ADDR_TOKEN));
                    }
                    return Err(status);
                }
            };
            let (response_tx, response_rx) = oneshot::channel();
            let dispatch = LegacyForwardDispatchRequest {
                request: request.clone(),
                trace_context: trace_context.clone(),
                failpoint,
                response_tx,
            };
            let enqueue_result =
                tokio::time::timeout(LEGACY_FORWARD_PROXY_TIMEOUT, queue.send(dispatch)).await;
            match enqueue_result {
                Ok(Ok(())) => {}
                Ok(Err(_)) => {
                    self.forward_dispatchers
                        .remove_if_matches(&dispatch_host, &queue)
                        .await;
                    continue;
                }
                Err(_) => {
                    return Err(legacy_forward_timeout_status());
                }
            }
            let response_result =
                tokio::time::timeout(dispatch_response_timeout, response_rx).await;
            match response_result {
                Ok(Ok(response)) => return response,
                Ok(Err(_)) => {
                    self.forward_dispatchers
                        .remove_if_matches(&dispatch_host, &queue)
                        .await;
                    continue;
                }
                Err(_) => {
                    return Err(legacy_forward_timeout_status());
                }
            }
        }
        Err(Status::unavailable(
            "legacy forwarding dispatcher unavailable",
        ))
    }

    fn try_acquire_legacy_proxy_stream(
        &self,
        max_concurrent: i32,
    ) -> Result<LegacyForwardStreamPermit, Box<Status>> {
        if max_concurrent < 0 {
            return Ok(LegacyForwardStreamPermit::unbounded());
        }
        let current = self
            .concurrent_legacy_proxy_streamings
            .fetch_add(1, Ordering::AcqRel)
            + 1;
        if current > max_concurrent {
            self.concurrent_legacy_proxy_streamings
                .fetch_sub(1, Ordering::AcqRel);
            return Err(Box::new(Status::resource_exhausted(
                LEGACY_PROXY_MAX_CONCURRENT_STREAMINGS_EXCEEDED_TOKEN,
            )));
        }
        Ok(LegacyForwardStreamPermit::bounded(Arc::clone(
            &self.concurrent_legacy_proxy_streamings,
        )))
    }

    async fn legacy_proxy_runtime_knobs(&self) -> (i32, std::time::Duration, std::time::Duration) {
        let guard = self.inner.read().await;
        let runtime = guard.runtime_config();
        (
            runtime.max_concurrent_tso_proxy_streamings,
            runtime.tso_proxy_recv_from_client_timeout,
            runtime.tso_proxy_stream_idle_timeout,
        )
    }

    async fn recv_forwarded_legacy_request(
        incoming: &mut Streaming<pdpb::TsoRequest>,
        recv_timeout: std::time::Duration,
    ) -> Result<Option<pdpb::TsoRequest>, Status> {
        match tokio::time::timeout(recv_timeout, incoming.message()).await {
            Ok(Ok(Some(request))) => Ok(Some(request)),
            Ok(Ok(None)) => Err(Status::unknown(LEGACY_PROXY_RECV_FROM_CLIENT_EOF_TOKEN)),
            Ok(Err(status)) => Err(status),
            Err(_) => Err(Status::deadline_exceeded(
                LEGACY_PROXY_RECV_FROM_CLIENT_TIMEOUT_TOKEN,
            )),
        }
    }
}

impl LegacyPdGrpcService {
    async fn aggregate_get_min_ts(
        &self,
        request: pdpb::GetMinTsRequest,
        trace_context: TraceContext,
    ) -> pdpb::GetMinTsResponse {
        let header = request.header.unwrap_or_default();
        let fallback_cluster_id = {
            let guard = self.inner.read().await;
            guard.cluster_id()
        };
        let response_cluster_id = if header.cluster_id == 0 {
            fallback_cluster_id
        } else {
            header.cluster_id
        };

        let snapshot = legacy_discovery_snapshot(&self.inner).await;
        if snapshot.tso_urls.is_empty() {
            return pdpb_get_min_ts_error(
                response_cluster_id,
                "GetMinTS fan-out discovery returned no TSO URLs".to_string(),
            );
        }

        let mut calls = Vec::with_capacity(snapshot.tso_urls.len());
        for host in snapshot.tso_urls {
            let pool = self.forward_pool.clone();
            let trace_context = trace_context.clone();
            calls.push(async move {
                fetch_min_ts_from_host(
                    pool,
                    host,
                    header.cluster_id,
                    header.sender_id,
                    &trace_context,
                )
                .await
            });
        }

        let mut responses = Vec::new();
        for result in join_all(calls).await {
            match result {
                Ok(response) => responses.push(response),
                Err(status) => {
                    warn!(error = %status, "GetMinTS fan-out call failed for one pod");
                }
            }
        }

        aggregate_fanout_get_min_ts(response_cluster_id, responses)
    }
}

#[tonic::async_trait]
impl pdpb::pd_server::Pd for LegacyPdGrpcService {
    type TsoStream = LegacyTsoResponseStream;
    type RegionHeartbeatStream = PdResponseStream<pdpb::RegionHeartbeatResponse>;
    type QueryRegionStream = PdResponseStream<pdpb::QueryRegionResponse>;
    type WatchGCSafePointV2Stream = PdResponseStream<pdpb::WatchGcSafePointV2Response>;
    type SyncRegionsStream = PdResponseStream<pdpb::SyncRegionResponse>;
    type WatchGlobalConfigStream = PdResponseStream<pdpb::WatchGlobalConfigResponse>;

    async fn get_cluster_info(
        &self,
        request: Request<pdpb::GetClusterInfoRequest>,
    ) -> Result<Response<pdpb::GetClusterInfoResponse>, Status> {
        let trace_context = TraceContext::extract(request.metadata());
        let span = rpc_span("pdpb.PD/GetClusterInfo", &trace_context);
        let service = Arc::clone(&self.inner);
        let response = tokio::task::spawn_blocking(move || {
            let guard = service.blocking_read();
            legacy_get_cluster_info_response_locked(&guard)
        })
        .instrument(span)
        .await
        .map_err(join_error_to_status)?;
        Ok(Response::new(response))
    }

    async fn get_members(
        &self,
        request: Request<pdpb::GetMembersRequest>,
    ) -> Result<Response<pdpb::GetMembersResponse>, Status> {
        let trace_context = TraceContext::extract(request.metadata());
        let span = rpc_span("pdpb.PD/GetMembers", &trace_context);
        let service = Arc::clone(&self.inner);
        let response = tokio::task::spawn_blocking(move || {
            let guard = service.blocking_read();
            legacy_get_members_response_locked(&guard)
        })
        .instrument(span)
        .await
        .map_err(join_error_to_status)?;
        Ok(Response::new(response))
    }

    async fn tso(
        &self,
        request: Request<Streaming<pdpb::TsoRequest>>,
    ) -> Result<Response<Self::TsoStream>, Status> {
        let trace_context = TraceContext::extract(request.metadata());
        let span = rpc_span("pdpb.PD/Tso", &trace_context);
        let forwarded_host = forwarded_host_from_metadata(request.metadata());
        let forwarded_failpoint = forwarded_failpoint_from_metadata(request.metadata());
        let mut incoming = request.into_inner();
        let (tx, rx) = mpsc::channel::<StreamResult<pdpb::TsoResponse>>(32);

        if let Some(forwarded_host) = forwarded_host
            && self.should_forward_to_host(&forwarded_host).await
        {
            let (max_concurrent, recv_timeout, idle_timeout) =
                self.legacy_proxy_runtime_knobs().await;
            let stream_permit = self
                .try_acquire_legacy_proxy_stream(max_concurrent)
                .map_err(|status| *status)?;
            let service = self.clone();
            let trace_context = trace_context.clone();
            tokio::spawn(
                async move {
                    let _stream_permit = stream_permit;
                    loop {
                        let request =
                            match Self::recv_forwarded_legacy_request(&mut incoming, recv_timeout)
                                .await
                            {
                                Ok(Some(req)) => req,
                                Ok(None) => break,
                                Err(status) => {
                                    let _ = tx.send(Err(status)).await;
                                    break;
                                }
                            };
                        match service
                            .dispatch_legacy_forwarded_tso(
                                &forwarded_host,
                                request,
                                trace_context.clone(),
                                forwarded_failpoint,
                                idle_timeout,
                            )
                            .await
                        {
                            Ok(resp) => {
                                if tx.send(Ok(resp)).await.is_err() {
                                    break;
                                }
                            }
                            Err(status) => {
                                let _ = tx.send(Err(status)).await;
                                break;
                            }
                        }
                    }
                }
                .instrument(span),
            );
            return Ok(Response::new(Box::pin(ReceiverStream::new(rx))));
        }

        let executor = {
            let guard = self.inner.read().await;
            guard.tso_executor()
        };

        tokio::spawn(
            async move {
                let mut route_cache = TsoRequestRouteCache::default();
                loop {
                    let next = incoming.message().await;
                    let request = match next {
                        Ok(Some(req)) => req,
                        Ok(None) => break,
                        Err(status) => {
                            let _ = tx.send(Err(status)).await;
                            break;
                        }
                    };

                    let core_request = from_pdpb_tso_request(request);
                    let response = if tokio::runtime::Handle::current().runtime_flavor()
                        == tokio::runtime::RuntimeFlavor::MultiThread
                    {
                        // Legacy PD compatibility requests can block on the
                        // service's std::sync manager lock. Running that
                        // path in `block_in_place` keeps watchdog timers
                        // progressing when the legacy proxy intentionally
                        // exercises lock contention.
                        tokio::task::block_in_place(|| {
                            executor.handle_tso(now_ms(), core_request, &mut route_cache)
                        })
                    } else {
                        executor.handle_tso(now_ms(), core_request, &mut route_cache)
                    };

                    match response {
                        Ok(resp) => {
                            if tx.send(Ok(to_pdpb_tso_response(resp))).await.is_err() {
                                break;
                            }
                        }
                        Err(status) => {
                            let _ = tx.send(Err(status_from_core(status))).await;
                            break;
                        }
                    }
                }
            }
            .instrument(span),
        );

        Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
    }

    async fn get_min_ts(
        &self,
        request: Request<pdpb::GetMinTsRequest>,
    ) -> Result<Response<pdpb::GetMinTsResponse>, Status> {
        let trace_context = TraceContext::extract(request.metadata());
        let span = rpc_span("pdpb.PD/GetMinTS", &trace_context);
        let request = request.into_inner();
        let (authz_sender_id, cluster_id, started) = {
            let guard = self.inner.read().await;
            (
                guard.runtime_config().get_min_ts_authz_allow_sender_id,
                guard.cluster_id(),
                guard.is_started(),
            )
        };
        if !started {
            return Err(status_from_core(ServiceError::NotStarted.to_grpc_status()));
        }
        let request_header = request.header.as_ref();
        let request_cluster_id = request_header.map_or(0, |header| header.cluster_id);
        if request_cluster_id != cluster_id {
            return Err(status_from_core(
                ServiceError::ClusterMismatched {
                    expected: cluster_id,
                    got: request_cluster_id,
                }
                .to_grpc_status(),
            ));
        }
        let request_sender_id = request_header.map_or(0, |header| header.sender_id);
        if authz_sender_id != 0 && request_sender_id != authz_sender_id {
            metrics::observe_get_min_ts_request("authz_denied");
            metrics::observe_get_min_ts_authz_denied();
            let response_cluster_id = request_header.map_or(cluster_id, |header| {
                if header.cluster_id == 0 {
                    cluster_id
                } else {
                    header.cluster_id
                }
            });
            return Ok(Response::new(pdpb_get_min_ts_error(
                response_cluster_id,
                "get_min_ts caller is not authorized".to_string(),
            )));
        }
        let response = self
            .aggregate_get_min_ts(request, trace_context)
            .instrument(span)
            .await;
        if response
            .header
            .as_ref()
            .and_then(|header| header.error.as_ref())
            .is_none()
        {
            metrics::observe_get_min_ts_request("success");
        } else {
            metrics::observe_get_min_ts_request("error");
        }
        Ok(Response::new(response))
    }

    async fn bootstrap(
        &self,
        _request: Request<pdpb::BootstrapRequest>,
    ) -> std::result::Result<Response<pdpb::BootstrapResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.bootstrap is not supported by pd-tso",
        ))
    }

    async fn is_bootstrapped(
        &self,
        _request: Request<pdpb::IsBootstrappedRequest>,
    ) -> std::result::Result<Response<pdpb::IsBootstrappedResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.is_bootstrapped is not supported by pd-tso",
        ))
    }

    async fn alloc_id(
        &self,
        _request: Request<pdpb::AllocIdRequest>,
    ) -> std::result::Result<Response<pdpb::AllocIdResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.alloc_id is not supported by pd-tso",
        ))
    }

    async fn is_snapshot_recovering(
        &self,
        _request: Request<pdpb::IsSnapshotRecoveringRequest>,
    ) -> std::result::Result<Response<pdpb::IsSnapshotRecoveringResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.is_snapshot_recovering is not supported by pd-tso",
        ))
    }

    async fn get_store(
        &self,
        _request: Request<pdpb::GetStoreRequest>,
    ) -> std::result::Result<Response<pdpb::GetStoreResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_store is not supported by pd-tso",
        ))
    }

    async fn put_store(
        &self,
        _request: Request<pdpb::PutStoreRequest>,
    ) -> std::result::Result<Response<pdpb::PutStoreResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.put_store is not supported by pd-tso",
        ))
    }

    async fn get_all_stores(
        &self,
        _request: Request<pdpb::GetAllStoresRequest>,
    ) -> std::result::Result<Response<pdpb::GetAllStoresResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_all_stores is not supported by pd-tso",
        ))
    }

    async fn store_heartbeat(
        &self,
        _request: Request<pdpb::StoreHeartbeatRequest>,
    ) -> std::result::Result<Response<pdpb::StoreHeartbeatResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.store_heartbeat is not supported by pd-tso",
        ))
    }

    async fn region_heartbeat(
        &self,
        _request: Request<tonic::Streaming<pdpb::RegionHeartbeatRequest>>,
    ) -> std::result::Result<Response<Self::RegionHeartbeatStream>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.region_heartbeat is not supported by pd-tso",
        ))
    }

    async fn get_region(
        &self,
        _request: Request<pdpb::GetRegionRequest>,
    ) -> std::result::Result<Response<pdpb::GetRegionResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_region is not supported by pd-tso",
        ))
    }

    async fn get_prev_region(
        &self,
        _request: Request<pdpb::GetRegionRequest>,
    ) -> std::result::Result<Response<pdpb::GetRegionResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_prev_region is not supported by pd-tso",
        ))
    }

    async fn get_region_by_id(
        &self,
        _request: Request<pdpb::GetRegionByIdRequest>,
    ) -> std::result::Result<Response<pdpb::GetRegionResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_region_by_id is not supported by pd-tso",
        ))
    }

    async fn query_region(
        &self,
        _request: Request<tonic::Streaming<pdpb::QueryRegionRequest>>,
    ) -> std::result::Result<Response<Self::QueryRegionStream>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.query_region is not supported by pd-tso",
        ))
    }

    async fn scan_regions(
        &self,
        _request: Request<pdpb::ScanRegionsRequest>,
    ) -> std::result::Result<Response<pdpb::ScanRegionsResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.scan_regions is not supported by pd-tso",
        ))
    }

    async fn batch_scan_regions(
        &self,
        _request: Request<pdpb::BatchScanRegionsRequest>,
    ) -> std::result::Result<Response<pdpb::BatchScanRegionsResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.batch_scan_regions is not supported by pd-tso",
        ))
    }

    async fn ask_split(
        &self,
        _request: Request<pdpb::AskSplitRequest>,
    ) -> std::result::Result<Response<pdpb::AskSplitResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.ask_split is not supported by pd-tso",
        ))
    }

    async fn report_split(
        &self,
        _request: Request<pdpb::ReportSplitRequest>,
    ) -> std::result::Result<Response<pdpb::ReportSplitResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.report_split is not supported by pd-tso",
        ))
    }

    async fn ask_batch_split(
        &self,
        _request: Request<pdpb::AskBatchSplitRequest>,
    ) -> std::result::Result<Response<pdpb::AskBatchSplitResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.ask_batch_split is not supported by pd-tso",
        ))
    }

    async fn report_batch_split(
        &self,
        _request: Request<pdpb::ReportBatchSplitRequest>,
    ) -> std::result::Result<Response<pdpb::ReportBatchSplitResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.report_batch_split is not supported by pd-tso",
        ))
    }

    async fn get_cluster_config(
        &self,
        _request: Request<pdpb::GetClusterConfigRequest>,
    ) -> std::result::Result<Response<pdpb::GetClusterConfigResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_cluster_config is not supported by pd-tso",
        ))
    }

    async fn put_cluster_config(
        &self,
        _request: Request<pdpb::PutClusterConfigRequest>,
    ) -> std::result::Result<Response<pdpb::PutClusterConfigResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.put_cluster_config is not supported by pd-tso",
        ))
    }

    async fn scatter_region(
        &self,
        _request: Request<pdpb::ScatterRegionRequest>,
    ) -> std::result::Result<Response<pdpb::ScatterRegionResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.scatter_region is not supported by pd-tso",
        ))
    }

    async fn get_gc_safe_point(
        &self,
        _request: Request<pdpb::GetGcSafePointRequest>,
    ) -> std::result::Result<Response<pdpb::GetGcSafePointResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_gc_safe_point is not supported by pd-tso",
        ))
    }

    async fn update_gc_safe_point(
        &self,
        _request: Request<pdpb::UpdateGcSafePointRequest>,
    ) -> std::result::Result<Response<pdpb::UpdateGcSafePointResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.update_gc_safe_point is not supported by pd-tso",
        ))
    }

    async fn update_service_gc_safe_point(
        &self,
        _request: Request<pdpb::UpdateServiceGcSafePointRequest>,
    ) -> std::result::Result<Response<pdpb::UpdateServiceGcSafePointResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.update_service_gc_safe_point is not supported by pd-tso",
        ))
    }

    async fn get_gc_safe_point_v2(
        &self,
        _request: Request<pdpb::GetGcSafePointV2Request>,
    ) -> std::result::Result<Response<pdpb::GetGcSafePointV2Response>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_gc_safe_point_v2 is not supported by pd-tso",
        ))
    }

    async fn watch_gc_safe_point_v2(
        &self,
        _request: Request<pdpb::WatchGcSafePointV2Request>,
    ) -> std::result::Result<Response<Self::WatchGCSafePointV2Stream>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.watch_gc_safe_point_v2 is not supported by pd-tso",
        ))
    }

    async fn update_gc_safe_point_v2(
        &self,
        _request: Request<pdpb::UpdateGcSafePointV2Request>,
    ) -> std::result::Result<Response<pdpb::UpdateGcSafePointV2Response>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.update_gc_safe_point_v2 is not supported by pd-tso",
        ))
    }

    async fn update_service_safe_point_v2(
        &self,
        _request: Request<pdpb::UpdateServiceSafePointV2Request>,
    ) -> std::result::Result<Response<pdpb::UpdateServiceSafePointV2Response>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.update_service_safe_point_v2 is not supported by pd-tso",
        ))
    }

    async fn get_all_gc_safe_point_v2(
        &self,
        _request: Request<pdpb::GetAllGcSafePointV2Request>,
    ) -> std::result::Result<Response<pdpb::GetAllGcSafePointV2Response>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_all_gc_safe_point_v2 is not supported by pd-tso",
        ))
    }

    async fn advance_gc_safe_point(
        &self,
        _request: Request<pdpb::AdvanceGcSafePointRequest>,
    ) -> std::result::Result<Response<pdpb::AdvanceGcSafePointResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.advance_gc_safe_point is not supported by pd-tso",
        ))
    }

    async fn advance_txn_safe_point(
        &self,
        _request: Request<pdpb::AdvanceTxnSafePointRequest>,
    ) -> std::result::Result<Response<pdpb::AdvanceTxnSafePointResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.advance_txn_safe_point is not supported by pd-tso",
        ))
    }

    async fn set_gc_barrier(
        &self,
        _request: Request<pdpb::SetGcBarrierRequest>,
    ) -> std::result::Result<Response<pdpb::SetGcBarrierResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.set_gc_barrier is not supported by pd-tso",
        ))
    }

    async fn delete_gc_barrier(
        &self,
        _request: Request<pdpb::DeleteGcBarrierRequest>,
    ) -> std::result::Result<Response<pdpb::DeleteGcBarrierResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.delete_gc_barrier is not supported by pd-tso",
        ))
    }

    async fn set_global_gc_barrier(
        &self,
        _request: Request<pdpb::SetGlobalGcBarrierRequest>,
    ) -> std::result::Result<Response<pdpb::SetGlobalGcBarrierResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.set_global_gc_barrier is not supported by pd-tso",
        ))
    }

    async fn delete_global_gc_barrier(
        &self,
        _request: Request<pdpb::DeleteGlobalGcBarrierRequest>,
    ) -> std::result::Result<Response<pdpb::DeleteGlobalGcBarrierResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.delete_global_gc_barrier is not supported by pd-tso",
        ))
    }

    async fn get_gc_state(
        &self,
        _request: Request<pdpb::GetGcStateRequest>,
    ) -> std::result::Result<Response<pdpb::GetGcStateResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_gc_state is not supported by pd-tso",
        ))
    }

    async fn get_all_keyspaces_gc_states(
        &self,
        _request: Request<pdpb::GetAllKeyspacesGcStatesRequest>,
    ) -> std::result::Result<Response<pdpb::GetAllKeyspacesGcStatesResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_all_keyspaces_gc_states is not supported by pd-tso",
        ))
    }

    async fn sync_regions(
        &self,
        _request: Request<tonic::Streaming<pdpb::SyncRegionRequest>>,
    ) -> std::result::Result<Response<Self::SyncRegionsStream>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.sync_regions is not supported by pd-tso",
        ))
    }

    async fn get_operator(
        &self,
        _request: Request<pdpb::GetOperatorRequest>,
    ) -> std::result::Result<Response<pdpb::GetOperatorResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_operator is not supported by pd-tso",
        ))
    }

    async fn sync_max_ts(
        &self,
        _request: Request<pdpb::SyncMaxTsRequest>,
    ) -> std::result::Result<Response<pdpb::SyncMaxTsResponse>, Status> {
        let service = Arc::clone(&self.inner);
        let cluster_id = tokio::task::spawn_blocking(move || {
            let guard = service.blocking_read();
            guard.cluster_id()
        })
        .await
        .map_err(join_error_to_status)?;
        Ok(Response::new(pdpb::SyncMaxTsResponse {
            header: Some(pdpb::ResponseHeader {
                cluster_id,
                error: None,
            }),
            max_local_ts: Some(pdpb::Timestamp::default()),
            synced_dcs: Vec::new(),
        }))
    }

    async fn split_regions(
        &self,
        _request: Request<pdpb::SplitRegionsRequest>,
    ) -> std::result::Result<Response<pdpb::SplitRegionsResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.split_regions is not supported by pd-tso",
        ))
    }

    async fn split_and_scatter_regions(
        &self,
        _request: Request<pdpb::SplitAndScatterRegionsRequest>,
    ) -> std::result::Result<Response<pdpb::SplitAndScatterRegionsResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.split_and_scatter_regions is not supported by pd-tso",
        ))
    }

    async fn get_dc_location_info(
        &self,
        _request: Request<pdpb::GetDcLocationInfoRequest>,
    ) -> std::result::Result<Response<pdpb::GetDcLocationInfoResponse>, Status> {
        let service = Arc::clone(&self.inner);
        let cluster_id = tokio::task::spawn_blocking(move || {
            let guard = service.blocking_read();
            guard.cluster_id()
        })
        .await
        .map_err(join_error_to_status)?;
        Ok(Response::new(pdpb::GetDcLocationInfoResponse {
            header: Some(pdpb::ResponseHeader {
                cluster_id,
                error: None,
            }),
            suffix: 0,
            max_ts: Some(pdpb::Timestamp::default()),
        }))
    }

    async fn store_global_config(
        &self,
        _request: Request<pdpb::StoreGlobalConfigRequest>,
    ) -> std::result::Result<Response<pdpb::StoreGlobalConfigResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.store_global_config is not supported by pd-tso",
        ))
    }

    async fn load_global_config(
        &self,
        _request: Request<pdpb::LoadGlobalConfigRequest>,
    ) -> std::result::Result<Response<pdpb::LoadGlobalConfigResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.load_global_config is not supported by pd-tso",
        ))
    }

    async fn watch_global_config(
        &self,
        _request: Request<pdpb::WatchGlobalConfigRequest>,
    ) -> std::result::Result<Response<Self::WatchGlobalConfigStream>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.watch_global_config is not supported by pd-tso",
        ))
    }

    async fn report_buckets(
        &self,
        _request: Request<tonic::Streaming<pdpb::ReportBucketsRequest>>,
    ) -> std::result::Result<Response<pdpb::ReportBucketsResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.report_buckets is not supported by pd-tso",
        ))
    }

    async fn report_min_resolved_ts(
        &self,
        _request: Request<pdpb::ReportMinResolvedTsRequest>,
    ) -> std::result::Result<Response<pdpb::ReportMinResolvedTsResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.report_min_resolved_ts is not supported by pd-tso",
        ))
    }

    async fn set_external_timestamp(
        &self,
        _request: Request<pdpb::SetExternalTimestampRequest>,
    ) -> std::result::Result<Response<pdpb::SetExternalTimestampResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.set_external_timestamp is not supported by pd-tso",
        ))
    }

    async fn get_external_timestamp(
        &self,
        _request: Request<pdpb::GetExternalTimestampRequest>,
    ) -> std::result::Result<Response<pdpb::GetExternalTimestampResponse>, Status> {
        Err(Status::unimplemented(
            "pdpb.PD.get_external_timestamp is not supported by pd-tso",
        ))
    }
}

fn aggregate_fanout_get_min_ts(
    cluster_id: u64,
    responses: Vec<tsopb::GetMinTsResponse>,
) -> pdpb::GetMinTsResponse {
    if responses.is_empty() {
        return pdpb_get_min_ts_error(
            cluster_id,
            "GetMinTS fan-out returned no successful pod responses".to_string(),
        );
    }

    let mut accepted = Vec::new();
    for response in responses {
        let has_error = response
            .header
            .as_ref()
            .and_then(|header| header.error.as_ref())
            .is_some();
        if has_error {
            continue;
        }
        accepted.push(response);
    }
    if accepted.is_empty() {
        return pdpb_get_min_ts_error(
            cluster_id,
            "GetMinTS fan-out had no successful pod responses".to_string(),
        );
    }

    let expected_total = accepted[0].keyspace_groups_total;
    if accepted
        .iter()
        .any(|response| response.keyspace_groups_total != expected_total)
    {
        return pdpb_get_min_ts_error(
            cluster_id,
            "GetMinTS fan-out total mismatch across pod responses".to_string(),
        );
    }
    let serving_sum = accepted.iter().fold(0_u64, |acc, response| {
        acc.saturating_add(u64::from(response.keyspace_groups_serving))
    });
    if serving_sum != u64::from(expected_total) {
        return pdpb_get_min_ts_error(
            cluster_id,
            format!(
                "GetMinTS fan-out serving sum {} does not match total {}",
                serving_sum, expected_total
            ),
        );
    }

    let mut min_timestamp: Option<pdpb::Timestamp> = None;
    let mut response_cluster_id = cluster_id;
    for response in accepted {
        if let Some(header) = response.header.as_ref()
            && header.cluster_id != 0
        {
            response_cluster_id = header.cluster_id;
        }
        let Some(timestamp) = response.timestamp else {
            continue;
        };
        if !is_non_empty_timestamp(&timestamp) {
            continue;
        }
        if min_timestamp
            .as_ref()
            .is_none_or(|current| compare_pdpb_timestamp(&timestamp, current).is_lt())
        {
            min_timestamp = Some(timestamp);
        }
    }

    let Some(timestamp) = min_timestamp else {
        return pdpb_get_min_ts_error(
            cluster_id,
            "GetMinTS fan-out had no non-empty minimum timestamp".to_string(),
        );
    };
    pdpb::GetMinTsResponse {
        header: Some(pdpb::ResponseHeader {
            cluster_id: response_cluster_id,
            error: None,
        }),
        timestamp: Some(timestamp),
    }
}

fn pdpb_get_min_ts_error(cluster_id: u64, message: String) -> pdpb::GetMinTsResponse {
    pdpb::GetMinTsResponse {
        header: Some(pdpb::ResponseHeader {
            cluster_id,
            error: Some(pdpb::Error {
                r#type: pdpb::ErrorType::Unknown as i32,
                message,
            }),
        }),
        timestamp: Some(pdpb::Timestamp::default()),
    }
}

fn is_non_empty_timestamp(timestamp: &pdpb::Timestamp) -> bool {
    timestamp.physical != 0 || timestamp.logical != 0 || timestamp.suffix_bits != 0
}

fn compare_pdpb_timestamp(left: &pdpb::Timestamp, right: &pdpb::Timestamp) -> std::cmp::Ordering {
    (left.physical, left.logical, left.suffix_bits).cmp(&(
        right.physical,
        right.logical,
        right.suffix_bits,
    ))
}

#[derive(Default)]
struct LegacyDiscoverySnapshot {
    tso_urls: Vec<String>,
    members: Vec<pdpb::Member>,
    leader: Option<pdpb::Member>,
}

fn legacy_not_started_error() -> Option<pdpb::Error> {
    to_pdpb_error(Some(ServiceError::NotStarted.to_unary_error()))
}

fn legacy_get_cluster_info_response_locked(service: &TsoService) -> pdpb::GetClusterInfoResponse {
    let cluster_id = service.cluster_id();
    if !service.is_started() {
        return pdpb::GetClusterInfoResponse {
            header: Some(pdpb::ResponseHeader {
                cluster_id,
                error: legacy_not_started_error(),
            }),
            service_modes: Vec::new(),
            tso_urls: Vec::new(),
        };
    }

    let discovery = legacy_discovery_snapshot_locked(service);
    pdpb::GetClusterInfoResponse {
        header: Some(pdpb::ResponseHeader {
            cluster_id,
            error: None,
        }),
        service_modes: vec![pdpb::ServiceMode::ApiSvcMode as i32],
        tso_urls: discovery.tso_urls,
    }
}

fn legacy_get_members_response_locked(service: &TsoService) -> pdpb::GetMembersResponse {
    let cluster_id = service.cluster_id();
    if !service.is_started() {
        return pdpb::GetMembersResponse {
            header: Some(pdpb::ResponseHeader {
                cluster_id,
                error: legacy_not_started_error(),
            }),
            members: Vec::new(),
            leader: None,
            etcd_leader: None,
            tso_allocator_leaders: HashMap::new(),
        };
    }

    legacy_get_members_response(cluster_id, legacy_discovery_snapshot_locked(service))
}

fn legacy_get_members_response(
    cluster_id: u64,
    discovery: LegacyDiscoverySnapshot,
) -> pdpb::GetMembersResponse {
    pdpb::GetMembersResponse {
        header: Some(pdpb::ResponseHeader {
            cluster_id,
            error: None,
        }),
        members: discovery.members,
        leader: discovery.leader.clone(),
        etcd_leader: discovery.leader,
        tso_allocator_leaders: HashMap::new(),
    }
}

async fn legacy_discovery_snapshot(service: &Arc<RwLock<TsoService>>) -> LegacyDiscoverySnapshot {
    let guard = service.read().await;
    legacy_discovery_snapshot_locked(&guard)
}

fn legacy_discovery_snapshot_locked(service: &TsoService) -> LegacyDiscoverySnapshot {
    let manager = service.manager();
    let mut urls = manager.discovery_tso_urls();
    let mut members = urls
        .iter()
        .map(|url| member_from_client_url(url))
        .collect::<Vec<_>>();
    let mut leader = None;
    let mut ambiguous_leader = false;
    for group_id in manager.active_group_ids() {
        let Some(group) = manager.group(group_id) else {
            continue;
        };
        for member in &group.members {
            let raw_address = member.address.trim();
            if raw_address.is_empty() {
                continue;
            }
            let mapped_member = member_from_client_url(raw_address);
            if member.is_primary {
                match leader.as_ref() {
                    None => leader = Some(mapped_member.clone()),
                    Some(current) if current.member_id != mapped_member.member_id => {
                        ambiguous_leader = true;
                    }
                    _ => {}
                }
            }
        }
    }
    drop(manager);

    if urls.is_empty() {
        urls = legacy_discovery_tso_urls();
        members = urls.iter().map(|url| member_from_client_url(url)).collect();
        leader = None;
    } else if ambiguous_leader {
        leader = None;
    }

    LegacyDiscoverySnapshot {
        tso_urls: urls,
        members,
        leader,
    }
}

fn member_from_client_url(url: &str) -> pdpb::Member {
    let normalized = normalize_endpoint(url);
    let name = if normalized.is_empty() {
        url.to_string()
    } else {
        normalized.clone()
    };
    pdpb::Member {
        name,
        member_id: stable_member_id(&normalized),
        peer_urls: Vec::new(),
        client_urls: vec![url.to_string()],
        leader_priority: 0,
        deploy_path: String::new(),
        binary_version: String::new(),
        git_hash: String::new(),
        dc_location: String::new(),
    }
}

fn now_ms() -> i64 {
    let duration = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default();
    (duration.as_millis().min(i64::MAX as u128)) as i64
}

fn legacy_discovery_tso_urls() -> Vec<String> {
    if let Ok(raw) = std::env::var("PD_TSO_LEGACY_DISCOVERY_URLS") {
        let urls = raw
            .split(',')
            .map(str::trim)
            .filter(|value| !value.is_empty())
            .map(ToString::to_string)
            .collect::<Vec<_>>();
        if !urls.is_empty() {
            return urls;
        }
    }
    if let Ok(addr) = std::env::var("PD_TSO_ADVERTISE_LISTEN_ADDR")
        && !addr.trim().is_empty()
    {
        return vec![addr];
    }
    if let Ok(addr) = std::env::var("PD_TSO_LISTEN_ADDR")
        && !addr.trim().is_empty()
    {
        return vec![addr];
    }
    Vec::new()
}

fn legacy_local_tso_urls() -> Vec<String> {
    let mut urls = Vec::new();
    for key in ["PD_TSO_ADVERTISE_LISTEN_ADDR", "PD_TSO_LISTEN_ADDR"] {
        if let Ok(value) = std::env::var(key) {
            for part in value.split(',') {
                let part = part.trim();
                if !part.is_empty() {
                    urls.push(part.to_string());
                }
            }
        }
    }
    urls
}

fn forwarded_host_from_metadata(metadata: &tonic::metadata::MetadataMap) -> Option<String> {
    metadata
        .get("pd-forwarded-host")
        .and_then(|value| value.to_str().ok())
        .filter(|value| !value.is_empty())
        .map(ToString::to_string)
}

fn forwarded_failpoint_from_metadata(
    metadata: &tonic::metadata::MetadataMap,
) -> Option<LegacyForwardFailpoint> {
    metadata
        .get(LEGACY_FORWARD_FAILPOINT_METADATA_KEY)
        .and_then(|value| value.to_str().ok())
        .and_then(LegacyForwardFailpoint::parse)
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use pd_proto_grpc::{pdpb, tsopb};
    use pd_test_util::env::{EnvVarRestore, async_env_lock, remove_env, set_env};
    use tokio::sync::{RwLock, oneshot};
    use tonic::metadata::{MetadataMap, MetadataValue};

    use crate::allocator::AllocatorConfig;
    use crate::forwarding::{
        BootstrapMode, LEGACY_FORWARD_MERGED_COUNT_OVERFLOW_TOKEN, LegacyForwardDispatchRequest,
        LegacyForwardFailpoint, bootstrap_mode_from_env, map_tsopb_error_to_pdpb,
        merged_forward_count, parse_bootstrap_keyspace_id, split_forwarded_responses,
    };
    use crate::manager::KeyspaceGroupManager;
    use crate::persistence::InMemoryEtcdStore;
    use crate::service::TsoService;
    use crate::trace_context::TraceContext;

    use super::{
        LEGACY_FORWARD_FAILPOINT_METADATA_KEY, LEGACY_FORWARD_FAILPOINT_RECV_TIMEOUT_VALUE,
        LEGACY_FORWARD_FAILPOINT_SEND_TIMEOUT_VALUE, LegacyDiscoverySnapshot, LegacyPdGrpcService,
        aggregate_fanout_get_min_ts, forwarded_failpoint_from_metadata, from_pdpb_tso_request,
        legacy_get_cluster_info_response_locked, legacy_get_members_response,
        legacy_get_members_response_locked, legacy_local_tso_urls,
    };

    #[test]
    fn forwarding_error_mapping_preserves_ok_and_maps_non_ok_to_unknown() {
        let ok = map_tsopb_error_to_pdpb(Some(tsopb::Error {
            r#type: tsopb::ErrorType::Ok as i32,
            message: "ok".to_string(),
        }))
        .expect("mapped");
        assert_eq!(ok.r#type, pdpb::ErrorType::Ok as i32);

        let non_ok = map_tsopb_error_to_pdpb(Some(tsopb::Error {
            r#type: tsopb::ErrorType::ClusterMismatched as i32,
            message: "cluster".to_string(),
        }))
        .expect("mapped");
        assert_eq!(non_ok.r#type, pdpb::ErrorType::Unknown as i32);
    }

    #[test]
    fn legacy_pdpb_request_conversion_forces_default_empty_dc_location() {
        let converted = from_pdpb_tso_request(pdpb::TsoRequest {
            header: Some(pdpb::RequestHeader {
                cluster_id: 1,
                sender_id: 7,
                ..pdpb::RequestHeader::default()
            }),
            count: 3,
            dc_location: "zone-a".to_string(),
        });
        assert!(converted.dc_location.is_empty());
    }

    #[test]
    fn bootstrap_keyspace_id_parser_honors_mode_default_and_override() {
        assert_eq!(parse_bootstrap_keyspace_id(None, BootstrapMode::Classic), 0);
        assert_eq!(
            parse_bootstrap_keyspace_id(None, BootstrapMode::NextGen),
            16_777_214
        );
        assert_eq!(
            parse_bootstrap_keyspace_id(Some("abc".to_string()), BootstrapMode::NextGen),
            16_777_214
        );
        assert_eq!(
            parse_bootstrap_keyspace_id(Some("42".to_string()), BootstrapMode::Classic),
            42
        );
    }

    #[test]
    fn bootstrap_mode_parser_supports_classic_and_nextgen_inputs() {
        assert_eq!(
            bootstrap_mode_from_env(Some("nextgen".to_string()), None),
            BootstrapMode::NextGen
        );
        assert_eq!(
            bootstrap_mode_from_env(None, Some("next-gen".to_string())),
            BootstrapMode::NextGen
        );
        assert_eq!(
            bootstrap_mode_from_env(Some("classic".to_string()), None),
            BootstrapMode::Classic
        );
        assert_eq!(bootstrap_mode_from_env(None, None), BootstrapMode::Classic);
    }

    #[test]
    fn forwarded_failpoint_metadata_parser_accepts_known_tokens() {
        let mut send_metadata = MetadataMap::new();
        send_metadata.insert(
            LEGACY_FORWARD_FAILPOINT_METADATA_KEY,
            MetadataValue::from_static(LEGACY_FORWARD_FAILPOINT_SEND_TIMEOUT_VALUE),
        );
        assert_eq!(
            forwarded_failpoint_from_metadata(&send_metadata),
            Some(LegacyForwardFailpoint::SendToTsoTimeout)
        );

        let mut recv_metadata = MetadataMap::new();
        recv_metadata.insert(
            LEGACY_FORWARD_FAILPOINT_METADATA_KEY,
            MetadataValue::from_static(LEGACY_FORWARD_FAILPOINT_RECV_TIMEOUT_VALUE),
        );
        assert_eq!(
            forwarded_failpoint_from_metadata(&recv_metadata),
            Some(LegacyForwardFailpoint::RecvFromTsoTimeout)
        );
    }

    #[test]
    fn legacy_local_tso_urls_splits_and_trims_comma_separated_env_values() {
        let _guard = async_env_lock().blocking_lock();
        let [_advertise_restore, _listen_restore] =
            EnvVarRestore::capture_many(["PD_TSO_ADVERTISE_LISTEN_ADDR", "PD_TSO_LISTEN_ADDR"]);

        set_env(
            "PD_TSO_ADVERTISE_LISTEN_ADDR",
            "127.0.0.1:3379, 127.0.0.1:3380,",
        );
        set_env("PD_TSO_LISTEN_ADDR", " 127.0.0.1:3390 ");

        let urls = legacy_local_tso_urls();
        assert_eq!(
            urls,
            vec![
                "127.0.0.1:3379".to_string(),
                "127.0.0.1:3380".to_string(),
                "127.0.0.1:3390".to_string()
            ]
        );
    }

    #[tokio::test]
    async fn should_forward_to_host_treats_split_local_env_values_as_local() {
        let _guard = async_env_lock().lock().await;
        let [_advertise_restore, _listen_restore] =
            EnvVarRestore::capture_many(["PD_TSO_ADVERTISE_LISTEN_ADDR", "PD_TSO_LISTEN_ADDR"]);

        set_env(
            "PD_TSO_ADVERTISE_LISTEN_ADDR",
            "127.0.0.1:3379,127.0.0.1:3380",
        );
        remove_env("PD_TSO_LISTEN_ADDR");

        let store = Arc::new(InMemoryEtcdStore::new());
        let manager = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));
        let grpc = LegacyPdGrpcService::new(service);

        assert!(!grpc.should_forward_to_host("127.0.0.1:3380").await);
        assert!(grpc.should_forward_to_host("127.0.0.1:3399").await);
    }

    #[test]
    fn fanout_get_min_ts_empty_response_set_returns_handled_error() {
        let response = aggregate_fanout_get_min_ts(42, Vec::new());
        let header = response.header.expect("header");
        assert_eq!(header.cluster_id, 42);
        assert!(header.error.is_some());
        assert!(response.timestamp.is_some());
    }

    #[test]
    fn get_members_response_keeps_unresolved_leader_unset() {
        let response = legacy_get_members_response(9, LegacyDiscoverySnapshot::default());
        let header = response.header.expect("header");
        assert_eq!(header.cluster_id, 9);
        assert!(header.error.is_none());
        assert!(response.members.is_empty());
        assert!(response.leader.is_none());
        assert!(response.etcd_leader.is_none());
    }

    #[test]
    fn get_members_response_does_not_infer_leader_from_member_order() {
        let response = legacy_get_members_response(
            9,
            LegacyDiscoverySnapshot {
                tso_urls: vec!["a".to_string(), "b".to_string()],
                members: vec![
                    pdpb::Member {
                        name: "first".to_string(),
                        client_urls: vec!["a".to_string()],
                        ..pdpb::Member::default()
                    },
                    pdpb::Member {
                        name: "second".to_string(),
                        client_urls: vec!["b".to_string()],
                        ..pdpb::Member::default()
                    },
                ],
                leader: None,
            },
        );
        assert_eq!(response.members.len(), 2);
        assert!(response.leader.is_none());
        assert!(response.etcd_leader.is_none());
    }

    #[test]
    fn get_members_response_uses_explicit_resolved_leader() {
        let leader = pdpb::Member {
            name: "leader".to_string(),
            client_urls: vec!["a".to_string()],
            ..pdpb::Member::default()
        };
        let response = legacy_get_members_response(
            9,
            LegacyDiscoverySnapshot {
                tso_urls: vec!["a".to_string(), "b".to_string()],
                members: vec![
                    leader.clone(),
                    pdpb::Member {
                        name: "second".to_string(),
                        client_urls: vec!["b".to_string()],
                        ..pdpb::Member::default()
                    },
                ],
                leader: Some(leader),
            },
        );
        assert_eq!(response.members.len(), 2);
        assert_eq!(
            response
                .leader
                .as_ref()
                .expect("leader")
                .client_urls
                .first()
                .expect("leader url"),
            "a"
        );
        assert_eq!(
            response.etcd_leader.as_ref().expect("etcd leader").name,
            "leader"
        );
    }

    #[test]
    fn discovery_snapshot_resolves_leader_from_primary_marker() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut manager = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 0,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: vec![
                    pd_proto_tso::KeyspaceGroupMember {
                        address: "node-a:1234".to_string(),
                        is_primary: false,
                    },
                    pd_proto_tso::KeyspaceGroupMember {
                        address: "node-b:1234".to_string(),
                        is_primary: true,
                    },
                ],
                keyspaces: vec![0],
            })
            .expect("register");
        let service = TsoService::new(1, manager);

        let snapshot = super::legacy_discovery_snapshot_locked(&service);
        assert_eq!(snapshot.members.len(), 2);
        assert_eq!(
            snapshot
                .leader
                .as_ref()
                .expect("resolved leader")
                .client_urls
                .first()
                .expect("leader url"),
            "node-b:1234"
        );
    }

    #[test]
    fn discovery_snapshot_keeps_leader_unset_without_primary_marker() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut manager = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 0,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: vec![pd_proto_tso::KeyspaceGroupMember {
                    address: "node-a:1234".to_string(),
                    is_primary: false,
                }],
                keyspaces: vec![0],
            })
            .expect("register");
        let service = TsoService::new(1, manager);

        let snapshot = super::legacy_discovery_snapshot_locked(&service);
        assert_eq!(snapshot.members.len(), 1);
        assert!(snapshot.leader.is_none());
    }

    #[test]
    fn discovery_cluster_info_returns_handled_not_bootstrapped_when_stopped() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let manager = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        let service = TsoService::new(1, manager);
        service.stop();

        let response = legacy_get_cluster_info_response_locked(&service);
        let header = response.header.expect("header");
        let error = header.error.expect("error");
        assert_eq!(error.r#type, pdpb::ErrorType::Unknown as i32);
        assert!(response.tso_urls.is_empty());
    }

    #[test]
    fn discovery_get_members_returns_handled_not_bootstrapped_when_stopped() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let manager = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        let service = TsoService::new(1, manager);
        service.stop();

        let response = legacy_get_members_response_locked(&service);
        let header = response.header.expect("header");
        let error = header.error.expect("error");
        assert_eq!(error.r#type, pdpb::ErrorType::Unknown as i32);
        assert!(response.members.is_empty());
        assert!(response.leader.is_none());
        assert!(response.etcd_leader.is_none());
    }

    #[test]
    fn merged_forward_count_rejects_overflow_with_parity_token() {
        let (tx_a, _rx_a) = oneshot::channel();
        let (tx_b, _rx_b) = oneshot::channel();
        let requests = vec![
            LegacyForwardDispatchRequest {
                request: pdpb::TsoRequest {
                    header: Some(pdpb::RequestHeader {
                        cluster_id: 1,
                        ..pdpb::RequestHeader::default()
                    }),
                    count: u32::MAX,
                    dc_location: String::new(),
                },
                trace_context: TraceContext::default(),
                failpoint: None,
                response_tx: tx_a,
            },
            LegacyForwardDispatchRequest {
                request: pdpb::TsoRequest {
                    header: Some(pdpb::RequestHeader {
                        cluster_id: 2,
                        ..pdpb::RequestHeader::default()
                    }),
                    count: 1,
                    dc_location: String::new(),
                },
                trace_context: TraceContext::default(),
                failpoint: None,
                response_tx: tx_b,
            },
        ];

        let err = merged_forward_count(&requests).expect_err("overflow should be rejected");
        assert_eq!(err.code(), tonic::Code::Unknown);
        assert!(
            err.message()
                .contains(LEGACY_FORWARD_MERGED_COUNT_OVERFLOW_TOKEN)
        );
    }

    #[test]
    fn split_forwarded_responses_preserve_order_header_and_logical_arithmetic() {
        let (tx_a, _rx_a) = oneshot::channel();
        let (tx_b, _rx_b) = oneshot::channel();
        let requests = vec![
            LegacyForwardDispatchRequest {
                request: pdpb::TsoRequest {
                    header: Some(pdpb::RequestHeader {
                        cluster_id: 11,
                        ..pdpb::RequestHeader::default()
                    }),
                    count: 2,
                    dc_location: String::new(),
                },
                trace_context: TraceContext::default(),
                failpoint: None,
                response_tx: tx_a,
            },
            LegacyForwardDispatchRequest {
                request: pdpb::TsoRequest {
                    header: Some(pdpb::RequestHeader {
                        cluster_id: 22,
                        ..pdpb::RequestHeader::default()
                    }),
                    count: 3,
                    dc_location: String::new(),
                },
                trace_context: TraceContext::default(),
                failpoint: None,
                response_tx: tx_b,
            },
        ];
        let merged = pdpb::TsoResponse {
            header: Some(pdpb::ResponseHeader {
                cluster_id: 99,
                error: None,
            }),
            count: 5,
            timestamp: Some(pdpb::Timestamp {
                physical: 7,
                logical: 30,
                suffix_bits: 4,
            }),
        };

        let split = split_forwarded_responses(&requests, merged, 5);
        assert_eq!(split.len(), 2);
        assert_eq!(split[0].header.as_ref().expect("header").cluster_id, 11);
        assert_eq!(split[1].header.as_ref().expect("header").cluster_id, 22);
        assert!(split[0].header.as_ref().expect("header").error.is_none());
        assert!(split[1].header.as_ref().expect("header").error.is_none());

        let ts0 = split[0].timestamp.as_ref().expect("timestamp0");
        let ts1 = split[1].timestamp.as_ref().expect("timestamp1");
        assert_eq!(ts0.physical, 7);
        assert_eq!(ts1.physical, 7);
        assert_eq!(ts0.suffix_bits, 4);
        assert_eq!(ts1.suffix_bits, 4);
        assert_eq!(ts0.logical, 27);
        assert_eq!(ts1.logical, 30);
    }
}