use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex as StdMutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use etcd_client::{
    Client, Compare, CompareOp, GetOptions, LeaseKeepAliveStream, LeaseKeeper, PutOptions, Txn,
    TxnOp, WatchOptions,
};
use once_cell::sync::Lazy;
use pd_proto_grpc::tsopb;
use prost::Message;
use serde::Deserialize;
use tokio::sync::{Mutex, RwLock};
use tokio::task::JoinHandle;
use tracing::{debug, info, warn};

use crate::manager::{
    GroupElectionMeta, KeyspaceGroupManager, expected_primary_matches_local_identity,
};
use crate::metrics;
use crate::service::TsoService;

const WATCH_RETRY_BACKOFF_MAX_EXP_SHIFT: u32 = 6;
const ASSIGNED_GROUP_QPS_STALE_THRESHOLD_MS: i64 = 5 * 60 * 1_000;
const NTP_PACKET_SIZE: usize = 48;
const NTP_UNIX_EPOCH_OFFSET_SECS: i64 = 2_208_988_800;

#[derive(Debug, Clone, Copy)]
struct AssignedGroupQpsSample {
    observed_at: Instant,
    total_requests: u64,
}

static ASSIGNED_GROUP_QPS_SAMPLE: Lazy<StdMutex<Option<AssignedGroupQpsSample>>> =
    Lazy::new(|| StdMutex::new(None));

#[derive(Debug)]
struct LeaseKeepAliveSession {
    lease_id: i64,
    keeper: LeaseKeeper,
    stream: LeaseKeepAliveStream,
}

#[derive(Debug)]
struct ExternalClockSkewSampler {
    next_poll_at: Instant,
    last_observed_skew_seconds: Option<f64>,
}

impl ExternalClockSkewSampler {
    fn new(now: Instant) -> Self {
        Self {
            next_poll_at: now,
            last_observed_skew_seconds: None,
        }
    }
}

/// Runtime configuration for background HA/watch loops.
#[derive(Debug, Clone)]
pub struct TsoRuntimeConfig {
    /// Cluster id used for etcd key layout and header validation.
    pub cluster_id: u64,
    /// Etcd client endpoints for election, watch, and timestamp persistence.
    pub etcd_endpoints: Vec<String>,
    /// Group-election campaign loop cadence.
    pub campaign_interval: Duration,
    /// Lease keepalive and maintenance loop cadence.
    pub keepalive_interval: Duration,
    /// Backoff base used for watch/campaign reconnect retry.
    pub watch_retry_backoff: Duration,
    /// Etcd lease TTL used by primary campaign/keepalive.
    pub lease_ttl_secs: i64,
    /// Optional external NTP server (`host:port`) for dual-source clock-skew monitoring.
    pub external_ntp_server: Option<String>,
    /// Minimum cadence between external NTP samples.
    pub external_ntp_refresh_interval: Duration,
    /// Timeout for one external NTP sampling operation.
    pub external_ntp_timeout: Duration,
}

impl Default for TsoRuntimeConfig {
    fn default() -> Self {
        Self {
            cluster_id: 0,
            etcd_endpoints: vec!["http://127.0.0.1:2379".to_string()],
            campaign_interval: Duration::from_millis(500),
            keepalive_interval: Duration::from_millis(200),
            watch_retry_backoff: Duration::from_millis(500),
            lease_ttl_secs: 5,
            external_ntp_server: None,
            external_ntp_refresh_interval: Duration::from_secs(30),
            external_ntp_timeout: Duration::from_secs(1),
        }
    }
}

/// Background runtime that drives membership watch and election/lease loops.
pub struct TsoRuntime {
    handles: Vec<JoinHandle<()>>,
}

impl TsoRuntime {
    /// Starts background loops.
    pub fn start(service: Arc<RwLock<TsoService>>, cfg: TsoRuntimeConfig) -> Self {
        let leases = Arc::new(Mutex::new(HashMap::<u32, i64>::new()));
        let external_skew_sampler =
            Arc::new(Mutex::new(ExternalClockSkewSampler::new(Instant::now())));
        let watch_handle =
            tokio::spawn(run_membership_watch_loop(Arc::clone(&service), cfg.clone()));
        let campaign_handle = tokio::spawn(run_campaign_loop(
            Arc::clone(&service),
            Arc::clone(&leases),
            Arc::clone(&external_skew_sampler),
            cfg.clone(),
        ));
        let keepalive_handle = tokio::spawn(run_lease_refresh_loop(
            service,
            leases,
            external_skew_sampler,
            cfg,
        ));
        Self {
            handles: vec![watch_handle, campaign_handle, keepalive_handle],
        }
    }
}

impl Drop for TsoRuntime {
    fn drop(&mut self) {
        for handle in &self.handles {
            handle.abort();
        }
    }
}

async fn run_membership_watch_loop(service: Arc<RwLock<TsoService>>, cfg: TsoRuntimeConfig) {
    let prefix = membership_prefix(cfg.cluster_id);
    let mut retry_attempt = 0_u32;

    loop {
        let started = Instant::now();
        let connect = Client::connect(cfg.etcd_endpoints.clone(), None).await;
        observe_etcd_result("connect", started, &connect);
        let mut client = match connect {
            Ok(client) => client,
            Err(_) => {
                force_groups_non_serving_for_watch_resync(&service, None, "watch connect failed")
                    .await;
                retry_attempt = retry_attempt.saturating_add(1);
                tokio::time::sleep(watch_retry_sleep_duration(
                    cfg.watch_retry_backoff,
                    retry_attempt,
                    watch_retry_entropy(),
                ))
                .await;
                continue;
            }
        };

        let watch_start_revision =
            match sync_membership_snapshot(&mut client, Arc::clone(&service), &prefix).await {
                Ok(revision) => {
                    let guard = service.write().await;
                    guard.set_membership_resync_blocking(false);
                    revision
                }
                Err(_) => {
                    force_groups_non_serving_for_watch_resync(
                        &service,
                        None,
                        "membership snapshot sync failed",
                    )
                    .await;
                    retry_attempt = retry_attempt.saturating_add(1);
                    tokio::time::sleep(watch_retry_sleep_duration(
                        cfg.watch_retry_backoff,
                        retry_attempt,
                        watch_retry_entropy(),
                    ))
                    .await;
                    continue;
                }
            };

        let mut watch_options = WatchOptions::new().with_prefix();
        if watch_start_revision > 0 {
            watch_options = watch_options.with_start_revision(watch_start_revision);
        }

        let watch_started = Instant::now();
        let watch = client.watch(prefix.clone(), Some(watch_options)).await;
        observe_etcd_result("watch", watch_started, &watch);
        let (_watcher, mut stream) = match watch {
            Ok(pair) => pair,
            Err(_) => {
                force_groups_non_serving_for_watch_resync(
                    &service,
                    None,
                    "membership watch start failed",
                )
                .await;
                retry_attempt = retry_attempt.saturating_add(1);
                tokio::time::sleep(watch_retry_sleep_duration(
                    cfg.watch_retry_backoff,
                    retry_attempt,
                    watch_retry_entropy(),
                ))
                .await;
                continue;
            }
        };

        let mut healthy_cycle = false;
        loop {
            match stream.message().await {
                Ok(Some(resp)) => {
                    if resp.canceled() {
                        force_groups_non_serving_for_watch_resync(
                            &service,
                            None,
                            "membership watch canceled",
                        )
                        .await;
                        break;
                    }

                    let mut resync_required = false;
                    let mut affected_groups = HashSet::new();
                    {
                        let guard = service.write().await;
                        let mut manager = guard.manager_mut();
                        if let Some(header) = resp.header() {
                            apply_mod_revision(&mut manager, header.revision());
                            metrics::set_etcd_watch_current_revision(&prefix, header.revision());
                        }
                        for event in resp.events() {
                            match event.event_type() {
                                etcd_client::EventType::Put => {
                                    if let Some(kv) = event.kv() {
                                        apply_mod_revision(&mut manager, kv.mod_revision());
                                        metrics::set_etcd_watch_current_revision(
                                            &prefix,
                                            kv.mod_revision(),
                                        );
                                        if let Some(group_id) =
                                            parse_group_id_from_membership_key(kv.key())
                                        {
                                            let meta = match decode_membership_value(
                                                group_id,
                                                kv.value(),
                                            ) {
                                                Ok(decoded) => decoded,
                                                Err(_) => {
                                                    warn!(
                                                        group_id,
                                                        "membership payload decode failed; forcing snapshot resync"
                                                    );
                                                    affected_groups.insert(group_id);
                                                    resync_required = true;
                                                    break;
                                                }
                                            };
                                            if let Err(err) = manager.upsert_group(meta) {
                                                warn!(
                                                    group_id,
                                                    error = %err,
                                                    "membership upsert failed; forcing snapshot resync"
                                                );
                                                affected_groups.insert(group_id);
                                                resync_required = true;
                                                break;
                                            }
                                        }
                                    }
                                }
                                etcd_client::EventType::Delete => {
                                    if let Some(kv) = event.kv() {
                                        apply_mod_revision(&mut manager, kv.mod_revision());
                                        metrics::set_etcd_watch_current_revision(
                                            &prefix,
                                            kv.mod_revision(),
                                        );
                                        if let Some(group_id) =
                                            parse_group_id_from_membership_key(kv.key())
                                            && let Err(err) = manager.remove_group(group_id)
                                        {
                                            warn!(
                                                group_id,
                                                error = %err,
                                                "membership remove failed; forcing snapshot resync"
                                            );
                                            affected_groups.insert(group_id);
                                            resync_required = true;
                                            break;
                                        }
                                    }
                                }
                            }
                        }
                        if !resync_required {
                            finalize_ready_merge_targets(&mut manager, now_ms());
                        }
                    }

                    if resync_required {
                        force_groups_non_serving_for_watch_resync(
                            &service,
                            Some(&affected_groups),
                            "membership snapshot resync required",
                        )
                        .await;
                        break;
                    }
                    healthy_cycle = true;
                }
                Ok(None) => {
                    force_groups_non_serving_for_watch_resync(
                        &service,
                        None,
                        "membership watch stream ended",
                    )
                    .await;
                    break;
                }
                Err(_) => {
                    force_groups_non_serving_for_watch_resync(
                        &service,
                        None,
                        "membership watch stream failed",
                    )
                    .await;
                    break;
                }
            }
        }
        if healthy_cycle {
            retry_attempt = 0;
        } else {
            retry_attempt = retry_attempt.saturating_add(1);
        }
        tokio::time::sleep(watch_retry_sleep_duration(
            cfg.watch_retry_backoff,
            retry_attempt,
            watch_retry_entropy(),
        ))
        .await;
    }
}

