use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicI64, AtomicU32, AtomicU64, Ordering as AtomicOrdering};
use std::sync::{Arc, Mutex};

use pd_proto_tso::{KeyspaceGroup, MergeState, SplitState, Timestamp};

use crate::allocator::{AllocatorConfig, GroupAllocator};
use crate::error::ServiceError;
use crate::keypath::{DEFAULT_KEYSPACE_GROUP_ID, MAX_KEYSPACE_GROUP_COUNT_IN_USE};
use crate::persistence::TimestampStore;

const DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP: u32 = 64;
const OVERLOAD_RETRY_AFTER_MS: u64 = 100;

struct GroupState {
    meta: KeyspaceGroup,
    keyspace_set: HashSet<u32>,
    runtime: Arc<Mutex<GroupRuntimeState>>,
    in_flight_requests: Arc<AtomicU32>,
    request_count_total: Arc<AtomicU64>,
    last_request_ms: Arc<AtomicI64>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PrefetchDiagnosticState {
    Idle,
    InFlight,
    Backoff,
}

impl PrefetchDiagnosticState {
    fn as_str(self) -> &'static str {
        match self {
            Self::Idle => "Idle",
            Self::InFlight => "InFlight",
            Self::Backoff => "Backoff",
        }
    }
}

struct GroupRuntimeState {
    allocator: GroupAllocator,
    initialized: bool,
    active: bool,
    lease_id: i64,
    config_epoch: u64,
    prefetch_state: PrefetchDiagnosticState,
    prefetch_last_attempt_ms: i64,
    prefetch_last_error: String,
}

#[derive(Default)]
struct KeyspaceRouting {
    by_keyspace: HashMap<u32, u32>,
}

impl KeyspaceRouting {
    fn bind_group(&mut self, group_id: u32, keyspaces: &[u32]) {
        for keyspace in keyspaces {
            self.by_keyspace.insert(*keyspace, group_id);
        }
    }

    fn rebind_keyspace(&mut self, keyspace_id: u32, group_id: u32) {
        self.by_keyspace.insert(keyspace_id, group_id);
    }

    fn unbind_group_if_owned(&mut self, group_id: u32, keyspaces: &[u32]) {
        for keyspace in keyspaces {
            if self.by_keyspace.get(keyspace).copied() == Some(group_id) {
                self.by_keyspace.remove(keyspace);
            }
        }
    }

    fn resolve(&self, keyspace_id: u32) -> Option<u32> {
        self.by_keyspace.get(&keyspace_id).copied()
    }
}

#[derive(Default)]
struct DeletedGroupTracker {
    groups: HashSet<u32>,
}

impl DeletedGroupTracker {
    fn mark(&mut self, group_id: u32) {
        self.groups.insert(group_id);
    }

    fn list(&self) -> Vec<u32> {
        self.groups.iter().copied().collect()
    }

    fn snapshot(&self) -> Vec<u32> {
        self.groups.iter().copied().collect()
    }

    fn remove(&mut self, group_id: u32) {
        self.groups.remove(&group_id);
    }
}

impl GroupState {
    fn new(meta: KeyspaceGroup, allocator: GroupAllocator) -> Self {
        let keyspace_set = meta.keyspaces.iter().copied().collect();
        let runtime = Arc::new(Mutex::new(GroupRuntimeState {
            allocator,
            initialized: false,
            active: true,
            lease_id: 0,
            config_epoch: 0,
            prefetch_state: PrefetchDiagnosticState::Idle,
            prefetch_last_attempt_ms: 0,
            prefetch_last_error: String::new(),
        }));
        let in_flight_requests = Arc::new(AtomicU32::new(0));
        let request_count_total = Arc::new(AtomicU64::new(0));
        let last_request_ms = Arc::new(AtomicI64::new(0));
        Self {
            meta,
            keyspace_set,
            runtime,
            in_flight_requests,
            request_count_total,
            last_request_ms,
        }
    }
}

/// Prepared request context for TSO fast path.
#[derive(Clone)]
pub(crate) struct PreparedTsoRequest {
    resolved_group_id: u32,
    runtime: Arc<Mutex<GroupRuntimeState>>,
    in_flight_requests: Arc<AtomicU32>,
    request_count_total: Arc<AtomicU64>,
    last_request_ms: Arc<AtomicI64>,
    requires_slow_path: bool,
}

impl PreparedTsoRequest {
    #[must_use]
    pub fn resolved_group_id(&self) -> u32 {
        self.resolved_group_id
    }

    #[must_use]
    pub fn requires_slow_path(&self) -> bool {
        self.requires_slow_path
    }

    /// Allocates one timestamp range on the prepared group runtime.
    pub fn generate_tso(
        &self,
        now_ms: i64,
        count: u32,
        overload_hardening: bool,
        overload_max_in_flight_per_group: u32,
    ) -> Result<Timestamp, ServiceError> {
        let _in_flight_guard = acquire_in_flight_request(
            &self.in_flight_requests,
            self.resolved_group_id,
            overload_hardening,
            overload_max_in_flight_per_group,
        )?;
        let mut runtime = lock_runtime(&self.runtime);
        if !runtime.active {
            return Err(ServiceError::NotServed {
                group_id: self.resolved_group_id,
            });
        }
        ensure_runtime_initialized(&mut runtime, now_ms, self.resolved_group_id)?;
        let timestamp = runtime.allocator.generate_tso(now_ms, count)?;
        drop(runtime);
        record_group_request_sample(&self.request_count_total, &self.last_request_ms, now_ms);
        Ok(timestamp)
    }
}

/// Result of keyspace-group lookup.
#[derive(Debug, Clone)]
pub struct GroupLookupResult {
    /// Resolved keyspace-group metadata.
    pub group: KeyspaceGroup,
    /// Group id selected by routing/lookup.
    pub resolved_group_id: u32,
    /// Membership `mod_revision` used for the lookup decision.
    pub mod_revision: u64,
}

/// Aggregated result for `GetMinTS`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MinTsSummary {
    /// Minimum timestamp across currently served groups.
    pub timestamp: Timestamp,
    /// Number of groups currently serving and included in the scan.
    pub keyspace_groups_serving: u32,
    /// Total groups observed by the manager.
    pub keyspace_groups_total: u32,
}

/// Read-only group-state snapshot for admin surfaces.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GroupDebugState {
    /// Whether this node is primary for the group.
    pub is_primary: bool,
    /// Current in-memory timestamp state.
    pub in_memory_timestamp: Timestamp,
    /// Etcd lease id currently bound to local primary state.
    pub lease_id: i64,
    /// Lease expiration in milliseconds.
    pub lease_expire_at_ms: i64,
    /// Persisted save-window upper bound in milliseconds.
    pub save_window_end_ms: i64,
    /// Last persisted physical milliseconds component.
    pub last_persisted_physical_ms: i64,
    /// Last persisted logical component.
    pub last_persisted_logical: i64,
    /// Current actor mailbox depth for this group.
    pub mailbox_depth: u32,
    /// Current in-flight request count for this group.
    pub in_flight_requests: u32,
    /// Human-readable prefetch state.
    pub prefetch_state: String,
    /// Unix milliseconds of the last prefetch attempt.
    pub prefetch_last_attempt_ms: i64,
    /// Last prefetch error message, if any.
    pub prefetch_last_error: String,
}

/// Runtime-only group snapshot used by HA and metrics loops.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct GroupRuntimeSnapshot {
    /// Whether this node is currently primary for the group.
    pub is_primary: bool,
    /// Current in-memory timestamp state.
    pub in_memory_timestamp: Timestamp,
}

/// Aggregated prefetch execution stats for one runtime tick.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PrefetchRunStats {
    /// Number of groups considered for prefetch this tick.
    pub attempted_groups: u32,
    /// Number of groups successfully prefetched.
    pub completed_groups: u32,
    /// Number of groups that stalled and fell back to sync path.
    pub stalled_groups: u32,
    /// Number of groups that failed prefetch.
    pub error_groups: u32,
}

/// Election metadata required by runtime HA loops.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GroupElectionMeta {
    /// Group id for campaign/lease operations.
    pub group_id: u32,
    /// Etcd leader key for CAS campaign operations.
    pub leader_key: String,
    /// Etcd expected-primary key used for transfer/fencing checks.
    pub expected_primary_key: String,
    /// Serialized local member identity value written into etcd.
    pub member_value: String,
}

/// In-memory manager for keyspace-group routing and allocator ownership.
pub struct KeyspaceGroupManager {
    cluster_id: u64,
    store: Arc<dyn TimestampStore>,
    allocator_config: AllocatorConfig,
    member_name_prefix: String,
    discovery_tso_urls: Vec<String>,

    entries: Vec<Option<GroupState>>,
    routing: KeyspaceRouting,
    deleted_groups: DeletedGroupTracker,
    mod_revision: u64,
    route_epoch: Arc<AtomicU64>,
}