async fn sync_membership_snapshot(
    client: &mut Client,
    service: Arc<RwLock<TsoService>>,
    prefix: &str,
) -> Result<i64, ()> {
    let started = Instant::now();
    let resp = client
        .get(prefix, Some(GetOptions::new().with_prefix()))
        .await;
    observe_etcd_result("get_membership_snapshot", started, &resp);
    let resp = resp.map_err(|_| ())?;
    let snapshot_revision = resp.header().map_or(0, |header| header.revision());
    metrics::set_etcd_watch_current_revision(prefix, snapshot_revision);

    let mut groups = Vec::new();
    let mut snapshot_ids = HashSet::new();
    for kv in resp.kvs() {
        if let Some(group_id) = parse_group_id_from_membership_key(kv.key()) {
            snapshot_ids.insert(group_id);
            let group = decode_membership_value(group_id, kv.value())?;
            groups.push(group);
        }
    }

    let guard = service.write().await;
    {
        let mut manager = guard.manager_mut();
        reconcile_membership_snapshot(&mut manager, groups, snapshot_ids, snapshot_revision)?;
    }
    Ok(snapshot_revision.saturating_add(1))
}

async fn run_campaign_loop(
    service: Arc<RwLock<TsoService>>,
    leases: Arc<Mutex<HashMap<u32, i64>>>,
    external_skew_sampler: Arc<Mutex<ExternalClockSkewSampler>>,
    cfg: TsoRuntimeConfig,
) {
    loop {
        let started = Instant::now();
        let connect = Client::connect(cfg.etcd_endpoints.clone(), None).await;
        observe_etcd_result("connect", started, &connect);
        let mut client = match connect {
            Ok(client) => client,
            Err(_) => {
                tokio::time::sleep(cfg.watch_retry_backoff).await;
                continue;
            }
        };

        let mut ticker = tokio::time::interval(cfg.campaign_interval);
        loop {
            ticker.tick().await;
            reconcile_lease_entries_with_active_groups(&service, &leases).await;
            let metas = collect_group_election_meta(&service).await;
            update_runtime_status_metrics(&service, &leases, &cfg, &external_skew_sampler).await;
            if metas.is_empty() {
                continue;
            }

            let mut reconnect = false;
            for meta in metas {
                let expected_primary =
                    match sync_expected_primary(&mut client, &service, &meta).await {
                        Ok(value) => value,
                        Err(_) => {
                            clear_group_lease_tracking(
                                &service,
                                &leases,
                                meta.group_id,
                                true,
                                "expected_primary_sync_error",
                            )
                            .await;
                            reconnect = true;
                            break;
                        }
                    };

                if expected_primary.as_deref().is_some_and(|expected| {
                    !expected_primary_matches_local_identity(expected, &meta.member_value)
                }) {
                    clear_group_lease_tracking(
                        &service,
                        &leases,
                        meta.group_id,
                        true,
                        "expected_primary_mismatch",
                    )
                    .await;
                    metrics::observe_failover_event("expected_primary_mismatch");
                    continue;
                }

                let lease_id = if let Some(existing) = {
                    let lease_guard = leases.lock().await;
                    lease_guard.get(&meta.group_id).copied()
                } {
                    existing
                } else {
                    let lease_started = Instant::now();
                    let lease_grant = client.lease_grant(cfg.lease_ttl_secs, None).await;
                    observe_etcd_result("lease_grant", lease_started, &lease_grant);
                    let lease_id = match lease_grant {
                        Ok(resp) => resp.id(),
                        Err(_) => {
                            clear_group_lease_tracking(
                                &service,
                                &leases,
                                meta.group_id,
                                true,
                                "lease_grant_failed",
                            )
                            .await;
                            reconnect = true;
                            break;
                        }
                    };
                    {
                        let mut lease_guard = leases.lock().await;
                        lease_guard.insert(meta.group_id, lease_id);
                    }
                    {
                        let guard = service.write().await;
                        let mut manager = guard.manager_mut();
                        if let Err(err) = manager.set_group_lease_id(meta.group_id, Some(lease_id))
                        {
                            warn!(
                                group_id = meta.group_id,
                                error = %err,
                                "failed to record granted lease id"
                            );
                        }
                    }
                    lease_id
                };

                let campaign_started = Instant::now();
                match try_campaign_group(&mut client, &meta, lease_id).await {
                    Ok(true) => {
                        let campaign_err = {
                            let guard = service.write().await;
                            let mut manager = guard.manager_mut();
                            manager.campaign_group_primary_external(meta.group_id, now_ms())
                        };
                        if let Err(err) = campaign_err {
                            metrics::observe_primary_takeover_duration(
                                false,
                                campaign_started.elapsed().as_secs_f64(),
                            );
                            warn!(
                                group_id = meta.group_id,
                                error = %err,
                                "campaign succeeded in etcd but local campaign state update failed"
                            );
                            clear_group_lease_tracking(
                                &service,
                                &leases,
                                meta.group_id,
                                true,
                                "campaign_state_update_failed",
                            )
                            .await;
                        } else {
                            metrics::observe_election_campaign();
                            metrics::observe_primary_takeover_duration(
                                true,
                                campaign_started.elapsed().as_secs_f64(),
                            );
                        }
                    }
                    Ok(false) => match load_key_value(&mut client, &meta.leader_key).await {
                        Ok(Some(value)) if value == meta.member_value => {
                            let campaign_err = {
                                let guard = service.write().await;
                                let mut manager = guard.manager_mut();
                                manager.campaign_group_primary_external(meta.group_id, now_ms())
                            };
                            if let Err(err) = campaign_err {
                                metrics::observe_primary_takeover_duration(
                                    false,
                                    campaign_started.elapsed().as_secs_f64(),
                                );
                                warn!(
                                    group_id = meta.group_id,
                                    error = %err,
                                    "failed to refresh local campaign state from observed leader value"
                                );
                                clear_group_lease_tracking(
                                    &service,
                                    &leases,
                                    meta.group_id,
                                    true,
                                    "campaign_refresh_failed",
                                )
                                .await;
                            } else {
                                metrics::observe_election_campaign();
                                metrics::observe_primary_takeover_duration(
                                    true,
                                    campaign_started.elapsed().as_secs_f64(),
                                );
                            }
                        }
                        Ok(_) => {
                            clear_group_lease_tracking(
                                &service,
                                &leases,
                                meta.group_id,
                                true,
                                "lost_campaign",
                            )
                            .await;
                            metrics::observe_failover_event("lost_campaign");
                        }
                        Err(_) => {
                            clear_group_lease_tracking(
                                &service,
                                &leases,
                                meta.group_id,
                                true,
                                "leader_read_failed",
                            )
                            .await;
                            reconnect = true;
                            break;
                        }
                    },
                    Err(_) => {
                        clear_group_lease_tracking(
                            &service,
                            &leases,
                            meta.group_id,
                            true,
                            "campaign_txn_failed",
                        )
                        .await;
                        reconnect = true;
                        break;
                    }
                }
            }

            if reconnect {
                break;
            }
        }

        tokio::time::sleep(cfg.watch_retry_backoff).await;
    }
}

async fn run_lease_refresh_loop(
    service: Arc<RwLock<TsoService>>,
    leases: Arc<Mutex<HashMap<u32, i64>>>,
    external_skew_sampler: Arc<Mutex<ExternalClockSkewSampler>>,
    cfg: TsoRuntimeConfig,
) {
    loop {
        let started = Instant::now();
        let connect = Client::connect(cfg.etcd_endpoints.clone(), None).await;
        observe_etcd_result("connect", started, &connect);
        let mut client = match connect {
            Ok(client) => client,
            Err(_) => {
                tokio::time::sleep(cfg.watch_retry_backoff).await;
                continue;
            }
        };

        let mut keepalive_sessions = HashMap::<u32, LeaseKeepAliveSession>::new();
        let mut ticker = tokio::time::interval(cfg.keepalive_interval);
        loop {
            ticker.tick().await;
            reconcile_lease_entries_with_active_groups(&service, &leases).await;
            prune_stale_keepalive_sessions(&leases, &mut keepalive_sessions).await;
            run_deleted_group_cleanup(&service).await;
            run_prefetch_if_enabled(&service).await;
            let metas = collect_group_election_meta(&service).await;
            update_runtime_status_metrics(&service, &leases, &cfg, &external_skew_sampler).await;
            if metas.is_empty() {
                continue;
            }

            let mut reconnect = false;
            for meta in metas {
                let lease_id = { leases.lock().await.get(&meta.group_id).copied() };
                let Some(lease_id) = lease_id else {
                    keepalive_sessions.remove(&meta.group_id);
                    continue;
                };

                let local_runtime = {
                    let guard = service.read().await;
                    let manager = guard.manager();
                    manager.group_runtime_snapshot(meta.group_id, now_ms())
                };
                let is_primary = match local_runtime {
                    Ok(state) => state.is_primary,
                    Err(err) => {
                        warn!(
                            group_id = meta.group_id,
                            error = %err,
                            "skip lease keepalive because local runtime snapshot failed"
                        );
                        continue;
                    }
                };
                if !is_primary {
                    keepalive_sessions.remove(&meta.group_id);
                    clear_group_lease_tracking(
                        &service,
                        &leases,
                        meta.group_id,
                        false,
                        "local_primary_check_failed",
                    )
                    .await;
                    metrics::observe_failover_event("local_primary_check_failed");
                    continue;
                }

                match keep_group_lease_alive(
                    &mut client,
                    &mut keepalive_sessions,
                    meta.group_id,
                    lease_id,
                )
                .await
                {
                    Ok(true) => {
                        let keepalive_err = {
                            let guard = service.write().await;
                            let mut manager = guard.manager_mut();
                            match manager.keep_group_primary_alive(meta.group_id, now_ms()) {
                                Ok(()) => manager.set_group_lease_id(meta.group_id, Some(lease_id)),
                                Err(err) => Err(err),
                            }
                        };
                        if let Err(err) = keepalive_err {
                            warn!(
                                group_id = meta.group_id,
                                error = %err,
                                "lease keepalive succeeded but local keepalive update failed"
                            );
                            keepalive_sessions.remove(&meta.group_id);
                            clear_group_lease_tracking(
                                &service,
                                &leases,
                                meta.group_id,
                                true,
                                "keepalive_failed",
                            )
                            .await;
                        }
                    }
                    Ok(false) => {
                        keepalive_sessions.remove(&meta.group_id);
                        clear_group_lease_tracking(
                            &service,
                            &leases,
                            meta.group_id,
                            true,
                            "lease_rejected",
                        )
                        .await;
                        metrics::observe_failover_event("lease_rejected");
                    }
                    Err(_) => {
                        keepalive_sessions.remove(&meta.group_id);
                        clear_group_lease_tracking(
                            &service,
                            &leases,
                            meta.group_id,
                            true,
                            "keepalive_error",
                        )
                        .await;
                        reconnect = true;
                        break;
                    }
                }
            }

            if reconnect {
                keepalive_sessions.clear();
                break;
            }
        }

        tokio::time::sleep(cfg.watch_retry_backoff).await;
    }
}

async fn collect_group_election_meta(service: &Arc<RwLock<TsoService>>) -> Vec<GroupElectionMeta> {
    let guard = service.read().await;
    if guard.is_draining() || guard.is_membership_resync_blocking() {
        return Vec::new();
    }
    let mut metas = Vec::new();
    let active_ids = {
        let manager = guard.manager();
        manager.active_group_ids()
    };
    for group_id in active_ids {
        let meta = {
            let manager = guard.manager();
            manager.group_election_meta(group_id)
        };
        match meta {
            Ok(meta) => metas.push(meta),
            Err(err) => {
                warn!(
                    group_id,
                    error = %err,
                    "failed to build group election metadata; skipping group in this cycle"
                );
            }
        }
    }
    metas
}

async fn sync_expected_primary(
    client: &mut Client,
    service: &Arc<RwLock<TsoService>>,
    meta: &GroupElectionMeta,
) -> Result<Option<String>, ()> {
    let expected = load_key_value(client, &meta.expected_primary_key).await?;
    let result = {
        let guard = service.write().await;
        let mut manager = guard.manager_mut();
        manager.apply_group_expected_primary(meta.group_id, expected.clone(), now_ms())
    };
    if let Err(err) = result {
        warn!(
            group_id = meta.group_id,
            error = %err,
            "failed to apply expected-primary metadata to local manager state"
        );
    }
    Ok(expected)
}