impl KeyspaceGroupManager {
    /// Creates a manager with fixed in-memory slot capacity for strict serving-range parity.
    #[must_use]
    pub fn new(
        cluster_id: u64,
        store: Arc<dyn TimestampStore>,
        allocator_config: AllocatorConfig,
        member_name_prefix: impl Into<String>,
    ) -> Self {
        let size = MAX_KEYSPACE_GROUP_COUNT_IN_USE as usize;
        Self {
            cluster_id,
            store,
            allocator_config,
            member_name_prefix: member_name_prefix.into(),
            discovery_tso_urls: Vec::new(),
            entries: (0..size).map(|_| None).collect(),
            routing: KeyspaceRouting::default(),
            deleted_groups: DeletedGroupTracker::default(),
            mod_revision: 0,
            route_epoch: Arc::new(AtomicU64::new(1)),
        }
    }

    /// Returns the latest observed membership mod-revision applied to routing state.
    #[must_use]
    pub fn mod_revision(&self) -> u64 {
        self.mod_revision
    }

    /// Returns the current route-epoch used by request-path cache revalidation.
    #[must_use]
    pub fn route_epoch(&self) -> u64 {
        self.route_epoch.load(AtomicOrdering::Acquire)
    }

    /// Returns a shared route-epoch counter for lock-free request-path checks.
    #[must_use]
    pub fn route_epoch_counter(&self) -> Arc<AtomicU64> {
        Arc::clone(&self.route_epoch)
    }

    /// Returns current allocator configuration used by existing/new groups.
    #[must_use]
    pub fn allocator_config(&self) -> AllocatorConfig {
        self.allocator_config
    }

    /// Advances the in-memory membership revision fence monotonically.
    pub fn set_mod_revision(&mut self, revision: u64) {
        if self.mod_revision < revision {
            self.mod_revision = revision;
            self.bump_route_epoch();
        }
    }

    /// Returns immutable metadata for one served group id, if present.
    #[must_use]
    pub fn group(&self, group_id: u32) -> Option<&KeyspaceGroup> {
        self.state(group_id).ok().map(|state| &state.meta)
    }

    /// Returns active in-memory group ids currently tracked by this node.
    #[must_use]
    pub fn active_group_ids(&self) -> Vec<u32> {
        let mut ids = Vec::new();
        for (group_id, entry) in self.entries.iter().enumerate() {
            if entry.is_some() {
                ids.push(group_id as u32);
            }
        }
        ids
    }

    /// Returns discovery URL snapshot in deterministic manager-maintained order.
    #[must_use]
    pub fn discovery_tso_urls(&self) -> Vec<String> {
        self.discovery_tso_urls.clone()
    }

    /// Returns whether one address maps to local member identity aliases.
    #[must_use]
    pub fn is_local_member_address(&self, address: &str) -> bool {
        expected_primary_matches_local_identity(address, self.member_name_prefix.as_str())
    }

    /// Returns discovery URLs that map to the local node identity.
    #[must_use]
    pub fn local_discovery_tso_urls(&self) -> Vec<String> {
        self.discovery_tso_urls
            .iter()
            .filter(|url| self.is_local_member_address(url))
            .cloned()
            .collect()
    }

    /// Returns groups that currently satisfy local primary + lease serving checks.
    #[must_use]
    pub fn serving_group_ids(&self, now_ms: i64) -> Vec<u32> {
        let mut ids = Vec::new();
        for (group_id, entry) in self.entries.iter().enumerate() {
            if let Some(state) = entry {
                let runtime = lock_runtime(&state.runtime);
                if runtime.active && runtime.allocator.is_serving(now_ms) {
                    ids.push(group_id as u32);
                }
            }
        }
        ids
    }

    /// Returns group ids marked for deferred timestamp-key cleanup.
    #[must_use]
    pub fn deleted_groups(&self) -> Vec<u32> {
        self.deleted_groups.list()
    }

    /// Prepares routing/runtime context for one TSO request.
    ///
    /// Normal requests can use the returned runtime handle on a per-group lock
    /// without holding manager lock for timestamp generation.
    pub(crate) fn prepare_tso_request(
        &self,
        keyspace_id: u32,
        requested_group_id: u32,
    ) -> Result<PreparedTsoRequest, ServiceError> {
        validate_serving_group_id(requested_group_id)?;
        let resolved_group_id = self.resolve_group_id(keyspace_id, requested_group_id)?;
        let state = self.state(resolved_group_id)?;
        if state.meta.merge_state.is_some() {
            return Err(ServiceError::KeyspaceGroupMerging {
                group_id: resolved_group_id,
            });
        }

        Ok(PreparedTsoRequest {
            resolved_group_id,
            runtime: Arc::clone(&state.runtime),
            in_flight_requests: Arc::clone(&state.in_flight_requests),
            request_count_total: Arc::clone(&state.request_count_total),
            last_request_ms: Arc::clone(&state.last_request_ms),
            // Only split targets require split timeline checks on request path.
            requires_slow_path: is_split_target(&state.meta),
        })
    }

    /// Registers a keyspace-group slot and binds all owned keyspaces to it.
    pub fn register_group(&mut self, mut meta: KeyspaceGroup) -> Result<(), ServiceError> {
        validate_serving_group_id(meta.id)?;
        let idx = meta.id as usize;
        let member_value = format!("{}:{}", self.member_name_prefix, meta.id);
        let allocator = GroupAllocator::new(
            self.cluster_id,
            meta.id,
            member_value,
            Arc::clone(&self.store),
            self.allocator_config,
        );

        // Preserve deterministic keyspace ordering in metadata.
        meta.keyspaces.sort_unstable();
        meta.keyspaces.dedup();

        self.routing.bind_group(meta.id, &meta.keyspaces);
        self.entries[idx] = Some(GroupState::new(meta, allocator));
        self.refresh_discovery_tso_urls();
        self.bump_route_epoch();
        Ok(())
    }

    /// Upserts group metadata and allocator slot.
    pub fn upsert_group(&mut self, mut meta: KeyspaceGroup) -> Result<(), ServiceError> {
        validate_serving_group_id(meta.id)?;
        let idx = meta.id as usize;

        // Keep deterministic keyspace ordering.
        meta.keyspaces.sort_unstable();
        meta.keyspaces.dedup();

        // Remove stale lookup bindings if the slot already exists.
        if let Some(existing) = self.entries[idx].as_ref() {
            self.routing
                .unbind_group_if_owned(meta.id, &existing.meta.keyspaces);
        }

        if let Some(existing) = self.entries[idx].as_mut() {
            existing.meta = meta.clone();
            existing.keyspace_set = meta.keyspaces.iter().copied().collect();
        } else {
            let member_value = format!("{}:{}", self.member_name_prefix, meta.id);
            let allocator = GroupAllocator::new(
                self.cluster_id,
                meta.id,
                member_value,
                Arc::clone(&self.store),
                self.allocator_config,
            );
            self.entries[idx] = Some(GroupState::new(meta.clone(), allocator));
        }

        self.routing.bind_group(meta.id, &meta.keyspaces);
        self.refresh_discovery_tso_urls();
        self.bump_route_epoch();
        Ok(())
    }

    /// Removes a group slot and associated keyspace routing entries.
    pub fn remove_group(&mut self, group_id: u32) -> Result<(), ServiceError> {
        validate_serving_group_id(group_id)?;
        let idx = group_id as usize;
        let Some(state) = self.entries[idx].take() else {
            return Ok(());
        };

        self.routing
            .unbind_group_if_owned(group_id, &state.meta.keyspaces);
        {
            let mut runtime = lock_runtime(&state.runtime);
            runtime.active = false;
            runtime.lease_id = 0;
            runtime.initialized = false;
            runtime.allocator.step_down();
            clear_prefetch_diagnostics(&mut runtime);
        }
        state.in_flight_requests.store(0, AtomicOrdering::Release);
        state.request_count_total.store(0, AtomicOrdering::Release);
        state.last_request_ms.store(0, AtomicOrdering::Release);
        self.deleted_groups.mark(group_id);
        self.refresh_discovery_tso_urls();
        self.bump_route_epoch();
        Ok(())
    }

    /// Campaigns a group primary through local in-memory election state.
    /// This is primarily used by in-memory test mode; production etcd mode
    /// should use external lease-bound CAS election and `campaign_group_primary_external`.
    pub fn campaign_group_primary(
        &mut self,
        group_id: u32,
        now_ms: i64,
    ) -> Result<(), ServiceError> {
        let state = self.state_mut(group_id)?;
        let mut runtime = lock_runtime(&state.runtime);
        if runtime.allocator.campaign_primary(now_ms)? {
            runtime.allocator.initialize(now_ms)?;
            runtime.initialized = true;
        }
        runtime.active = true;
        Ok(())
    }