async fn reconcile_lease_entries_with_active_groups(
    service: &Arc<RwLock<TsoService>>,
    leases: &Arc<Mutex<HashMap<u32, i64>>>,
) {
    let active_group_ids = {
        let guard = service.read().await;
        let manager = guard.manager();
        manager
            .active_group_ids()
            .into_iter()
            .collect::<HashSet<_>>()
    };
    let stale_group_ids = {
        let mut lease_guard = leases.lock().await;
        let stale = lease_guard
            .keys()
            .copied()
            .filter(|group_id| !active_group_ids.contains(group_id))
            .collect::<Vec<_>>();
        for group_id in &stale {
            lease_guard.remove(group_id);
        }
        stale
    };
    if stale_group_ids.is_empty() {
        return;
    }
    let guard = service.write().await;
    let mut manager = guard.manager_mut();
    for group_id in stale_group_ids {
        if let Err(err) = manager.set_group_lease_id(group_id, None) {
            debug!(
                group_id,
                error = %err,
                "failed to clear lease id while pruning stale runtime lease entry"
            );
        }
    }
}

async fn clear_group_lease_tracking(
    service: &Arc<RwLock<TsoService>>,
    leases: &Arc<Mutex<HashMap<u32, i64>>>,
    group_id: u32,
    step_down_primary: bool,
    restart_reason: &'static str,
) {
    leases.lock().await.remove(&group_id);
    let guard = service.write().await;
    let mut manager = guard.manager_mut();
    if step_down_primary {
        metrics::observe_group_actor_restart(restart_reason);
        if let Err(err) = manager.step_down_group_primary(group_id) {
            warn!(
                group_id,
                error = %err,
                "failed to step down local primary while clearing lease tracking"
            );
        }
    }
    if let Err(err) = manager.set_group_lease_id(group_id, None) {
        warn!(
            group_id,
            error = %err,
            "failed to clear lease id while clearing lease tracking"
        );
    }
}

async fn try_campaign_group(
    client: &mut Client,
    meta: &GroupElectionMeta,
    lease_id: i64,
) -> Result<bool, ()> {
    let txn = Txn::new()
        .when(vec![Compare::create_revision(
            meta.leader_key.clone(),
            CompareOp::Equal,
            0,
        )])
        .and_then(vec![TxnOp::put(
            meta.leader_key.clone(),
            meta.member_value.clone(),
            Some(PutOptions::new().with_lease(lease_id)),
        )]);
    let started = Instant::now();
    let resp = client.txn(txn).await;
    observe_etcd_result("txn_campaign", started, &resp);
    let resp = resp.map_err(|_| ())?;
    Ok(resp.succeeded())
}

async fn open_keepalive_session(
    client: &mut Client,
    lease_id: i64,
) -> Result<LeaseKeepAliveSession, ()> {
    let started = Instant::now();
    let lease_keep_alive = client.lease_keep_alive(lease_id).await;
    observe_etcd_result("lease_keep_alive_open", started, &lease_keep_alive);
    let (keeper, stream) = lease_keep_alive.map_err(|_| ())?;
    Ok(LeaseKeepAliveSession {
        lease_id,
        keeper,
        stream,
    })
}

async fn keep_group_lease_alive(
    client: &mut Client,
    keepalive_sessions: &mut HashMap<u32, LeaseKeepAliveSession>,
    group_id: u32,
    lease_id: i64,
) -> Result<bool, ()> {
    let mut session = match keepalive_sessions.remove(&group_id) {
        Some(existing) if existing.lease_id == lease_id => existing,
        _ => open_keepalive_session(client, lease_id).await?,
    };

    let send_started = Instant::now();
    let send = session.keeper.keep_alive().await;
    observe_etcd_result("lease_keep_alive_send", send_started, &send);
    send.map_err(|_| ())?;

    let recv_started = Instant::now();
    let message = session.stream.message().await;
    observe_etcd_result("lease_keep_alive_recv", recv_started, &message);
    match message.map_err(|_| ())? {
        Some(response) if response.ttl() > 0 => {
            keepalive_sessions.insert(group_id, session);
            Ok(true)
        }
        Some(_) | None => Ok(false),
    }
}

async fn prune_stale_keepalive_sessions(
    leases: &Arc<Mutex<HashMap<u32, i64>>>,
    keepalive_sessions: &mut HashMap<u32, LeaseKeepAliveSession>,
) {
    if keepalive_sessions.is_empty() {
        return;
    }
    let lease_guard = leases.lock().await;
    keepalive_sessions
        .retain(|group_id, session| lease_guard.get(group_id).copied() == Some(session.lease_id));
}

async fn run_deleted_group_cleanup(service: &Arc<RwLock<TsoService>>) {
    let guard = service.write().await;
    let mut manager = guard.manager_mut();
    if let Err(err) = manager.cleanup_deleted_group_timestamps() {
        warn!(
            error = %err,
            "deleted-group timestamp cleanup failed in runtime loop"
        );
    }
}

async fn run_prefetch_if_enabled(service: &Arc<RwLock<TsoService>>) {
    let enabled = {
        let guard = service.read().await;
        guard.runtime_config().enable_async_tso_window_prefetch
    };
    if !enabled {
        return;
    }
    let started = Instant::now();
    let guard = service.read().await;
    if !guard.runtime_config().enable_async_tso_window_prefetch {
        metrics::observe_prefetch_epoch_mismatch();
        metrics::observe_prefetch_stall(1);
        metrics::observe_prefetch_duration(false, started.elapsed().as_secs_f64());
        return;
    }
    let manager = guard.manager();
    let stats = manager.prefetch_serving_group_windows(now_ms());

    metrics::observe_prefetch_stall(stats.stalled_groups);
    metrics::observe_prefetch_error("prefetch_failed", stats.error_groups);
    metrics::observe_prefetch_duration(stats.error_groups == 0, started.elapsed().as_secs_f64());
}

async fn load_key_value(client: &mut Client, key: &str) -> Result<Option<String>, ()> {
    let started = Instant::now();
    let resp = client.get(key, None).await;
    observe_etcd_result("get_value", started, &resp);
    let resp = resp.map_err(|_| ())?;
    Ok(resp
        .kvs()
        .first()
        .map(|kv| String::from_utf8_lossy(kv.value()).to_string()))
}

fn watch_retry_sleep_duration(base: Duration, retry_attempt: u32, entropy: u64) -> Duration {
    let base_ms = (base.as_millis().min(u64::MAX as u128) as u64).max(1);
    let factor = 1_u64 << retry_attempt.min(WATCH_RETRY_BACKOFF_MAX_EXP_SHIFT);
    let max_ms = base_ms.saturating_mul(factor);
    let span = max_ms.saturating_sub(base_ms);
    let jitter = if span == 0 {
        0
    } else {
        entropy % span.saturating_add(1)
    };
    Duration::from_millis(base_ms.saturating_add(jitter))
}

fn watch_retry_entropy() -> u64 {
    let duration = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default();
    let nanos = (duration.as_nanos().min(u64::MAX as u128)) as u64;
    nanos ^ ((std::process::id() as u64) << 32)
}

async fn force_groups_non_serving_for_watch_resync(
    service: &Arc<RwLock<TsoService>>,
    affected_group_ids: Option<&HashSet<u32>>,
    reason: &'static str,
) {
    let guard = service.write().await;
    let groups_to_step_down = {
        let manager = guard.manager();
        if let Some(group_ids) = affected_group_ids {
            if group_ids.is_empty() {
                manager.active_group_ids()
            } else {
                group_ids.iter().copied().collect()
            }
        } else {
            manager.active_group_ids()
        }
    };
    {
        let mut manager = guard.manager_mut();
        for group_id in groups_to_step_down {
            if let Err(err) = manager.step_down_group_primary(group_id) {
                warn!(
                    group_id,
                    error = %err,
                    reason,
                    "failed to step down group for fail-closed membership resync handling"
                );
            }
        }
    }
    guard.set_membership_resync_blocking(true);
}

fn membership_prefix(cluster_id: u64) -> String {
    format!("/pd/{cluster_id}/tso/keyspace_groups/membership/")
}

fn reconcile_membership_snapshot(
    manager: &mut KeyspaceGroupManager,
    groups: Vec<pd_proto_tso::KeyspaceGroup>,
    snapshot_ids: HashSet<u32>,
    snapshot_revision: i64,
) -> Result<(), ()> {
    let active_ids = manager.active_group_ids();
    for group_id in active_ids {
        if !snapshot_ids.contains(&group_id)
            && let Err(err) = manager.remove_group(group_id)
        {
            warn!(
                group_id,
                error = %err,
                "snapshot reconciliation failed to remove stale group"
            );
            return Err(());
        }
    }

    for group in groups {
        let group_id = group.id;
        if let Err(err) = manager.upsert_group(group) {
            warn!(
                group_id,
                error = %err,
                "snapshot reconciliation failed to upsert group"
            );
            return Err(());
        }
    }
    apply_mod_revision(manager, snapshot_revision);
    finalize_ready_merge_targets(manager, now_ms());
    Ok(())
}

fn apply_mod_revision(manager: &mut KeyspaceGroupManager, revision: i64) {
    if revision > 0 {
        manager.set_mod_revision(revision as u64);
    }
}

fn finalize_ready_merge_targets(manager: &mut KeyspaceGroupManager, now_ms: i64) {
    let mut ready_targets = Vec::new();
    for group_id in manager.active_group_ids() {
        let is_merge_target = manager
            .group(group_id)
            .and_then(|group| group.merge_state.as_ref())
            .is_some();
        if !is_merge_target {
            continue;
        }

        let is_primary = manager
            .group_runtime_snapshot(group_id, now_ms)
            .map(|state| state.is_primary)
            .unwrap_or(false);
        if !is_primary {
            continue;
        }

        let all_sources_removed = manager
            .group(group_id)
            .and_then(|group| group.merge_state.as_ref())
            .map(|merge_state| {
                merge_state
                    .merge_list
                    .iter()
                    .all(|source_id| manager.group(*source_id).is_none())
            })
            .unwrap_or(false);
        if all_sources_removed {
            ready_targets.push(group_id);
        }
    }

    for group_id in ready_targets {
        match manager.finish_merge_keyspace_group(group_id, now_ms) {
            Ok(()) => info!(group_id, "finalized merge state for keyspace group"),
            Err(err) => {
                debug!(
                    group_id,
                    error = %err,
                    "merge target finalize attempt did not complete in this cycle"
                );
            }
        }
    }
}

#[derive(Debug, Deserialize)]
struct JsonSplitState {
    #[serde(rename = "split-source")]
    split_source: u32,
}

#[derive(Debug, Deserialize)]
struct JsonMergeState {
    #[serde(rename = "merge-list", default)]
    merge_list: Vec<u32>,
}

#[derive(Debug, Deserialize)]
struct JsonMember {
    address: String,
}

#[derive(Debug, Deserialize)]
struct JsonKeyspaceGroup {
    #[serde(rename = "user-kind", default)]
    user_kind: String,
    #[serde(rename = "split-state")]
    split_state: Option<JsonSplitState>,
    #[serde(rename = "merge-state")]
    merge_state: Option<JsonMergeState>,
    #[serde(default)]
    members: Vec<JsonMember>,
    #[serde(default)]
    keyspaces: Vec<u32>,
}

fn decode_membership_value(group_id: u32, raw: &[u8]) -> Result<pd_proto_tso::KeyspaceGroup, ()> {
    if let Ok(parsed) = serde_json::from_slice::<JsonKeyspaceGroup>(raw) {
        let mut keyspaces = parsed.keyspaces;
        keyspaces.sort_unstable();
        keyspaces.dedup();
        if keyspaces.is_empty() {
            keyspaces.push(group_id);
        }

        return Ok(pd_proto_tso::KeyspaceGroup {
            id: group_id,
            user_kind: if parsed.user_kind.is_empty() {
                "basic".to_string()
            } else {
                parsed.user_kind
            },
            split_state: parsed.split_state.map(|state| pd_proto_tso::SplitState {
                split_source: state.split_source,
            }),
            merge_state: parsed.merge_state.map(|state| pd_proto_tso::MergeState {
                merge_list: state.merge_list,
            }),
            members: parsed
                .members
                .into_iter()
                .map(|member| pd_proto_tso::KeyspaceGroupMember {
                    address: member.address,
                    is_primary: false,
                })
                .collect(),
            keyspaces,
        });
    }

    // Legacy fallback keeps backward compatibility with protobuf-encoded payloads.
    if let Ok(parsed) = tsopb::KeyspaceGroup::decode(raw) {
        return Ok(pd_proto_tso::KeyspaceGroup {
            id: group_id,
            user_kind: parsed.user_kind,
            split_state: parsed.split_state.map(|state| pd_proto_tso::SplitState {
                split_source: state.split_source,
            }),
            merge_state: None,
            members: parsed
                .members
                .into_iter()
                .map(|member| pd_proto_tso::KeyspaceGroupMember {
                    address: member.address,
                    is_primary: member.is_primary,
                })
                .collect(),
            // Current tsopb.KeyspaceGroup payload does not carry explicit keyspace list.
            keyspaces: vec![parsed.id],
        });
    }
    Err(())
}

fn parse_group_id_from_membership_key(key: &[u8]) -> Option<u32> {
    let key = std::str::from_utf8(key).ok()?;
    let suffix = key.rsplit('/').next()?;
    suffix.parse::<u32>().ok()
}

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 observe_etcd_result<T, E>(operation: &str, started: Instant, result: &Result<T, E>) {
    metrics::observe_etcd_request(operation, result.is_ok(), started.elapsed().as_secs_f64());
}

async fn update_runtime_status_metrics(
    service: &Arc<RwLock<TsoService>>,
    leases: &Arc<Mutex<HashMap<u32, i64>>>,
    cfg: &TsoRuntimeConfig,
    external_skew_sampler: &Arc<Mutex<ExternalClockSkewSampler>>,
) {
    let (
        has_primary,
        mailbox_depth,
        clock_skew_seconds,
        total_requests,
        stale_groups,
        physical_bounds_ms,
    ) = {
        let sample_now_ms = now_ms();
        let guard = service.read().await;
        let manager = guard.manager();
        let mut has_primary = false;
        let mut max_abs_skew_ms = 0_u64;
        let mut min_physical_ms = i64::MAX;
        let mut max_physical_ms = i64::MIN;
        let mut has_physical_sample = false;
        for group_id in manager.active_group_ids() {
            if let Ok(state) = manager.group_runtime_snapshot(group_id, sample_now_ms) {
                if state.is_primary {
                    has_primary = true;
                }
                let skew_ms = state.in_memory_timestamp.physical.abs_diff(sample_now_ms);
                max_abs_skew_ms = max_abs_skew_ms.max(skew_ms);
                min_physical_ms = min_physical_ms.min(state.in_memory_timestamp.physical);
                max_physical_ms = max_physical_ms.max(state.in_memory_timestamp.physical);
                has_physical_sample = true;
            }
        }
        let (total_requests, stale_groups) =
            manager.assigned_group_qps_inputs(sample_now_ms, ASSIGNED_GROUP_QPS_STALE_THRESHOLD_MS);
        (
            has_primary,
            manager.total_in_flight_requests(),
            (max_abs_skew_ms as f64) / 1_000.0,
            total_requests,
            stale_groups,
            if has_physical_sample {
                Some((min_physical_ms, max_physical_ms))
            } else {
                None
            },
        )
    };

    let has_lease = !leases.lock().await.is_empty();
    metrics::set_group_actor_mailbox_size(mailbox_depth);
    metrics::set_election_and_lease_status(has_primary, has_lease);
    metrics::set_physical_time_clock_skew_seconds(clock_skew_seconds);
    maybe_update_external_clock_skew_metrics(cfg, physical_bounds_ms, external_skew_sampler).await;
    observe_assigned_group_qps_metrics(total_requests, stale_groups);
}