    /// Campaigns local allocator state after successful external election CAS.
    pub fn campaign_group_primary_external(
        &mut self,
        group_id: u32,
        now_ms: i64,
    ) -> Result<(), ServiceError> {
        let state = self.state_mut(group_id)?;
        let mut runtime = lock_runtime(&state.runtime);
        if runtime.allocator.campaign_primary_external(now_ms)? {
            runtime.allocator.initialize(now_ms)?;
            runtime.initialized = true;
        }
        runtime.active = true;
        Ok(())
    }

    /// Renews one group's lease-serving state after a keepalive succeeds.
    pub fn keep_group_primary_alive(
        &mut self,
        group_id: u32,
        now_ms: i64,
    ) -> Result<(), ServiceError> {
        let state = self.state_mut(group_id)?;
        let mut runtime = lock_runtime(&state.runtime);
        runtime.allocator.keep_alive(now_ms);
        Ok(())
    }

    /// Performs background timestamp-window update for serving groups.
    ///
    /// This is used by the async-prefetch extension path to reduce request-path
    /// persistence/update pressure by warming save windows ahead of overflow.
    pub fn prefetch_serving_group_windows(&self, now_ms: i64) -> PrefetchRunStats {
        let mut stats = PrefetchRunStats::default();
        for (group_id, entry) in self.entries.iter().enumerate() {
            let Some(state) = entry.as_ref() else {
                continue;
            };
            let mut runtime = lock_runtime(&state.runtime);
            if !runtime.active || !runtime.allocator.is_serving(now_ms) {
                continue;
            }
            stats.attempted_groups = stats.attempted_groups.saturating_add(1);
            runtime.prefetch_state = PrefetchDiagnosticState::InFlight;
            runtime.prefetch_last_attempt_ms = now_ms;
            if let Err(err) = ensure_runtime_initialized(&mut runtime, now_ms, group_id as u32) {
                runtime.prefetch_state = PrefetchDiagnosticState::Backoff;
                runtime.prefetch_last_error = err.to_string();
                stats.stalled_groups = stats.stalled_groups.saturating_add(1);
                stats.error_groups = stats.error_groups.saturating_add(1);
                continue;
            }
            if let Err(err) = runtime.allocator.update_tso(now_ms) {
                // Prefetch is a best-effort extension path; request path remains authoritative.
                tracing::debug!(
                    group_id,
                    error = %err,
                    "prefetch update failed for serving group"
                );
                runtime.prefetch_state = PrefetchDiagnosticState::Backoff;
                runtime.prefetch_last_error = err.to_string();
                stats.stalled_groups = stats.stalled_groups.saturating_add(1);
                stats.error_groups = stats.error_groups.saturating_add(1);
                continue;
            }
            runtime.prefetch_state = PrefetchDiagnosticState::Idle;
            runtime.prefetch_last_error.clear();
            stats.completed_groups = stats.completed_groups.saturating_add(1);
        }
        stats
    }

    /// Updates observed lease id for one group.
    pub fn set_group_lease_id(
        &mut self,
        group_id: u32,
        lease_id: Option<i64>,
    ) -> Result<(), ServiceError> {
        let state = self.state_mut(group_id)?;
        let mut runtime = lock_runtime(&state.runtime);
        runtime.lease_id = lease_id.unwrap_or_default();
        Ok(())
    }

    /// Updates expected-primary metadata for one group.
    pub fn set_group_expected_primary(
        &mut self,
        group_id: u32,
        expected: Option<String>,
    ) -> Result<(), ServiceError> {
        let state = self.state_mut(group_id)?;
        let mut runtime = lock_runtime(&state.runtime);
        runtime.allocator.set_expected_primary(expected);
        Ok(())
    }

    /// Applies externally observed expected-primary and steps down if required.
    pub fn apply_group_expected_primary(
        &mut self,
        group_id: u32,
        expected: Option<String>,
        now_ms: i64,
    ) -> Result<(), ServiceError> {
        let state = self.state_mut(group_id)?;
        let mut runtime = lock_runtime(&state.runtime);
        runtime.allocator.set_expected_primary(expected.clone());
        if let Some(expected_member) = expected
            && !expected_primary_matches_local_identity(
                &expected_member,
                runtime.allocator.member_value(),
            )
            && runtime.allocator.is_serving(now_ms)
        {
            runtime.allocator.step_down();
            runtime.lease_id = 0;
            clear_prefetch_diagnostics(&mut runtime);
        }
        Ok(())
    }

    /// Steps down one group's local primary state.
    pub fn step_down_group_primary(&mut self, group_id: u32) -> Result<(), ServiceError> {
        let state = self.state_mut(group_id)?;
        let mut runtime = lock_runtime(&state.runtime);
        runtime.allocator.step_down();
        runtime.lease_id = 0;
        clear_prefetch_diagnostics(&mut runtime);
        Ok(())
    }

    /// Steps down all locally served primaries.
    pub fn step_down_all_primaries(&mut self) {
        for entry in &mut self.entries {
            if let Some(state) = entry.as_mut() {
                let mut runtime = lock_runtime(&state.runtime);
                runtime.allocator.step_down();
                runtime.lease_id = 0;
                state.in_flight_requests.store(0, AtomicOrdering::Release);
                clear_prefetch_diagnostics(&mut runtime);
            }
        }
    }

    /// Returns total in-flight mailbox depth across active groups.
    #[must_use]
    pub fn total_in_flight_requests(&self) -> u32 {
        let mut total = 0_u32;
        for entry in &self.entries {
            if let Some(state) = entry.as_ref() {
                total = total.saturating_add(load_in_flight_requests(&state.in_flight_requests));
            }
        }
        total
    }

    /// Returns cumulative request counter and stale-sample group count for capacity admission metrics.
    ///
    /// A group sample is considered stale when no request has been observed in
    /// `stale_after_ms`, or when the group has never served a request.
    #[must_use]
    pub fn assigned_group_qps_inputs(&self, now_ms: i64, stale_after_ms: i64) -> (u64, u32) {
        let mut total_requests = 0_u64;
        let mut stale_groups = 0_u32;
        for entry in &self.entries {
            let Some(state) = entry.as_ref() else {
                continue;
            };
            total_requests = total_requests
                .saturating_add(state.request_count_total.load(AtomicOrdering::Acquire));
            let last_request_ms = state.last_request_ms.load(AtomicOrdering::Acquire);
            if last_request_ms == 0
                || now_ms.saturating_sub(last_request_ms).max(0) > stale_after_ms.max(0)
            {
                stale_groups = stale_groups.saturating_add(1);
            }
        }
        (total_requests, stale_groups)
    }

    /// Returns a detailed debug snapshot for one keyspace group.
    pub fn group_debug_state(
        &self,
        group_id: u32,
        now_ms: i64,
    ) -> Result<GroupDebugState, ServiceError> {
        let state = self.state(group_id)?;
        let (
            is_primary,
            in_memory_timestamp,
            lease_id,
            lease_expire_at_ms,
            save_window_end_ms,
            prefetch_state,
            prefetch_last_attempt_ms,
            prefetch_last_error,
        ) = {
            let runtime = lock_runtime(&state.runtime);
            (
                runtime.allocator.is_serving(now_ms),
                runtime.allocator.current_timestamp(),
                runtime.lease_id,
                runtime.allocator.lease_expire_at_ms(),
                runtime.allocator.save_window_end_ms(),
                runtime.prefetch_state.as_str().to_string(),
                runtime.prefetch_last_attempt_ms,
                runtime.prefetch_last_error.clone(),
            )
        };
        let persisted_nanos = self
            .load_persisted_timestamp_nanos(group_id)?
            .unwrap_or_default();
        let in_flight_requests = load_in_flight_requests(&state.in_flight_requests);
        Ok(GroupDebugState {
            is_primary,
            in_memory_timestamp,
            lease_id,
            lease_expire_at_ms,
            save_window_end_ms,
            last_persisted_physical_ms: persisted_nanos.saturating_div(1_000_000),
            last_persisted_logical: 0,
            mailbox_depth: in_flight_requests,
            in_flight_requests,
            prefetch_state,
            prefetch_last_attempt_ms,
            prefetch_last_error,
        })
    }

    /// Returns an in-memory runtime snapshot without persisted-store reads.
    pub fn group_runtime_snapshot(
        &self,
        group_id: u32,
        now_ms: i64,
    ) -> Result<GroupRuntimeSnapshot, ServiceError> {
        let state = self.state(group_id)?;
        let runtime = lock_runtime(&state.runtime);
        Ok(GroupRuntimeSnapshot {
            is_primary: runtime.allocator.is_serving(now_ms),
            in_memory_timestamp: runtime.allocator.current_timestamp(),
        })
    }