async fn maybe_update_external_clock_skew_metrics(
    cfg: &TsoRuntimeConfig,
    physical_bounds_ms: Option<(i64, i64)>,
    sampler: &Arc<Mutex<ExternalClockSkewSampler>>,
) {
    let Some(ntp_server) = cfg.external_ntp_server.as_deref() else {
        return;
    };
    let Some(physical_bounds_ms) = physical_bounds_ms else {
        return;
    };

    let now = Instant::now();
    {
        let mut guard = sampler.lock().await;
        if now < guard.next_poll_at {
            if let Some(skew_seconds) = guard.last_observed_skew_seconds {
                metrics::set_physical_time_external_clock_skew_seconds(skew_seconds);
            }
            return;
        }
        guard.next_poll_at = now + cfg.external_ntp_refresh_interval;
    }

    match sample_external_ntp_time_millis(ntp_server, cfg.external_ntp_timeout).await {
        Ok(external_time_ms) => {
            let skew_seconds = external_clock_skew_seconds(physical_bounds_ms, external_time_ms);
            {
                let mut guard = sampler.lock().await;
                guard.last_observed_skew_seconds = Some(skew_seconds);
            }
            metrics::set_physical_time_external_clock_skew_seconds(skew_seconds);
        }
        Err(err) => {
            metrics::observe_physical_time_external_clock_source_error();
            warn!(
                ntp_server,
                error = %err,
                "external NTP skew sampling failed"
            );
            if let Some(skew_seconds) = sampler.lock().await.last_observed_skew_seconds {
                metrics::set_physical_time_external_clock_skew_seconds(skew_seconds);
            }
        }
    }
}

fn external_clock_skew_seconds(physical_bounds_ms: (i64, i64), external_time_ms: i64) -> f64 {
    let (min_physical_ms, max_physical_ms) = physical_bounds_ms;
    let min_skew_ms = min_physical_ms.abs_diff(external_time_ms);
    let max_skew_ms = max_physical_ms.abs_diff(external_time_ms);
    (min_skew_ms.max(max_skew_ms) as f64) / 1_000.0
}

async fn sample_external_ntp_time_millis(
    ntp_server: &str,
    timeout: Duration,
) -> Result<i64, String> {
    let server = ntp_server.trim();
    if server.is_empty() {
        return Err("external NTP server is empty".to_string());
    }

    let socket = tokio::time::timeout(timeout, tokio::net::UdpSocket::bind("0.0.0.0:0"))
        .await
        .map_err(|_| "NTP socket bind timed out".to_string())?
        .map_err(|err| format!("NTP socket bind failed: {err}"))?;
    tokio::time::timeout(timeout, socket.connect(server))
        .await
        .map_err(|_| "NTP connect timed out".to_string())?
        .map_err(|err| format!("NTP connect failed: {err}"))?;

    let mut request = [0_u8; NTP_PACKET_SIZE];
    request[0] = 0x1B;
    let sent = tokio::time::timeout(timeout, socket.send(&request))
        .await
        .map_err(|_| "NTP request send timed out".to_string())?
        .map_err(|err| format!("NTP request send failed: {err}"))?;
    if sent != NTP_PACKET_SIZE {
        return Err(format!(
            "NTP request expected {NTP_PACKET_SIZE} bytes, sent {sent}"
        ));
    }

    let mut response = [0_u8; NTP_PACKET_SIZE];
    let received = tokio::time::timeout(timeout, socket.recv(&mut response))
        .await
        .map_err(|_| "NTP response recv timed out".to_string())?
        .map_err(|err| format!("NTP response recv failed: {err}"))?;
    if received < NTP_PACKET_SIZE {
        return Err(format!(
            "NTP response expected at least {NTP_PACKET_SIZE} bytes, got {received}"
        ));
    }

    decode_ntp_unix_time_millis(&response)
        .ok_or_else(|| "NTP response had invalid transmit timestamp".to_string())
}

fn decode_ntp_unix_time_millis(packet: &[u8]) -> Option<i64> {
    if packet.len() < NTP_PACKET_SIZE {
        return None;
    }

    let secs = u32::from_be_bytes(packet[40..44].try_into().ok()?) as i64;
    let frac = u32::from_be_bytes(packet[44..48].try_into().ok()?) as u64;
    if secs < NTP_UNIX_EPOCH_OFFSET_SECS {
        return None;
    }

    let unix_secs = secs - NTP_UNIX_EPOCH_OFFSET_SECS;
    let frac_ms = ((frac * 1_000) >> 32) as i64;
    unix_secs.checked_mul(1_000)?.checked_add(frac_ms)
}

fn observe_assigned_group_qps_metrics(total_requests: u64, stale_groups: u32) {
    let sampled_at = Instant::now();
    let qps = {
        let mut sample = ASSIGNED_GROUP_QPS_SAMPLE
            .lock()
            .unwrap_or_else(|poison| poison.into_inner());
        let previous = sample.replace(AssignedGroupQpsSample {
            observed_at: sampled_at,
            total_requests,
        });
        if let Some(previous) = previous {
            let elapsed = sampled_at
                .saturating_duration_since(previous.observed_at)
                .as_secs_f64();
            if elapsed > 0.0 {
                (total_requests.saturating_sub(previous.total_requests) as f64) / elapsed
            } else {
                0.0
            }
        } else {
            0.0
        }
    };
    metrics::set_assigned_group_qps_on_node(qps);
    metrics::set_assigned_group_qps_stale_groups(u64::from(stale_groups));
}

#[cfg(test)]
mod tests {
    use std::collections::{HashMap, HashSet};
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};

    use crate::allocator::AllocatorConfig;
    use crate::keypath::timestamp_path;
    use crate::manager::KeyspaceGroupManager;
    use crate::metrics;
    use crate::persistence::{
        InMemoryEtcdStore, StoreError, TimestampStore, encode_timestamp_nanos,
    };
    use crate::service::TsoService;
    use tokio::sync::{Mutex, RwLock};

    use super::{
        NTP_PACKET_SIZE, NTP_UNIX_EPOCH_OFFSET_SECS, TsoRuntimeConfig, clear_group_lease_tracking,
        decode_membership_value, decode_ntp_unix_time_millis, external_clock_skew_seconds,
        finalize_ready_merge_targets, force_groups_non_serving_for_watch_resync, membership_prefix,
        parse_group_id_from_membership_key, reconcile_lease_entries_with_active_groups,
        reconcile_membership_snapshot, run_deleted_group_cleanup, run_prefetch_if_enabled,
        watch_retry_sleep_duration,
    };

    #[derive(Default)]
    struct SaveFailingStore {
        inner: InMemoryEtcdStore,
        save_calls: AtomicUsize,
    }

    impl TimestampStore for SaveFailingStore {
        fn load_timestamp(
            &self,
            cluster_id: u64,
            group_id: u32,
        ) -> Result<Option<i64>, StoreError> {
            self.inner.load_timestamp(cluster_id, group_id)
        }

        fn save_timestamp_if_leader(
            &self,
            cluster_id: u64,
            group_id: u32,
            ts_nanos: i64,
            leader_key: &str,
            expected_leader_value: &str,
        ) -> Result<(), StoreError> {
            let save_call = self.save_calls.fetch_add(1, AtomicOrdering::AcqRel);
            if save_call > 0 {
                return Err(StoreError::Internal("forced save failure".to_string()));
            }
            self.inner.save_timestamp_if_leader(
                cluster_id,
                group_id,
                ts_nanos,
                leader_key,
                expected_leader_value,
            )
        }

        fn delete_timestamp(&self, cluster_id: u64, group_id: u32) -> Result<(), StoreError> {
            self.inner.delete_timestamp(cluster_id, group_id)
        }

        fn set_leader_value(&self, leader_key: &str, leader_value: &str) -> Result<(), StoreError> {
            self.inner.set_leader_value(leader_key, leader_value)
        }
    }

    async fn enable_prefetch(service: &Arc<RwLock<TsoService>>) {
        let guard = service.write().await;
        let mut config = guard.runtime_config();
        config.enable_async_tso_window_prefetch = true;
        guard.apply_runtime_config(config);
    }

    fn build_prefetch_failure_service() -> Arc<RwLock<TsoService>> {
        let store = Arc::new(SaveFailingStore::default());
        let mut allocator = AllocatorConfig::default();
        allocator.oracle.save_interval_ms = 1;
        allocator.oracle.update_physical_interval_ms = 1;
        let mut manager = KeyspaceGroupManager::new(1, store, allocator, "node");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 0,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![0],
            })
            .expect("register group 0");
        manager
            .campaign_group_primary(0, super::now_ms())
            .expect("campaign group 0");
        Arc::new(RwLock::new(TsoService::new(1, manager)))
    }

    #[test]
    fn parse_group_id_handles_expected_path() {
        let key = b"/pd/1/tso/keyspace_groups/membership/00007";
        assert_eq!(parse_group_id_from_membership_key(key), Some(7));
    }

    #[test]
    fn runtime_defaults_are_non_zero() {
        let cfg = TsoRuntimeConfig::default();
        assert!(!cfg.etcd_endpoints.is_empty());
        assert!(cfg.lease_ttl_secs > 0);
        assert!(cfg.external_ntp_server.is_none());
        assert!(cfg.external_ntp_refresh_interval > std::time::Duration::ZERO);
        assert!(cfg.external_ntp_timeout > std::time::Duration::ZERO);
        assert_eq!(
            membership_prefix(42),
            "/pd/42/tso/keyspace_groups/membership/"
        );
    }

    #[test]
    fn decode_ntp_unix_time_millis_parses_transmit_timestamp() {
        let unix_ms = 1_700_000_123_456_i64;
        let unix_secs = unix_ms / 1_000;
        let unix_frac_ms = (unix_ms % 1_000) as u64;
        let ntp_secs = (unix_secs + NTP_UNIX_EPOCH_OFFSET_SECS) as u32;
        let ntp_frac = (unix_frac_ms << 32) / 1_000;

        let mut packet = [0_u8; NTP_PACKET_SIZE];
        packet[40..44].copy_from_slice(&ntp_secs.to_be_bytes());
        packet[44..48].copy_from_slice(&(ntp_frac as u32).to_be_bytes());

        let decoded = decode_ntp_unix_time_millis(&packet).expect("decoded NTP timestamp");
        assert!(
            (decoded - unix_ms).abs() <= 1,
            "decoded={decoded}, expected={unix_ms}"
        );
    }

    #[test]
    fn external_clock_skew_uses_furthest_physical_bound() {
        let centered = external_clock_skew_seconds((1_000, 1_400), 1_200);
        assert!((centered - 0.2).abs() < f64::EPSILON);

        let shifted = external_clock_skew_seconds((1_000, 1_400), 500);
        assert!((shifted - 0.9).abs() < f64::EPSILON);
    }

    #[test]
    fn watch_retry_backoff_uses_jittered_exponential_bounds() {
        let base = std::time::Duration::from_millis(100);
        let first = watch_retry_sleep_duration(base, 1, 0);
        let first_max = watch_retry_sleep_duration(base, 1, 100);
        let third = watch_retry_sleep_duration(base, 3, 700);
        let capped = watch_retry_sleep_duration(base, 99, 6_300);

        assert!(first >= base);
        assert!(first <= std::time::Duration::from_millis(200));
        assert!(first_max <= std::time::Duration::from_millis(200));

        assert!(third >= base);
        assert!(third <= std::time::Duration::from_millis(800));

        // Capped at base * 2^6 = 6400ms.
        assert!(capped >= base);
        assert!(capped <= std::time::Duration::from_millis(6_400));
    }

    #[test]
    fn decode_membership_value_prefers_json_endpoint_shape() {
        let raw = br#"{
          "id":7,
          "user-kind":"standard",
          "split-state":{"split-source":5},
          "merge-state":{"merge-list":[2,3]},
          "members":[{"address":"127.0.0.1:1234"}],
          "keyspaces":[20,10,20]
        }"#;
        let parsed = decode_membership_value(7, raw).expect("decode json");
        assert_eq!(parsed.id, 7);
        assert_eq!(parsed.user_kind, "standard");
        assert_eq!(parsed.split_state.expect("split").split_source, 5);
        assert_eq!(parsed.merge_state.expect("merge").merge_list, vec![2, 3]);
        assert_eq!(parsed.members.len(), 1);
        assert_eq!(parsed.members[0].address, "127.0.0.1:1234");
        assert_eq!(parsed.keyspaces, vec![10, 20]);
    }

    #[test]
    fn decode_membership_value_rejects_malformed_payload() {
        assert!(decode_membership_value(7, b"not-json-and-not-proto").is_err());
    }

    #[test]
    fn reconcile_membership_snapshot_removes_stale_groups_and_advances_revision() {
        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::new(),
                keyspaces: vec![0],
            })
            .expect("register g0");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 1,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![1],
            })
            .expect("register g1");

        let mut snapshot_ids = HashSet::new();
        snapshot_ids.insert(0);
        reconcile_membership_snapshot(
            &mut manager,
            vec![pd_proto_tso::KeyspaceGroup {
                id: 0,
                user_kind: "standard".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![0],
            }],
            snapshot_ids,
            123,
        )
        .expect("reconcile snapshot");

        assert!(manager.group(1).is_none());
        assert_eq!(manager.group(0).expect("group 0").user_kind, "standard");
        assert_eq!(manager.mod_revision(), 123);
    }

    #[test]
    fn finalize_ready_merge_targets_clears_merge_state_when_sources_are_absent() {
        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::new(),
                keyspaces: vec![0],
            })
            .expect("register merge target");
        manager
            .campaign_group_primary(0, 10_000)
            .expect("campaign target");
        manager
            .upsert_group(pd_proto_tso::KeyspaceGroup {
                id: 0,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: Some(pd_proto_tso::MergeState {
                    merge_list: vec![1],
                }),
                members: Vec::new(),
                keyspaces: vec![0, 1],
            })
            .expect("mark merge state");

        finalize_ready_merge_targets(&mut manager, 11_000);

        assert!(
            manager
                .group(0)
                .expect("target group")
                .merge_state
                .is_none()
        );
    }

    #[tokio::test]
    async fn runtime_deleted_group_cleanup_loop_removes_timestamp_keys() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut manager =
            KeyspaceGroupManager::new(1, store.clone(), 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::new(),
                keyspaces: vec![0],
            })
            .expect("register target");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 1,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![1],
            })
            .expect("register source");
        manager
            .campaign_group_primary(0, 10_000)
            .expect("campaign target");
        manager
            .campaign_group_primary(1, 10_000)
            .expect("campaign source");
        manager.merge_keyspace_groups(0, vec![1]).expect("merge");

        store
            .put_raw(
                timestamp_path(1, 1),
                encode_timestamp_nanos(12_000_000_000).to_vec(),
            )
            .expect("seed source timestamp");

        manager
            .finish_merge_keyspace_group(0, 11_000)
            .expect("finalize merge");

        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));
        run_deleted_group_cleanup(&service).await;

        let guard = service.read().await;
        assert!(guard.manager().deleted_groups().is_empty());
        drop(guard);
        assert!(
            store
                .load_timestamp(1, 1)
                .expect("load source timestamp")
                .is_none(),
            "cleanup loop should remove merged source timestamp key"
        );
    }

    #[tokio::test]
    async fn runtime_cleanup_does_not_remove_merge_source_before_finalize() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut manager =
            KeyspaceGroupManager::new(1, store.clone(), 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::new(),
                keyspaces: vec![0],
            })
            .expect("register target");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 1,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![1],
            })
            .expect("register source");
        manager.merge_keyspace_groups(0, vec![1]).expect("merge");

        store
            .put_raw(
                timestamp_path(1, 1),
                encode_timestamp_nanos(13_000_000_000).to_vec(),
            )
            .expect("seed source timestamp");

        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));
        run_deleted_group_cleanup(&service).await;

        assert!(
            store
                .load_timestamp(1, 1)
                .expect("load source timestamp")
                .is_some(),
            "cleanup must wait until merge finalize before deleting source key"
        );
    }

    #[tokio::test]
    async fn runtime_reconciles_stale_lease_entries_against_active_groups() {
        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::new(),
                keyspaces: vec![0],
            })
            .expect("register group 0");
        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));
        let leases = Arc::new(Mutex::new(HashMap::from([
            (0_u32, 11_i64),
            (9_u32, 99_i64),
        ])));

        reconcile_lease_entries_with_active_groups(&service, &leases).await;

        let lease_guard = leases.lock().await;
        assert!(lease_guard.contains_key(&0));
        assert!(!lease_guard.contains_key(&9));
    }

    #[tokio::test]
    async fn runtime_clear_group_lease_tracking_clears_map_and_primary_state() {
        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::new(),
                keyspaces: vec![0],
            })
            .expect("register group 0");
        manager
            .campaign_group_primary(0, super::now_ms())
            .expect("campaign group 0");
        manager
            .set_group_lease_id(0, Some(123))
            .expect("set lease id");
        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));
        let leases = Arc::new(Mutex::new(HashMap::from([(0_u32, 123_i64)])));

        clear_group_lease_tracking(&service, &leases, 0, true, "test").await;

        let lease_guard = leases.lock().await;
        assert!(!lease_guard.contains_key(&0));
        drop(lease_guard);

        let guard = service.read().await;
        let state = guard
            .manager()
            .group_debug_state(0, super::now_ms())
            .expect("group state");
        assert!(!state.is_primary);
        assert_eq!(state.lease_id, 0);
        let names = metrics::gather_metric_names();
        assert!(
            names
                .iter()
                .any(|name| name == "group_actor_restarts_total")
        );
    }

    #[tokio::test]
    async fn runtime_watch_resync_block_steps_down_affected_groups_only() {
        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::new(),
                keyspaces: vec![0],
            })
            .expect("register group 0");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 1,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![1],
            })
            .expect("register group 1");
        manager
            .campaign_group_primary(0, super::now_ms())
            .expect("campaign group 0");
        manager
            .campaign_group_primary(1, super::now_ms())
            .expect("campaign group 1");
        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));

        let affected = HashSet::from([0_u32]);
        force_groups_non_serving_for_watch_resync(
            &service,
            Some(&affected),
            "unit test resync affected group",
        )
        .await;

        let guard = service.read().await;
        assert!(guard.is_membership_resync_blocking());
        let state0 = guard
            .manager()
            .group_debug_state(0, super::now_ms())
            .expect("group 0 state");
        let state1 = guard
            .manager()
            .group_debug_state(1, super::now_ms())
            .expect("group 1 state");
        assert!(!state0.is_primary);
        assert!(state1.is_primary);
    }

    #[tokio::test]
    async fn runtime_watch_resync_block_without_scope_steps_down_all_groups() {
        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::new(),
                keyspaces: vec![0],
            })
            .expect("register group 0");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 1,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![1],
            })
            .expect("register group 1");
        manager
            .campaign_group_primary(0, super::now_ms())
            .expect("campaign group 0");
        manager
            .campaign_group_primary(1, super::now_ms())
            .expect("campaign group 1");
        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));

        force_groups_non_serving_for_watch_resync(&service, None, "unit test global resync").await;

        let guard = service.read().await;
        assert!(guard.is_membership_resync_blocking());
        let state0 = guard
            .manager()
            .group_debug_state(0, super::now_ms())
            .expect("group 0 state");
        let state1 = guard
            .manager()
            .group_debug_state(1, super::now_ms())
            .expect("group 1 state");
        assert!(!state0.is_primary);
        assert!(!state1.is_primary);
    }

    #[tokio::test]
    async fn runtime_prefetch_runs_only_when_enabled() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut allocator = AllocatorConfig::default();
        allocator.oracle.update_physical_interval_ms = 1;
        let mut manager = KeyspaceGroupManager::new(1, store, allocator, "node");
        manager
            .register_group(pd_proto_tso::KeyspaceGroup {
                id: 0,
                user_kind: "basic".to_string(),
                split_state: None,
                merge_state: None,
                members: Vec::new(),
                keyspaces: vec![0],
            })
            .expect("register group 0");
        manager
            .campaign_group_primary(0, super::now_ms())
            .expect("campaign group 0");
        let service = Arc::new(RwLock::new(TsoService::new(1, manager)));

        let before = {
            let guard = service.read().await;
            let state = guard
                .manager()
                .group_debug_state(0, super::now_ms())
                .expect("group state before prefetch");
            state.in_memory_timestamp.physical
        };

        run_prefetch_if_enabled(&service).await;

        let after_disabled = {
            let guard = service.read().await;
            let state = guard
                .manager()
                .group_debug_state(0, super::now_ms())
                .expect("group state with prefetch disabled");
            state.in_memory_timestamp.physical
        };
        assert_eq!(
            after_disabled, before,
            "prefetch-disabled path must not mutate allocator window"
        );

        let duration_samples_before = metrics::TSO_PREFETCH_DURATION_SECONDS
            .with_label_values(&["success"])
            .get_sample_count();
        tokio::time::sleep(std::time::Duration::from_millis(3)).await;
        enable_prefetch(&service).await;
        run_prefetch_if_enabled(&service).await;

        let after_enabled = {
            let guard = service.read().await;
            let state = guard
                .manager()
                .group_debug_state(0, super::now_ms())
                .expect("group state with prefetch enabled");
            state.in_memory_timestamp.physical
        };
        assert!(
            after_enabled > after_disabled,
            "enabled prefetch should advance allocator state on interval boundary"
        );
        assert!(
            metrics::TSO_PREFETCH_DURATION_SECONDS
                .with_label_values(&["success"])
                .get_sample_count()
                > duration_samples_before,
            "enabled prefetch should record a duration sample"
        );
    }

    #[tokio::test]
    async fn runtime_prefetch_records_stall_and_error_metrics_on_failed_prefetch() {
        let service = build_prefetch_failure_service();
        let stalled_before = metrics::TSO_PREFETCH_STALL_TOTAL.get();
        let errors_before = metrics::TSO_PREFETCH_ERROR_TOTAL
            .with_label_values(&["prefetch_failed"])
            .get();

        tokio::time::sleep(std::time::Duration::from_millis(3)).await;
        enable_prefetch(&service).await;
        run_prefetch_if_enabled(&service).await;

        let stalled_after = metrics::TSO_PREFETCH_STALL_TOTAL.get();
        let errors_after = metrics::TSO_PREFETCH_ERROR_TOTAL
            .with_label_values(&["prefetch_failed"])
            .get();
        assert!(
            stalled_after > stalled_before,
            "failed prefetch should increment stall metric"
        );
        assert!(
            errors_after > errors_before,
            "failed prefetch should increment prefetch_failed error metric"
        );

        let guard = service.read().await;
        let state = guard
            .manager()
            .group_debug_state(0, super::now_ms())
            .expect("group state after failed prefetch");
        assert_eq!(state.prefetch_state, "Backoff");
        assert!(!state.prefetch_last_error.is_empty());
    }
}