    /// Returns election key/value metadata for runtime etcd HA loops.
    pub fn group_election_meta(&self, group_id: u32) -> Result<GroupElectionMeta, ServiceError> {
        let state = self.state(group_id)?;
        let runtime = lock_runtime(&state.runtime);
        Ok(GroupElectionMeta {
            group_id,
            leader_key: runtime.allocator.leader_key().to_string(),
            expected_primary_key: runtime.allocator.expected_primary_key().to_string(),
            member_value: runtime.allocator.member_value().to_string(),
        })
    }

    /// Returns actor ids that must acknowledge config epoch apply in the current snapshot.
    #[must_use]
    pub fn required_config_apply_actors(&self) -> Vec<u32> {
        self.active_group_ids()
    }

    /// Applies candidate config epoch only when all required actors are present/active.
    ///
    /// This models strict commit-or-abort behavior: if any required actor cannot
    /// acknowledge, the operation aborts without mutating any actor epoch.
    pub fn apply_config_epoch(
        &mut self,
        candidate_epoch: u64,
        required_actors: &[u32],
    ) -> Result<(), ServiceError> {
        self.apply_config_epoch_with_allocator(candidate_epoch, required_actors, None)
    }

    /// Applies candidate config epoch and optional allocator runtime config atomically.
    pub fn apply_config_epoch_with_allocator(
        &mut self,
        candidate_epoch: u64,
        required_actors: &[u32],
        allocator_config: Option<AllocatorConfig>,
    ) -> Result<(), ServiceError> {
        for group_id in required_actors {
            let state = self.state(*group_id)?;
            let runtime = lock_runtime(&state.runtime);
            if !runtime.active {
                return Err(ServiceError::Storage(format!(
                    "config apply actor {group_id} is inactive"
                )));
            }
        }

        for group_id in required_actors {
            let state = self.state_mut(*group_id)?;
            let mut runtime = lock_runtime(&state.runtime);
            runtime.config_epoch = candidate_epoch;
            if let Some(config) = allocator_config {
                runtime.allocator.apply_config(config);
            }
        }
        if let Some(config) = allocator_config {
            self.allocator_config = config;
        }
        Ok(())
    }

    /// Returns current applied config epoch for one group runtime actor.
    pub fn group_config_epoch(&self, group_id: u32) -> Result<u64, ServiceError> {
        let state = self.state(group_id)?;
        let runtime = lock_runtime(&state.runtime);
        Ok(runtime.config_epoch)
    }

    /// Handles one `Tso` request and returns `(timestamp, resolved_group_id)`.
    pub fn handle_tso_request(
        &mut self,
        now_ms: i64,
        keyspace_id: u32,
        requested_group_id: u32,
        count: u32,
        overload_hardening: bool,
        overload_max_in_flight_per_group: u32,
    ) -> Result<(Timestamp, u32), ServiceError> {
        validate_serving_group_id(requested_group_id)?;

        let resolved_group_id = self.resolve_group_id(keyspace_id, requested_group_id)?;
        self.check_tso_split(now_ms, resolved_group_id)?;
        self.check_group_merge(resolved_group_id)?;

        let ts = self.generate_group_tso(
            now_ms,
            resolved_group_id,
            count,
            overload_hardening,
            overload_max_in_flight_per_group,
        )?;
        Ok((ts, resolved_group_id))
    }

    /// Finds the keyspace group by keyspace id with current in-memory mapping.
    pub fn find_group_by_keyspace_id(
        &self,
        _now_ms: i64,
        keyspace_id: u32,
        requested_group_id: u32,
    ) -> Result<GroupLookupResult, ServiceError> {
        validate_api_group_id(requested_group_id)?;
        let resolved_group_id =
            self.resolve_group_id_for_api_lookup(keyspace_id, requested_group_id)?;
        let state = self.state(resolved_group_id)?;

        let mut group = state.meta.clone();
        let runtime = lock_runtime(&state.runtime);
        let has_elected_primary = runtime.allocator.is_primary_elected();
        let local_member_address = self.member_name_prefix.as_str();
        for member in &mut group.members {
            member.is_primary = has_elected_primary
                && expected_primary_matches_local_identity(&member.address, local_member_address);
        }

        Ok(GroupLookupResult {
            group,
            resolved_group_id,
            mod_revision: self.mod_revision,
        })
    }

    /// Samples minimum timestamp across local-serving groups.
    pub fn get_min_ts(&mut self, now_ms: i64) -> Result<MinTsSummary, ServiceError> {
        let mut min_ts: Option<Timestamp> = None;
        let mut serving = 0_u32;
        let mut total = 0_u32;

        for (group_id, entry) in self.entries.iter().enumerate() {
            let Some(state) = entry.as_ref() else {
                continue;
            };
            total = total.saturating_add(1);
            let mut runtime = lock_runtime(&state.runtime);

            if !runtime.active || !runtime.allocator.is_primary_elected() {
                return Err(ServiceError::Storage(format!(
                    "keyspace group {group_id}'s primary is not elected"
                )));
            }
            if !runtime.allocator.is_serving(now_ms) {
                continue;
            }
            serving = serving.saturating_add(1);

            if is_split_target(&state.meta) {
                continue;
            }

            ensure_runtime_initialized(&mut runtime, now_ms, group_id as u32)?;
            let _in_flight_guard = acquire_in_flight_request(
                &state.in_flight_requests,
                group_id as u32,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )?;
            let sampled = runtime.allocator.generate_tso(now_ms, 1)?;
            if min_ts
                .as_ref()
                .is_none_or(|current| compare_timestamp(sampled, *current) == Ordering::Less)
            {
                min_ts = Some(sampled);
            }
        }

        Ok(MinTsSummary {
            timestamp: min_ts.unwrap_or_default(),
            keyspace_groups_serving: serving,
            keyspace_groups_total: total,
        })
    }

    /// Creates a split-target group and marks source/target in split state.
    pub fn split_keyspace_group(
        &mut self,
        split_source_id: u32,
        split_target_id: u32,
        keyspaces: Vec<u32>,
    ) -> Result<(), ServiceError> {
        validate_serving_group_id(split_source_id)?;
        validate_serving_group_id(split_target_id)?;
        if split_source_id == split_target_id {
            return Err(ServiceError::Storage(
                "split source and target must differ".to_string(),
            ));
        }
        if keyspaces.is_empty() {
            return Err(ServiceError::Storage(
                "split target keyspace list must be non-empty".to_string(),
            ));
        }
        if self.state(split_target_id).is_ok() {
            return Err(ServiceError::Storage(
                "split target already exists".to_string(),
            ));
        }

        let source_snapshot = self.state(split_source_id)?.meta.clone();
        if source_snapshot.merge_state.is_some() {
            return Err(ServiceError::Storage(format!(
                "keyspace group {} is merging",
                split_source_id
            )));
        }
        if source_snapshot.split_state.is_some() {
            return Err(ServiceError::Storage(format!(
                "keyspace group {} is already splitting",
                split_source_id
            )));
        }

        let requested: HashSet<u32> = keyspaces.iter().copied().collect();
        if !requested
            .iter()
            .all(|ks| source_snapshot.keyspaces.contains(ks))
        {
            return Err(ServiceError::Storage(
                "requested split keyspace is not in source group".to_string(),
            ));
        }

        let remaining: Vec<u32> = source_snapshot
            .keyspaces
            .iter()
            .copied()
            .filter(|ks| !requested.contains(ks))
            .collect();
        if remaining.is_empty() {
            return Err(ServiceError::Storage(
                "split source cannot become empty".to_string(),
            ));
        }

        {
            let source = self.state_mut(split_source_id)?;
            source.meta.keyspaces = remaining.clone();
            source.meta.split_state = Some(SplitState {
                split_source: split_source_id,
            });
            source.keyspace_set = remaining.into_iter().collect();
        }

        let mut target_meta = KeyspaceGroup {
            id: split_target_id,
            user_kind: source_snapshot.user_kind,
            split_state: Some(SplitState {
                split_source: split_source_id,
            }),
            merge_state: None,
            members: source_snapshot.members,
            keyspaces: keyspaces.to_vec(),
        };
        target_meta.keyspaces.sort_unstable();

        let member_value = format!("{}:{}", self.member_name_prefix, split_target_id);
        let target_allocator = GroupAllocator::new(
            self.cluster_id,
            split_target_id,
            member_value,
            Arc::clone(&self.store),
            self.allocator_config,
        );

        self.entries[split_target_id as usize] =
            Some(GroupState::new(target_meta, target_allocator));

        self.routing.bind_group(split_target_id, &keyspaces);
        self.refresh_discovery_tso_urls();
        self.bump_route_epoch();
        Ok(())
    }

    /// Clears split states after split safety checks pass.
    pub fn finish_split_keyspace_group(
        &mut self,
        split_target_id: u32,
    ) -> Result<(), ServiceError> {
        let split_source_id = {
            let target = self.state(split_target_id)?;
            let Some(split_state) = target.meta.split_state.as_ref() else {
                return Err(ServiceError::Storage(
                    "split target is not in split state".to_string(),
                ));
            };
            split_state.split_source
        };

        self.state_mut(split_target_id)?.meta.split_state = None;
        if let Ok(source) = self.state_mut(split_source_id) {
            source.meta.split_state = None;
        }
        self.bump_route_epoch();
        Ok(())
    }

    /// Marks merge target and removes merge sources from active groups.
    pub fn merge_keyspace_groups(
        &mut self,
        merge_target_id: u32,
        merge_list: Vec<u32>,
    ) -> Result<(), ServiceError> {
        validate_serving_group_id(merge_target_id)?;
        if merge_list.is_empty() {
            return Ok(());
        }
        if merge_list.contains(&merge_target_id) {
            return Err(ServiceError::Storage(
                "merge list must not include target".to_string(),
            ));
        }

        let mut merged_keyspaces: HashSet<u32> = self
            .state(merge_target_id)?
            .meta
            .keyspaces
            .iter()
            .copied()
            .collect();

        for source_id in &merge_list {
            validate_serving_group_id(*source_id)?;
            let source = self.state(*source_id)?;
            if source.meta.split_state.is_some() {
                return Err(ServiceError::Storage(format!(
                    "keyspace group {} is splitting",
                    source_id
                )));
            }
            if source.meta.merge_state.is_some() {
                return Err(ServiceError::Storage(format!(
                    "keyspace group {} is merging",
                    source_id
                )));
            }
            merged_keyspaces.extend(source.meta.keyspaces.iter().copied());
        }

        let mut merged = merged_keyspaces.into_iter().collect::<Vec<_>>();
        merged.sort_unstable();

        {
            let target = self.state_mut(merge_target_id)?;
            target.meta.keyspaces = merged.clone();
            target.meta.merge_state = Some(MergeState {
                merge_list: merge_list.clone(),
            });
            target.keyspace_set = merged.iter().copied().collect();
        }

        for source_id in &merge_list {
            let Some(source_state) = self.entries[*source_id as usize].take() else {
                continue;
            };
            {
                let mut runtime = lock_runtime(&source_state.runtime);
                runtime.active = false;
                runtime.lease_id = 0;
                runtime.initialized = false;
                runtime.allocator.step_down();
                clear_prefetch_diagnostics(&mut runtime);
            }
            source_state
                .in_flight_requests
                .store(0, AtomicOrdering::Release);
            source_state
                .request_count_total
                .store(0, AtomicOrdering::Release);
            source_state
                .last_request_ms
                .store(0, AtomicOrdering::Release);
            self.deleted_groups.mark(*source_id);
            for keyspace in source_state.meta.keyspaces {
                self.routing.rebind_keyspace(keyspace, merge_target_id);
            }
        }

        for keyspace in merged {
            self.routing.rebind_keyspace(keyspace, merge_target_id);
        }
        self.refresh_discovery_tso_urls();
        self.bump_route_epoch();
        Ok(())
    }

    /// Finalizes merge state after all source groups are gone.
    pub fn finish_merge_keyspace_group(
        &mut self,
        merge_target_id: u32,
        now_ms: i64,
    ) -> Result<(), ServiceError> {
        let merge_list = {
            let target = self.state(merge_target_id)?;
            let Some(merge_state) = target.meta.merge_state.as_ref() else {
                return Err(ServiceError::Storage(
                    "merge target is not in merging state".to_string(),
                ));
            };
            merge_state.merge_list.clone()
        };

        for source_id in &merge_list {
            if self.state(*source_id).is_ok() {
                return Err(ServiceError::Storage(format!(
                    "merge source {source_id} still exists"
                )));
            }
        }

        // Merge finalize must keep the target timeline above all removed sources.
        let target_nanos = self
            .load_persisted_timestamp_nanos(merge_target_id)?
            .unwrap_or_default();
        let mut max_source_nanos = target_nanos;
        for source_id in &merge_list {
            if let Some(source_nanos) = self.load_persisted_timestamp_nanos(*source_id)?
                && source_nanos > max_source_nanos
            {
                max_source_nanos = source_nanos;
            }
        }

        if max_source_nanos >= target_nanos {
            let target = self.state(merge_target_id)?;
            let mut runtime = lock_runtime(&target.runtime);
            ensure_runtime_initialized(&mut runtime, now_ms, merge_target_id)?;
            let rebased_physical_ms = max_source_nanos.saturating_add(1_000_000) / 1_000_000;
            let rebased = Timestamp {
                physical: rebased_physical_ms,
                logical: 0,
                suffix_bits: 0,
            };
            runtime
                .allocator
                .set_tso(now_ms, rebased.compose(), true, true)?;
        }

        self.state_mut(merge_target_id)?.meta.merge_state = None;
        self.bump_route_epoch();
        Ok(())
    }

    /// Cleans persisted timestamp keys for merged/deleted groups.
    pub fn cleanup_deleted_group_timestamps(&mut self) -> Result<(), ServiceError> {
        let groups = self.deleted_groups.snapshot();
        for group_id in groups {
            if self.state(group_id).is_ok() {
                self.deleted_groups.remove(group_id);
                continue;
            }
            if self.is_pending_merge_source(group_id) {
                continue;
            }
            self.store
                .delete_timestamp(self.cluster_id, group_id)
                .map_err(|err| ServiceError::Storage(err.to_string()))?;
            self.deleted_groups.remove(group_id);
        }
        Ok(())
    }

    fn check_tso_split(&mut self, now_ms: i64, target_group_id: u32) -> Result<(), ServiceError> {
        let source_group_id = match self
            .state(target_group_id)?
            .meta
            .split_state
            .as_ref()
            .map(|state| state.split_source)
        {
            Some(id) => id,
            None => return Ok(()),
        };
        if source_group_id == target_group_id {
            // Split source requests remain on normal serving path; split timeline
            // safety check is only required when target serves split keyspaces.
            return Ok(());
        }

        let (target, source) = self.two_states_mut(target_group_id, source_group_id)?;
        let mut target_runtime = lock_runtime(&target.runtime);
        let mut source_runtime = lock_runtime(&source.runtime);
        ensure_runtime_initialized(&mut target_runtime, now_ms, target_group_id)?;
        ensure_runtime_initialized(&mut source_runtime, now_ms, source_group_id)?;
        let target_ts = target_runtime.allocator.generate_tso(now_ms, 1)?;
        let source_ts = source_runtime.allocator.generate_tso(now_ms, 1)?;

        if compare_timestamp(source_ts, target_ts) == Ordering::Greater {
            let rebased = Timestamp {
                physical: source_ts.physical + 1,
                logical: source_ts.logical,
                suffix_bits: 0,
            };
            target_runtime
                .allocator
                .set_tso(now_ms, rebased.compose(), true, true)?;
        }
        drop(source_runtime);
        drop(target_runtime);

        // Both success cases finalize split state.
        self.finish_split_keyspace_group(target_group_id)
    }

    fn check_group_merge(&self, group_id: u32) -> Result<(), ServiceError> {
        if self.state(group_id)?.meta.merge_state.is_some() {
            return Err(ServiceError::KeyspaceGroupMerging { group_id });
        }
        Ok(())
    }

    fn is_pending_merge_source(&self, source_group_id: u32) -> bool {
        for entry in &self.entries {
            let Some(state) = entry.as_ref() else {
                continue;
            };
            let Some(merge_state) = state.meta.merge_state.as_ref() else {
                continue;
            };
            if merge_state.merge_list.contains(&source_group_id) {
                return true;
            }
        }
        false
    }

    fn refresh_discovery_tso_urls(&mut self) {
        let mut observed = Vec::new();
        let mut seen = HashSet::new();
        for entry in &self.entries {
            let Some(state) = entry.as_ref() else {
                continue;
            };
            for member in &state.meta.members {
                let address = member.address.trim();
                if address.is_empty() {
                    continue;
                }
                if seen.insert(address.to_string()) {
                    observed.push(address.to_string());
                }
            }
        }

        if self.discovery_tso_urls.is_empty() {
            self.discovery_tso_urls = observed;
            return;
        }

        // Keep existing emission order for stable addresses, then append new additions.
        let observed_lookup: HashSet<&str> = observed.iter().map(String::as_str).collect();
        let mut next = Vec::new();
        let mut next_seen = HashSet::new();
        for existing in &self.discovery_tso_urls {
            if observed_lookup.contains(existing.as_str()) && next_seen.insert(existing.clone()) {
                next.push(existing.clone());
            }
        }
        for address in observed {
            if next_seen.insert(address.clone()) {
                next.push(address);
            }
        }
        self.discovery_tso_urls = next;
    }

    fn bump_route_epoch(&self) {
        let _ = self.route_epoch.fetch_add(1, AtomicOrdering::AcqRel);
    }

    fn resolve_group_id(
        &self,
        keyspace_id: u32,
        requested_group_id: u32,
    ) -> Result<u32, ServiceError> {
        validate_serving_group_id(requested_group_id)?;
        if let Ok(state) = self.state(requested_group_id)
            && state.keyspace_set.contains(&keyspace_id)
        {
            return Ok(requested_group_id);
        }

        if let Some(group_id) = self.routing.resolve(keyspace_id) {
            if self.state(group_id).is_ok() {
                return Ok(group_id);
            }
            return Err(ServiceError::NotServed {
                group_id: requested_group_id,
            });
        }

        if requested_group_id != DEFAULT_KEYSPACE_GROUP_ID {
            return Err(ServiceError::KeyspaceNotAssigned { keyspace_id });
        }

        if self.state(DEFAULT_KEYSPACE_GROUP_ID).is_ok() {
            return Ok(DEFAULT_KEYSPACE_GROUP_ID);
        }

        Err(ServiceError::KeyspaceNotAssigned { keyspace_id })
    }

    fn resolve_group_id_for_api_lookup(
        &self,
        keyspace_id: u32,
        requested_group_id: u32,
    ) -> Result<u32, ServiceError> {
        validate_api_group_id(requested_group_id)?;

        if requested_group_id < MAX_KEYSPACE_GROUP_COUNT_IN_USE
            && let Ok(state) = self.state(requested_group_id)
            && state.keyspace_set.contains(&keyspace_id)
        {
            return Ok(requested_group_id);
        }

        if let Some(group_id) = self.routing.resolve(keyspace_id) {
            if self.state(group_id).is_ok() {
                return Ok(group_id);
            }
            return Err(ServiceError::NotServed {
                group_id: requested_group_id,
            });
        }

        if requested_group_id != DEFAULT_KEYSPACE_GROUP_ID {
            return Err(ServiceError::KeyspaceNotAssigned { keyspace_id });
        }

        if self.state(DEFAULT_KEYSPACE_GROUP_ID).is_ok() {
            return Ok(DEFAULT_KEYSPACE_GROUP_ID);
        }

        Err(ServiceError::KeyspaceNotAssigned { keyspace_id })
    }

    fn generate_group_tso(
        &self,
        now_ms: i64,
        group_id: u32,
        count: u32,
        overload_hardening: bool,
        overload_max_in_flight_per_group: u32,
    ) -> Result<Timestamp, ServiceError> {
        let state = self.state(group_id)?;
        let _in_flight_guard = acquire_in_flight_request(
            &state.in_flight_requests,
            group_id,
            overload_hardening,
            overload_max_in_flight_per_group,
        )?;
        let mut runtime = lock_runtime(&state.runtime);
        if !runtime.active {
            return Err(ServiceError::NotServed { group_id });
        }
        ensure_runtime_initialized(&mut runtime, now_ms, group_id)?;
        let timestamp = runtime.allocator.generate_tso(now_ms, count)?;
        drop(runtime);
        record_group_request_sample(&state.request_count_total, &state.last_request_ms, now_ms);
        Ok(timestamp)
    }

    fn state(&self, group_id: u32) -> Result<&GroupState, ServiceError> {
        validate_serving_group_id(group_id)?;
        self.entries[group_id as usize]
            .as_ref()
            .ok_or(ServiceError::NotServed { group_id })
    }

    fn state_mut(&mut self, group_id: u32) -> Result<&mut GroupState, ServiceError> {
        validate_serving_group_id(group_id)?;
        self.entries[group_id as usize]
            .as_mut()
            .ok_or(ServiceError::NotServed { group_id })
    }

    fn load_persisted_timestamp_nanos(&self, group_id: u32) -> Result<Option<i64>, ServiceError> {
        self.store
            .load_timestamp(self.cluster_id, group_id)
            .map_err(|err| ServiceError::Storage(err.to_string()))
    }

    fn two_states_mut(
        &mut self,
        group_a: u32,
        group_b: u32,
    ) -> Result<(&mut GroupState, &mut GroupState), ServiceError> {
        validate_serving_group_id(group_a)?;
        validate_serving_group_id(group_b)?;
        if group_a == group_b {
            return Err(ServiceError::Storage(
                "source and target group must differ".to_string(),
            ));
        }

        let idx_a = group_a as usize;
        let idx_b = group_b as usize;
        let (low_idx, high_idx, swapped) = if idx_a < idx_b {
            (idx_a, idx_b, false)
        } else {
            (idx_b, idx_a, true)
        };

        let (left, right) = self.entries.split_at_mut(high_idx);
        let low = left[low_idx].as_mut().ok_or(ServiceError::NotServed {
            group_id: low_idx as u32,
        })?;
        let high = right[0].as_mut().ok_or(ServiceError::NotServed {
            group_id: high_idx as u32,
        })?;

        if swapped {
            Ok((high, low))
        } else {
            Ok((low, high))
        }
    }
}

/// Returns whether an expected-primary identity string refers to the local node.
pub(crate) fn expected_primary_matches_local_identity(
    expected_member: &str,
    member_value: &str,
) -> bool {
    let normalized_expected = normalize_identity(expected_member);
    local_identity_aliases(member_value)
        .into_iter()
        .any(|alias| {
            alias == expected_member
                || alias == normalized_expected
                || normalize_identity(&alias) == normalized_expected
        })
}

fn local_identity_aliases(member_value: &str) -> Vec<String> {
    let mut aliases = vec![member_value.to_string()];
    for key in ["PD_TSO_ADVERTISE_LISTEN_ADDR", "PD_TSO_LISTEN_ADDR"] {
        let Ok(raw) = std::env::var(key) else {
            continue;
        };
        for part in raw.split(',') {
            let value = part.trim();
            if !value.is_empty() {
                aliases.push(value.to_string());
            }
        }
    }
    aliases
}

fn normalize_identity(value: &str) -> String {
    let trimmed = value.trim().trim_end_matches('/');
    if let Some(raw) = trimmed.strip_prefix("http://") {
        return raw.to_string();
    }
    if let Some(raw) = trimmed.strip_prefix("https://") {
        return raw.to_string();
    }
    trimmed.to_string()
}

fn is_split_target(group: &KeyspaceGroup) -> bool {
    group
        .split_state
        .as_ref()
        .is_some_and(|split| split.split_source != group.id)
}

fn compare_timestamp(left: Timestamp, right: Timestamp) -> Ordering {
    (left.physical, left.logical).cmp(&(right.physical, right.logical))
}

fn lock_runtime(
    runtime: &Arc<Mutex<GroupRuntimeState>>,
) -> std::sync::MutexGuard<'_, GroupRuntimeState> {
    runtime.lock().unwrap_or_else(|poison| poison.into_inner())
}

fn ensure_runtime_initialized(
    runtime: &mut GroupRuntimeState,
    now_ms: i64,
    group_id: u32,
) -> Result<(), ServiceError> {
    if !runtime.active {
        return Err(ServiceError::NotServed { group_id });
    }
    if !runtime.initialized {
        runtime.allocator.initialize(now_ms)?;
        runtime.initialized = true;
    }
    Ok(())
}

fn clear_prefetch_diagnostics(runtime: &mut GroupRuntimeState) {
    runtime.prefetch_state = PrefetchDiagnosticState::Idle;
    runtime.prefetch_last_error.clear();
}

fn record_group_request_sample(
    request_count_total: &Arc<AtomicU64>,
    last_request_ms: &Arc<AtomicI64>,
    now_ms: i64,
) {
    request_count_total.fetch_add(1, AtomicOrdering::AcqRel);
    last_request_ms.store(now_ms, AtomicOrdering::Release);
}

struct InFlightRequestGuard {
    in_flight_requests: Arc<AtomicU32>,
}

impl Drop for InFlightRequestGuard {
    fn drop(&mut self) {
        let previous = self.in_flight_requests.fetch_sub(1, AtomicOrdering::AcqRel);
        if previous == 0 {
            self.in_flight_requests.store(0, AtomicOrdering::Release);
        }
    }
}

fn acquire_in_flight_request(
    in_flight_requests: &Arc<AtomicU32>,
    group_id: u32,
    overload_hardening: bool,
    overload_max_in_flight_per_group: u32,
) -> Result<InFlightRequestGuard, ServiceError> {
    let max_in_flight = overload_max_in_flight_per_group.max(1);
    loop {
        let current = in_flight_requests.load(AtomicOrdering::Acquire);
        if overload_hardening && current >= max_in_flight {
            return Err(ServiceError::Overloaded {
                group_id,
                retry_after_ms: OVERLOAD_RETRY_AFTER_MS,
            });
        }
        if current == u32::MAX {
            return Err(ServiceError::Storage(format!(
                "group {group_id} in-flight counter overflow"
            )));
        }
        if in_flight_requests
            .compare_exchange_weak(
                current,
                current + 1,
                AtomicOrdering::AcqRel,
                AtomicOrdering::Acquire,
            )
            .is_ok()
        {
            return Ok(InFlightRequestGuard {
                in_flight_requests: Arc::clone(in_flight_requests),
            });
        }
    }
}

fn load_in_flight_requests(in_flight_requests: &Arc<AtomicU32>) -> u32 {
    in_flight_requests.load(AtomicOrdering::Acquire)
}

fn validate_serving_group_id(group_id: u32) -> Result<(), ServiceError> {
    if group_id < MAX_KEYSPACE_GROUP_COUNT_IN_USE {
        Ok(())
    } else {
        Err(ServiceError::InvalidGroupId { group_id })
    }
}

fn validate_api_group_id(group_id: u32) -> Result<(), ServiceError> {
    if group_id < MAX_KEYSPACE_GROUP_COUNT_IN_USE {
        Ok(())
    } else {
        Err(ServiceError::InvalidGroupId { group_id })
    }
}

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

    use pd_proto_tso::{KeyspaceGroup, KeyspaceGroupMember};
    use pd_test_util::env::{EnvVarRestore, async_env_lock, remove_env, set_env};

    use crate::allocator::AllocatorConfig;
    use crate::keypath::timestamp_path;
    use crate::persistence::{
        InMemoryEtcdStore, StoreError, TimestampStore, encode_timestamp_nanos,
    };

    use super::{
        DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP, KeyspaceGroupManager, ServiceError,
        expected_primary_matches_local_identity,
    };

    fn new_group(id: u32, keyspaces: Vec<u32>) -> KeyspaceGroup {
        KeyspaceGroup {
            id,
            user_kind: "basic".to_string(),
            split_state: None,
            merge_state: None,
            members: Vec::new(),
            keyspaces,
        }
    }

    #[derive(Default)]
    struct LoadFailingStore;

    impl TimestampStore for LoadFailingStore {
        fn load_timestamp(
            &self,
            _cluster_id: u64,
            _group_id: u32,
        ) -> Result<Option<i64>, StoreError> {
            Err(StoreError::Internal("forced load failure".to_string()))
        }

        fn save_timestamp_if_leader(
            &self,
            _cluster_id: u64,
            _group_id: u32,
            _ts_nanos: i64,
            _leader_key: &str,
            _expected_leader_value: &str,
        ) -> Result<(), StoreError> {
            Ok(())
        }

        fn delete_timestamp(&self, _cluster_id: u64, _group_id: u32) -> Result<(), StoreError> {
            Ok(())
        }

        fn set_leader_value(
            &self,
            _leader_key: &str,
            _leader_value: &str,
        ) -> Result<(), StoreError> {
            Ok(())
        }
    }

    #[test]
    fn routing_uses_corrected_group_when_requested_group_is_stale() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register default");
        mgr.register_group(new_group(1, vec![11]))
            .expect("register group 1");
        mgr.campaign_group_primary(0, 10_000).expect("campaign g0");
        mgr.campaign_group_primary(1, 10_000).expect("campaign g1");

        let (_ts, resolved) = mgr
            .handle_tso_request(
                10_000,
                11,
                0,
                1,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )
            .expect("route and alloc");
        assert_eq!(resolved, 1);
    }

    #[test]
    fn split_rebases_target_and_clears_split_state() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(1, vec![11, 12]))
            .expect("register source");
        mgr.campaign_group_primary(1, 10_000)
            .expect("campaign source");
        let _ = mgr
            .handle_tso_request(
                10_000,
                11,
                1,
                5,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )
            .expect("warm source");

        mgr.split_keyspace_group(1, 2, vec![12])
            .expect("split group");
        mgr.campaign_group_primary(2, 10_000)
            .expect("campaign target");

        let (_ts, resolved) = mgr
            .handle_tso_request(
                10_000,
                12,
                2,
                1,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )
            .expect("split check + alloc");
        assert_eq!(resolved, 2);

        assert!(mgr.group(2).expect("target exists").split_state.is_none());
        assert!(mgr.group(1).expect("source exists").split_state.is_none());
    }

    #[test]
    fn keyspace_group_serving_boundary_matches_strict_parity() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");

        mgr.register_group(new_group(4095, vec![4095]))
            .expect("group 4095 is in serving range");

        let err = mgr
            .register_group(new_group(4096, vec![4096]))
            .expect_err("group 4096 must be rejected by serving path");
        assert!(err.to_string().contains("out of serving range"));
    }

    #[test]
    fn api_and_serving_boundaries_both_reject_group_4096() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register default");
        mgr.campaign_group_primary(0, 10_000).expect("campaign g0");

        let lookup_err = mgr
            .find_group_by_keyspace_id(10_000, 0, 4096)
            .expect_err("api lookup must reject 4096 boundary");
        assert!(lookup_err.to_string().contains("out of serving range"));

        let err = mgr
            .handle_tso_request(
                10_000,
                0,
                4096,
                1,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )
            .expect_err("serving path must reject 4096");
        assert!(err.to_string().contains("out of serving range"));
    }

    #[test]
    fn split_source_requests_continue_serving_until_target_finalizes_split() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(1, vec![11, 12]))
            .expect("register source");
        mgr.campaign_group_primary(1, 10_000)
            .expect("campaign source");
        let _ = mgr
            .handle_tso_request(
                10_000,
                11,
                1,
                3,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )
            .expect("warm source");
        mgr.split_keyspace_group(1, 2, vec![12])
            .expect("split group");
        mgr.campaign_group_primary(2, 10_000)
            .expect("campaign target");

        let (_source_ts, resolved_source) = mgr
            .handle_tso_request(
                10_000,
                11,
                1,
                1,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )
            .expect("source request during split should succeed");
        assert_eq!(resolved_source, 1);
        assert!(mgr.group(1).expect("source group").split_state.is_some());
        assert!(mgr.group(2).expect("target group").split_state.is_some());

        let (_target_ts, resolved_target) = mgr
            .handle_tso_request(
                10_000,
                12,
                2,
                1,
                false,
                DEFAULT_OVERLOAD_MAX_IN_FLIGHT_PER_GROUP,
            )
            .expect("target request should finalize split");
        assert_eq!(resolved_target, 2);
        assert!(mgr.group(1).expect("source group").split_state.is_none());
        assert!(mgr.group(2).expect("target group").split_state.is_none());
    }

    #[test]
    fn overload_hardening_rejects_requests_when_in_flight_limit_is_reached() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register default");
        mgr.campaign_group_primary(0, 10_000)
            .expect("campaign default");

        let prepared = mgr.prepare_tso_request(0, 0).expect("prepare request");
        let max_in_flight = 4;
        let workers = 32;
        let barrier = Arc::new(Barrier::new(workers));
        let runtime_guard = prepared
            .runtime
            .lock()
            .unwrap_or_else(|poison| poison.into_inner());
        let mut handles = Vec::with_capacity(workers);
        for _ in 0..workers {
            let prepared = prepared.clone();
            let barrier = Arc::clone(&barrier);
            handles.push(thread::spawn(move || -> Option<ServiceError> {
                barrier.wait();
                prepared.generate_tso(10_000, 1, true, max_in_flight).err()
            }));
        }
        thread::sleep(std::time::Duration::from_millis(50));
        drop(runtime_guard);

        let mut overloaded = 0_usize;
        for handle in handles {
            match handle.join().expect("join worker") {
                Some(ServiceError::Overloaded { .. }) => overloaded = overloaded.saturating_add(1),
                Some(other) => panic!("unexpected error: {other}"),
                None => {}
            }
        }
        assert!(
            overloaded > 0,
            "overload hardening should reject some requests under concurrent load"
        );
    }

    #[test]
    fn finish_merge_rebases_target_above_removed_sources() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr =
            KeyspaceGroupManager::new(1, store.clone(), AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register target");
        mgr.register_group(new_group(1, vec![1]))
            .expect("register source");
        mgr.campaign_group_primary(0, 10_000)
            .expect("campaign target");
        mgr.campaign_group_primary(1, 10_000)
            .expect("campaign source");

        store
            .put_raw(
                timestamp_path(1, 0),
                encode_timestamp_nanos(10_000_000_000).to_vec(),
            )
            .expect("seed target timestamp");
        store
            .put_raw(
                timestamp_path(1, 1),
                encode_timestamp_nanos(30_000_000_000).to_vec(),
            )
            .expect("seed source timestamp");

        mgr.merge_keyspace_groups(0, vec![1]).expect("merge");
        mgr.finish_merge_keyspace_group(0, 12_000)
            .expect("finish merge");

        assert!(mgr.group(0).expect("target").merge_state.is_none());
        let persisted = store
            .load_timestamp(1, 0)
            .expect("load target persisted")
            .expect("target timestamp");
        assert!(
            persisted > 30_000_000_000,
            "target must be rebased above all source timelines"
        );
    }

    #[test]
    fn cleanup_skips_merge_sources_before_finalize_and_deletes_after_finalize() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr =
            KeyspaceGroupManager::new(1, store.clone(), AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register target");
        mgr.register_group(new_group(1, vec![1]))
            .expect("register source");
        mgr.campaign_group_primary(0, 10_000)
            .expect("campaign target");
        mgr.campaign_group_primary(1, 10_000)
            .expect("campaign source");

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

        mgr.merge_keyspace_groups(0, vec![1]).expect("merge");
        mgr.cleanup_deleted_group_timestamps()
            .expect("cleanup before finalize");
        assert!(
            store.load_timestamp(1, 1).expect("load source").is_some(),
            "cleanup must not delete source timestamp before merge finalize"
        );

        mgr.finish_merge_keyspace_group(0, 12_000)
            .expect("finish merge");
        mgr.cleanup_deleted_group_timestamps()
            .expect("cleanup after finalize");
        assert!(
            store.load_timestamp(1, 1).expect("load source").is_none(),
            "cleanup should delete source timestamp after finalize"
        );
    }

    #[test]
    fn remove_group_marks_deleted_group_for_cleanup() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr =
            KeyspaceGroupManager::new(1, store.clone(), AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register group");
        store
            .put_raw(
                timestamp_path(1, 0),
                encode_timestamp_nanos(11_000_000_000).to_vec(),
            )
            .expect("seed timestamp");

        mgr.remove_group(0).expect("remove group");
        assert!(
            mgr.deleted_groups().contains(&0),
            "removed group should be tracked for cleanup"
        );
        mgr.cleanup_deleted_group_timestamps()
            .expect("cleanup deleted group");
        assert!(
            store
                .load_timestamp(1, 0)
                .expect("load timestamp")
                .is_none(),
            "removed group timestamp should be cleaned up by periodic cleaner"
        );
    }

    #[test]
    fn config_epoch_apply_commits_when_all_required_actors_ack() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register group");
        mgr.apply_config_epoch(2, &[0]).expect("apply config epoch");
        assert_eq!(mgr.group_config_epoch(0).expect("epoch"), 2);
    }

    #[test]
    fn config_epoch_apply_updates_allocator_runtime_config() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register group");

        let mut next = AllocatorConfig::default();
        next.oracle.save_interval_ms = 9_000;
        next.oracle.update_physical_interval_ms = 200;
        mgr.apply_config_epoch_with_allocator(3, &[0], Some(next))
            .expect("apply config with allocator");

        assert_eq!(mgr.group_config_epoch(0).expect("epoch"), 3);
        assert_eq!(mgr.allocator_config().oracle.save_interval_ms, 9_000);
        assert_eq!(
            mgr.allocator_config().oracle.update_physical_interval_ms,
            200
        );
    }

    #[test]
    fn config_epoch_apply_aborts_on_missing_required_actor() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register group");
        let err = mgr
            .apply_config_epoch(2, &[0, 1])
            .expect_err("missing actor must fail");
        assert!(err.to_string().contains("not served"));
        assert_eq!(
            mgr.group_config_epoch(0).expect("epoch unchanged"),
            0,
            "epoch must not partially apply on abort"
        );
    }

    #[test]
    fn expected_primary_identity_match_accepts_local_address_aliases() {
        assert!(expected_primary_matches_local_identity(
            "node-a:7", "node-a:7"
        ));
        assert!(expected_primary_matches_local_identity(
            "http://127.0.0.1:3379",
            "127.0.0.1:3379"
        ));
        assert!(!expected_primary_matches_local_identity(
            "127.0.0.2:3379",
            "127.0.0.1:3379"
        ));
    }

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

        set_env(
            "PD_TSO_LEGACY_DISCOVERY_URLS",
            "http://203.0.113.10:3379,http://203.0.113.11:3379",
        );
        remove_env("PD_TSO_ADVERTISE_LISTEN_ADDR");
        remove_env("PD_TSO_LISTEN_ADDR");
        assert!(
            !expected_primary_matches_local_identity("http://203.0.113.10:3379", "node-a:7"),
            "legacy discovery URL list must not be treated as local identity alias"
        );
    }

    #[test]
    fn discovery_tso_urls_keep_existing_order_across_membership_updates() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");

        let mut g0 = new_group(0, vec![0]);
        g0.members = vec![
            KeyspaceGroupMember {
                address: "http://10.0.0.1:3379".to_string(),
                is_primary: false,
            },
            KeyspaceGroupMember {
                address: "http://10.0.0.2:3379".to_string(),
                is_primary: false,
            },
        ];
        mgr.upsert_group(g0).expect("upsert g0");
        assert_eq!(
            mgr.discovery_tso_urls(),
            vec![
                "http://10.0.0.1:3379".to_string(),
                "http://10.0.0.2:3379".to_string()
            ]
        );

        let mut g1 = new_group(1, vec![1]);
        g1.members = vec![
            KeyspaceGroupMember {
                address: "http://10.0.0.2:3379".to_string(),
                is_primary: false,
            },
            KeyspaceGroupMember {
                address: "http://10.0.0.3:3379".to_string(),
                is_primary: false,
            },
        ];
        mgr.upsert_group(g1).expect("upsert g1");
        assert_eq!(
            mgr.discovery_tso_urls(),
            vec![
                "http://10.0.0.1:3379".to_string(),
                "http://10.0.0.2:3379".to_string(),
                "http://10.0.0.3:3379".to_string()
            ]
        );

        let mut g0_reordered = new_group(0, vec![0]);
        g0_reordered.members = vec![
            KeyspaceGroupMember {
                address: "http://10.0.0.2:3379".to_string(),
                is_primary: false,
            },
            KeyspaceGroupMember {
                address: "http://10.0.0.1:3379".to_string(),
                is_primary: false,
            },
        ];
        mgr.upsert_group(g0_reordered).expect("reorder g0");
        assert_eq!(
            mgr.discovery_tso_urls(),
            vec![
                "http://10.0.0.1:3379".to_string(),
                "http://10.0.0.2:3379".to_string(),
                "http://10.0.0.3:3379".to_string()
            ],
            "existing discovery order should remain stable across membership updates"
        );

        mgr.remove_group(0).expect("remove g0");
        assert_eq!(
            mgr.discovery_tso_urls(),
            vec![
                "http://10.0.0.2:3379".to_string(),
                "http://10.0.0.3:3379".to_string()
            ]
        );
    }

    #[test]
    fn find_group_primary_member_mapping_uses_address_equivalence() {
        let store = Arc::new(InMemoryEtcdStore::new());
        let mut mgr =
            KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node-b:1234");
        mgr.register_group(KeyspaceGroup {
            id: 0,
            user_kind: "basic".to_string(),
            split_state: None,
            merge_state: None,
            members: vec![
                KeyspaceGroupMember {
                    address: "node-a:1234".to_string(),
                    is_primary: false,
                },
                KeyspaceGroupMember {
                    address: "http://node-b:1234".to_string(),
                    is_primary: false,
                },
            ],
            keyspaces: vec![0],
        })
        .expect("register group");
        mgr.campaign_group_primary(0, 10_000).expect("campaign g0");
        assert!(
            mgr.group_runtime_snapshot(0, 10_000)
                .expect("runtime snapshot")
                .is_primary,
            "campaigned group should be serving at lookup timestamp"
        );
        assert!(expected_primary_matches_local_identity(
            "http://node-b:1234",
            "node-b:1234"
        ));

        let lookup = mgr
            .find_group_by_keyspace_id(10_000, 0, 0)
            .expect("find group");
        assert_eq!(lookup.group.members.len(), 2);
        assert!(!lookup.group.members[0].is_primary);
        assert!(
            lookup.group.members[1].is_primary,
            "member mapping = {:?}",
            lookup.group.members
        );
    }

    #[test]
    fn runtime_snapshot_does_not_depend_on_persisted_timestamp_store_reads() {
        let store = Arc::new(LoadFailingStore);
        let mut mgr = KeyspaceGroupManager::new(1, store, AllocatorConfig::default(), "node");
        mgr.register_group(new_group(0, vec![0]))
            .expect("register group");

        let runtime_snapshot = mgr
            .group_runtime_snapshot(0, 10_000)
            .expect("runtime snapshot should not touch persisted timestamp store");
        assert!(!runtime_snapshot.is_primary);

        let debug_err = mgr
            .group_debug_state(0, 10_000)
            .expect_err("debug state should include persisted timestamp load");
        assert!(matches!(debug_err, ServiceError::Storage(_)));
    }
}