// Copyright 2026 PingCAP, Inc.

use std::{
    collections::HashSet,
    sync::{
        Arc, Condvar, Mutex, RwLock,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};

use dashmap::{DashMap, DashSet};
use http::StatusCode;
use kvproto::keyspacepb::KeyspaceState;
use pd_client::PdClient;
use security::RestfulClient;
use serde::Deserialize;
use tikv_util::{info, sys::thread::StdThreadBuildWrapper, warn};

use crate::{
    cpu_config::CpuThrottleConfig,
    cpu_throttle::CpuThrottleManager,
    metrics::{
        SERVICE_LIMIT_FETCH_DURATION, SERVICE_LIMIT_FETCH_TOTAL, SERVICE_LIMIT_KEYSPACE_QUOTA_US,
        SERVICE_LIMIT_KEYSPACE_VALUE, SERVICE_LIMIT_LAST_SUCCESS_TIMESTAMP, keyspace_label,
        remove_keyspace_label,
    },
};

#[derive(Debug, Deserialize)]
struct ServiceLimitResponse {
    service_limit: f64,
}

#[derive(Debug, Deserialize)]
struct KeyspaceMetaResponse {
    id: u32,
    name: String,
    state: String,
}

#[derive(Debug)]
struct KeyspaceMeta {
    name: String,
    state: KeyspaceState,
}

#[derive(Debug)]
struct NotFoundState {
    first_seen: Instant,
    count: u64,
}

pub struct ServiceLimitFetcher {
    pd_client: Arc<dyn PdClient>,
    cpu_throttle_manager: Arc<CpuThrottleManager>,
    // Hot-updatable configuration. These were originally snapshotted at
    // construction; they are now reconfigurable so a node that boots in tenant
    // cap mode (with a dormant fetcher) can later hot-transition into PD
    // service_limit mode using the up-to-date config instead of stale
    // startup-time values.
    fetch_interval_ms: AtomicU64,
    request_timeout_ms: AtomicU64,
    not_found_ttl: Duration,
    not_found_consecutive: u64,
    update_threshold: f64,
    pd_override_config: AtomicBool,
    fallback_endpoints: Vec<String>,
    config_keyspaces: RwLock<HashSet<u32>>,
    /// Bumped on every `reconfigure`. A fetch reads the generation before
    /// issuing requests and re-checks it (plus `can_apply_service_limit`)
    /// before applying results, so a response computed under an old config or
    /// before a mode switch is not applied after the switch.
    config_generation: AtomicU64,
    known_keyspaces: DashSet<u32>,
    last_quota_us: DashMap<u32, u64>,
    last_log_at: DashMap<u32, Instant>,
    apply_quota_log_at: DashMap<u32, Instant>,
    not_found_states: DashMap<u32, NotFoundState>,
    stopped: AtomicBool,
    sleep_mutex: Mutex<()>,
    sleep_cond: Condvar,
    thread_handle: Mutex<Option<std::thread::JoinHandle<()>>>,
}

impl ServiceLimitFetcher {
    pub fn start(
        pd_client: Arc<dyn PdClient>,
        cpu_throttle_manager: Arc<CpuThrottleManager>,
        config: CpuThrottleConfig,
        fallback_endpoints: Vec<String>,
    ) -> Arc<Self> {
        let config_keyspaces = CpuThrottleConfig::parse_keyspace_quotas(
            &config.keyspace_cpu_quotas,
            config.refill_interval_ms,
        )
        .keys()
        .copied()
        .collect::<HashSet<_>>();

        let fetcher = Arc::new(ServiceLimitFetcher {
            pd_client,
            cpu_throttle_manager,
            fetch_interval_ms: AtomicU64::new(config.pd_service_limit_fetch_interval_ms.max(1)),
            request_timeout_ms: AtomicU64::new(config.pd_service_limit_request_timeout_ms.max(1)),
            not_found_ttl: Duration::from_millis(config.pd_service_limit_ttl_ms.max(1)),
            not_found_consecutive: config.pd_service_limit_consecutive_404.max(1),
            update_threshold: config.pd_service_limit_update_threshold,
            pd_override_config: AtomicBool::new(config.pd_override_config),
            fallback_endpoints,
            config_keyspaces: RwLock::new(config_keyspaces),
            config_generation: AtomicU64::new(0),
            known_keyspaces: DashSet::new(),
            last_quota_us: DashMap::new(),
            last_log_at: DashMap::new(),
            apply_quota_log_at: DashMap::new(),
            not_found_states: DashMap::new(),
            stopped: AtomicBool::new(false),
            sleep_mutex: Mutex::new(()),
            sleep_cond: Condvar::new(),
            thread_handle: Mutex::new(None),
        });

        for keyspace_id in fetcher.config_keyspaces.read().unwrap().iter().copied() {
            fetcher.known_keyspaces.insert(keyspace_id);
        }

        let thread_fetcher = fetcher.clone();
        let handle = std::thread::Builder::new()
            .name("pd-service-limit-fetcher".to_string())
            .spawn_wrapper(move || {
                thread_fetcher.run();
            })
            .expect("failed to spawn pd-service-limit-fetcher");
        *fetcher.thread_handle.lock().unwrap() = Some(handle);

        fetcher
    }

    /// Build a fetcher without spawning the background thread. Test-only: lets
    /// unit tests drive `apply_quotas` / `remove_keyspace` / `reconfigure`
    /// deterministically without a running fetch loop or a live PD client.
    #[cfg(test)]
    fn new_for_test(
        pd_client: Arc<dyn PdClient>,
        cpu_throttle_manager: Arc<CpuThrottleManager>,
        config: CpuThrottleConfig,
    ) -> Arc<Self> {
        let config_keyspaces = CpuThrottleConfig::parse_keyspace_quotas(
            &config.keyspace_cpu_quotas,
            config.refill_interval_ms,
        )
        .keys()
        .copied()
        .collect::<HashSet<_>>();

        let fetcher = Arc::new(ServiceLimitFetcher {
            pd_client,
            cpu_throttle_manager,
            fetch_interval_ms: AtomicU64::new(config.pd_service_limit_fetch_interval_ms.max(1)),
            request_timeout_ms: AtomicU64::new(config.pd_service_limit_request_timeout_ms.max(1)),
            not_found_ttl: Duration::from_millis(config.pd_service_limit_ttl_ms.max(1)),
            not_found_consecutive: config.pd_service_limit_consecutive_404.max(1),
            update_threshold: config.pd_service_limit_update_threshold,
            pd_override_config: AtomicBool::new(config.pd_override_config),
            fallback_endpoints: Vec::new(),
            config_keyspaces: RwLock::new(config_keyspaces),
            config_generation: AtomicU64::new(0),
            known_keyspaces: DashSet::new(),
            last_quota_us: DashMap::new(),
            last_log_at: DashMap::new(),
            apply_quota_log_at: DashMap::new(),
            not_found_states: DashMap::new(),
            stopped: AtomicBool::new(false),
            sleep_mutex: Mutex::new(()),
            sleep_cond: Condvar::new(),
            thread_handle: Mutex::new(None),
        });
        for keyspace_id in fetcher.config_keyspaces.read().unwrap().iter().copied() {
            fetcher.known_keyspaces.insert(keyspace_id);
        }
        fetcher
    }

    pub fn register_keyspace(&self, keyspace_id: u32) {
        // Registration is called from the coprocessor admission path. Avoid a
        // write lock on every request once the keyspace is known.
        if self.known_keyspaces.contains(&keyspace_id) {
            return;
        }
        if self.known_keyspaces.insert(keyspace_id) {
            info!(
                "[ServiceLimitFetcher] register keyspace";
                "keyspace_id" => keyspace_id,
            );
        }
    }

    /// Hot-update the fetcher's configuration. The fetcher thread is long-lived
    /// (it stays resident even while dormant in tenant cap mode) so that a
    /// later transition into PD service_limit mode does not require a restart.
    /// Callers must invoke this *before* enabling PD service_limit so the first
    /// fetch after the switch uses the new config rather than startup-time
    /// values.
    pub fn reconfigure(&self, config: &CpuThrottleConfig) {
        self.fetch_interval_ms.store(
            config.pd_service_limit_fetch_interval_ms.max(1),
            Ordering::Release,
        );
        self.request_timeout_ms.store(
            config.pd_service_limit_request_timeout_ms.max(1),
            Ordering::Release,
        );
        self.pd_override_config
            .store(config.pd_override_config, Ordering::Release);

        let new_config_keyspaces = CpuThrottleConfig::parse_keyspace_quotas(
            &config.keyspace_cpu_quotas,
            config.refill_interval_ms,
        )
        .keys()
        .copied()
        .collect::<HashSet<_>>();
        {
            let mut guard = self.config_keyspaces.write().unwrap();
            *guard = new_config_keyspaces.clone();
        }
        // Merge into the live keyspace set without dropping keyspaces that were
        // discovered/registered at runtime.
        for keyspace_id in new_config_keyspaces {
            self.known_keyspaces.insert(keyspace_id);
        }

        // Bump last so the new config is fully visible before any fetch that
        // observes the new generation runs its post-fetch apply check.
        self.config_generation.fetch_add(1, Ordering::AcqRel);
        // Wake the fetcher so a cap -> PD service_limit transition starts its
        // first fetch with the new config promptly, instead of sleeping out the
        // (possibly stale/long) old fetch interval first.
        self.sleep_cond.notify_all();
        info!(
            "[ServiceLimitFetcher] reconfigured";
            "fetch_interval_ms" => config.pd_service_limit_fetch_interval_ms,
            "request_timeout_ms" => config.pd_service_limit_request_timeout_ms,
            "override_config" => config.pd_override_config,
        );
    }

    fn config_generation(&self) -> u64 {
        self.config_generation.load(Ordering::Acquire)
    }

    fn fetch_interval(&self) -> Duration {
        Duration::from_millis(self.fetch_interval_ms.load(Ordering::Acquire))
    }

    fn request_timeout(&self) -> Duration {
        Duration::from_millis(self.request_timeout_ms.load(Ordering::Acquire))
    }

    fn pd_override_config(&self) -> bool {
        self.pd_override_config.load(Ordering::Acquire)
    }

    fn config_keyspace_contains(&self, keyspace_id: u32) -> bool {
        self.config_keyspaces.read().unwrap().contains(&keyspace_id)
    }

    pub fn stop(&self) {
        self.stopped.store(true, Ordering::Release);
        self.sleep_cond.notify_all();
    }

    pub fn stop_and_join(&self) {
        self.stop();
        if let Some(handle) = self.thread_handle.lock().unwrap().take() {
            let _ = handle.join();
        }
    }

    fn run(self: Arc<Self>) {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        loop {
            if self.stopped.load(Ordering::Acquire) {
                break;
            }

            let _timer = SERVICE_LIMIT_FETCH_DURATION.start_timer();
            let result = runtime.block_on(self.fetch_and_update());
            match result {
                Ok(()) => {
                    SERVICE_LIMIT_FETCH_TOTAL
                        .with_label_values(&["success"])
                        .inc();
                    let ts = SystemTime::now()
                        .duration_since(UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_secs_f64();
                    SERVICE_LIMIT_LAST_SUCCESS_TIMESTAMP.set(ts);
                }
                Err(err) => {
                    SERVICE_LIMIT_FETCH_TOTAL
                        .with_label_values(&["error"])
                        .inc();
                    warn!(
                        "[ServiceLimitFetcher] fetch failed";
                        "error" => %err,
                    );
                }
            }

            let guard = self.sleep_mutex.lock().unwrap();
            let _ = self
                .sleep_cond
                .wait_timeout(guard, self.fetch_interval())
                .unwrap();
        }
    }

    async fn fetch_and_update(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if !self.can_apply_service_limit() {
            return Ok(());
        }

        // Snapshot the config generation before issuing any request. If a
        // concurrent `reconfigure` (or a mode switch) lands while we are
        // fetching, the generation changes and we skip applying this round's
        // (now potentially stale) results.
        let fetch_generation = self.config_generation();

        let client = match self.build_rest_client() {
            Some(client) => client,
            None => {
                return Err(tikv_util::box_err!(
                    "no PD endpoints available for service_limit fetch"
                ));
            }
        };

        let mut service_limits = std::collections::HashMap::new();

        let keyspace_ids: Vec<u32> = self.known_keyspaces.iter().map(|v| *v.key()).collect();
        for keyspace_id in keyspace_ids {
            let meta = match self
                .fetch_keyspace_meta(&client, keyspace_id, fetch_generation)
                .await
            {
                Ok(Some(meta)) => meta,
                Ok(None) => continue,
                Err(err) => {
                    warn!(
                        "[ServiceLimitFetcher] fetch keyspace meta failed";
                        "keyspace_id" => keyspace_id,
                        "error" => %err,
                    );
                    continue;
                }
            };
            if matches!(
                meta.state,
                KeyspaceState::Archived | KeyspaceState::Tombstone
            ) {
                info!(
                    "[ServiceLimitFetcher] keyspace is archived or tombstone, need to remove";
                    "keyspace_id" => keyspace_id,
                    "state" => ?meta.state,
                );
                self.remove_keyspace(keyspace_id, fetch_generation);
                continue;
            }
            if !self.pd_override_config() && self.config_keyspace_contains(keyspace_id) {
                continue;
            }

            let path = format!(
                "resource-manager/api/v1/config/keyspace/service-limit/{}",
                meta.name.as_str()
            );
            match client.get_with_status::<ServiceLimitResponse>(&path).await {
                Ok((status, resp)) => {
                    if status == StatusCode::NOT_FOUND {
                        self.on_not_found(keyspace_id, fetch_generation);
                        continue;
                    }
                    if !status.is_success() {
                        warn!(
                            "[ServiceLimitFetcher] request failed";
                            "keyspace_id" => keyspace_id,
                            "status" => ?status,
                        );
                        continue;
                    }
                    let resp = match resp {
                        Some(resp) => resp,
                        None => {
                            warn!(
                                "[ServiceLimitFetcher] request failed";
                                "keyspace_id" => keyspace_id,
                                "status" => ?status,
                            );
                            continue;
                        }
                    };
                    // Drop stale success responses before they touch any local
                    // state: clearing not_found_states / updating metrics / log
                    // timestamps from a pre-reconfigure response could cancel or
                    // delay a legitimate later removal decision.
                    if !self.can_apply_fetch_result(fetch_generation) {
                        continue;
                    }
                    self.not_found_states.remove(&keyspace_id);
                    let limit = resp.service_limit.max(0.0);
                    service_limits.insert(keyspace_id, limit);
                    let keyspace_label = keyspace_label(keyspace_id);
                    SERVICE_LIMIT_KEYSPACE_VALUE
                        .with_label_values(&[keyspace_label.as_ref()])
                        .set(limit);
                    let log_interval = Duration::from_secs(60);
                    let should_log = match self.last_log_at.get(&keyspace_id) {
                        Some(last) => last.elapsed() >= log_interval,
                        None => true,
                    };
                    if should_log {
                        self.last_log_at.insert(keyspace_id, Instant::now());
                        info!(
                            "[ServiceLimitFetcher] fetched service_limit";
                            "keyspace_id" => keyspace_id,
                            "keyspace_name" => meta.name.as_str(),
                            "service_limit" => limit,
                        );
                    }
                }
                Err(err) => {
                    warn!(
                        "[ServiceLimitFetcher] request failed";
                        "keyspace_id" => keyspace_id,
                        "error" => %err,
                    );
                }
            }
        }

        // Re-check before applying: the mode may have switched (cap re-enabled
        // / service_limit disabled) or the config may have been reconfigured
        // while we were fetching. In either case this round's results are stale
        // and must not be applied. `apply_quotas` re-checks this guard before
        // every individual manager mutation as well, so a reconfigure that
        // lands mid-apply is also caught.
        if !self.can_apply_fetch_result(fetch_generation) {
            info!(
                "[ServiceLimitFetcher] config changed during fetch, skip applying stale results";
            );
            return Ok(());
        }

        self.apply_quotas(service_limits, fetch_generation);
        Ok(())
    }

    fn build_rest_client(&self) -> Option<RestfulClient> {
        let mut endpoints = self.pd_client.get_leader_client_urls();
        if endpoints.is_empty() {
            endpoints = self.fallback_endpoints.clone();
        }
        if endpoints.is_empty() {
            return None;
        }
        let mut client = RestfulClient::new(
            "service_limit_fetcher",
            endpoints,
            self.pd_client.get_security_mgr(),
        )
        .ok()?;
        client.set_timeout(self.request_timeout(), self.request_timeout());
        Some(client)
    }

    async fn fetch_keyspace_meta(
        &self,
        client: &RestfulClient,
        keyspace_id: u32,
        fetch_generation: u64,
    ) -> Result<Option<KeyspaceMeta>, Box<dyn std::error::Error + Send + Sync>> {
        let path = format!("pd/api/v2/keyspaces/id/{}", keyspace_id);
        match client.get_with_status::<KeyspaceMetaResponse>(&path).await {
            Ok((status, resp)) => {
                if status == StatusCode::NOT_FOUND {
                    self.on_not_found(keyspace_id, fetch_generation);
                    return Ok(None);
                }
                if !status.is_success() {
                    warn!(
                        "[ServiceLimitFetcher] keyspace meta request failed";
                        "keyspace_id" => keyspace_id,
                        "status" => ?status,
                    );
                    return Ok(None);
                }
                let resp = match resp {
                    Some(resp) => resp,
                    None => {
                        warn!(
                            "[ServiceLimitFetcher] keyspace meta request failed";
                            "keyspace_id" => keyspace_id,
                            "status" => ?status,
                        );
                        return Ok(None);
                    }
                };
                if resp.id != keyspace_id {
                    warn!(
                        "[ServiceLimitFetcher] keyspace meta id mismatch";
                        "keyspace_id" => keyspace_id,
                        "resp_id" => resp.id,
                    );
                    return Ok(None);
                }
                let state = match parse_keyspace_state(&resp.state) {
                    Some(state) => state,
                    None => {
                        warn!(
                            "[ServiceLimitFetcher] keyspace meta has unknown state";
                            "keyspace_id" => keyspace_id,
                            "state" => resp.state.as_str(),
                        );
                        return Ok(None);
                    }
                };
                Ok(Some(KeyspaceMeta {
                    name: resp.name,
                    state,
                }))
            }
            Err(err) => Err(Box::new(err)),
        }
    }

    fn on_not_found(&self, keyspace_id: u32, fetch_generation: u64) -> bool {
        // Drop stale 404s entirely: a response from before a reconfigure / mode
        // switch must not even mutate the local not_found counter, otherwise a
        // later current-generation 404 could borrow the stale count and trip
        // `pd_service_limit_consecutive_404` early.
        if !self.can_apply_fetch_result(fetch_generation) {
            return false;
        }
        let now = Instant::now();
        let mut entry = self
            .not_found_states
            .entry(keyspace_id)
            .or_insert(NotFoundState {
                first_seen: now,
                count: 0,
            });
        entry.count = entry.count.saturating_add(1);
        let should_remove = entry.count >= self.not_found_consecutive
            || entry.first_seen.elapsed() >= self.not_found_ttl;
        drop(entry);

        if should_remove {
            self.remove_keyspace(keyspace_id, fetch_generation);
            true
        } else {
            false
        }
    }

    fn remove_keyspace(&self, keyspace_id: u32, fetch_generation: u64) {
        if !self.can_apply_fetch_result(fetch_generation) {
            info!(
                "[ServiceLimitFetcher] skip stale keyspace removal (config changed during fetch)";
                "keyspace_id" => keyspace_id,
            );
            return;
        }
        if !self
            .cpu_throttle_manager
            .remove_keyspace_for_pd_service_limit(keyspace_id)
        {
            info!(
                "[ServiceLimitFetcher] skip keyspace removal";
                "keyspace_id" => keyspace_id,
            );
            return;
        }

        info!(
            "[ServiceLimitFetcher] remove keyspace";
            "keyspace_id" => keyspace_id,
        );
        self.known_keyspaces.remove(&keyspace_id);
        self.last_quota_us.remove(&keyspace_id);
        self.last_log_at.remove(&keyspace_id);
        self.apply_quota_log_at.remove(&keyspace_id);
        self.not_found_states.remove(&keyspace_id);

        let keyspace_label = keyspace_label(keyspace_id);
        let keyspace_label = keyspace_label.as_ref();
        let _ = SERVICE_LIMIT_KEYSPACE_VALUE.remove_label_values(&[keyspace_label]);
        let _ = SERVICE_LIMIT_KEYSPACE_QUOTA_US.remove_label_values(&[keyspace_label]);
        remove_keyspace_label(keyspace_id);
    }

    fn apply_quotas(
        &self,
        service_limits: std::collections::HashMap<u32, f64>,
        fetch_generation: u64,
    ) {
        if !self.can_apply_fetch_result(fetch_generation) {
            return;
        }

        let mut limited_keyspaces = Vec::new();
        for (keyspace_id, limit) in &service_limits {
            if *limit > 0.0 {
                limited_keyspaces.push((*keyspace_id, *limit));
            }
        }

        let total_limit: f64 = limited_keyspaces.iter().map(|(_, limit)| *limit).sum();
        let global_refill_rate_us = self.cpu_throttle_manager.global_refill_rate();

        for (keyspace_id, limit) in service_limits {
            if !self.can_apply_fetch_result(fetch_generation) {
                return;
            }
            if limit <= 0.0
                && (self.pd_override_config() || !self.config_keyspace_contains(keyspace_id))
            {
                if self
                    .cpu_throttle_manager
                    .remove_keyspace_for_pd_service_limit(keyspace_id)
                {
                    self.last_quota_us.remove(&keyspace_id);
                    let keyspace_label = keyspace_label(keyspace_id);
                    SERVICE_LIMIT_KEYSPACE_QUOTA_US
                        .with_label_values(&[keyspace_label.as_ref()])
                        .set(0);
                } else {
                    return;
                }
            }
        }

        if total_limit <= 0.0 || global_refill_rate_us == 0 {
            return;
        }

        let apply_log_interval = Duration::from_secs(120);
        for (keyspace_id, limit) in limited_keyspaces {
            if !self.can_apply_fetch_result(fetch_generation) {
                return;
            }
            let ratio = limit / total_limit;
            let mut quota_us = (global_refill_rate_us as f64 * ratio) as u64;
            if quota_us == 0 {
                quota_us = 1;
            }

            let should_update = self.should_update_quota(keyspace_id, quota_us);
            let applied_quota = if should_update {
                if self
                    .cpu_throttle_manager
                    .update_keyspace_quota_for_pd_service_limit(keyspace_id, quota_us)
                {
                    self.last_quota_us.insert(keyspace_id, quota_us);
                    quota_us
                } else {
                    return;
                }
            } else {
                self.last_quota_us
                    .get(&keyspace_id)
                    .map(|v| *v)
                    .unwrap_or(quota_us)
            };
            SERVICE_LIMIT_KEYSPACE_QUOTA_US
                .with_label_values(&[keyspace_label(keyspace_id).as_ref()])
                .set(clamp_gauge_value(applied_quota));

            let should_log = match self.apply_quota_log_at.get(&keyspace_id) {
                Some(last) => last.elapsed() >= apply_log_interval,
                None => true,
            };
            if should_log {
                self.apply_quota_log_at.insert(keyspace_id, Instant::now());
                info!(
                    "[ServiceLimitFetcher] apply quota";
                    "keyspace_id" => keyspace_id,
                    "limit" => limit,
                    "applied_quota_us" => applied_quota,
                    "total_limit" => total_limit,
                    "global_refill_rate_us" => global_refill_rate_us,
                    "self ratio" => ratio,
                    "need update" => should_update,
                );
            }
        }
    }

    fn can_apply_service_limit(&self) -> bool {
        self.cpu_throttle_manager.pd_service_limit_enabled()
            && !self.cpu_throttle_manager.tenant_cap_enabled()
    }

    /// Guard for every manager mutation produced by a fetch round. Results may
    /// only be applied while service_limit is still applicable AND the config
    /// generation that the round was computed under is still current. If a
    /// `reconfigure()` (or a cap/enable switch) landed mid-fetch, the round is
    /// stale and must not mutate keyspace quotas — including the remove paths
    /// for archived/tombstone/404 keyspaces.
    fn can_apply_fetch_result(&self, fetch_generation: u64) -> bool {
        self.can_apply_service_limit() && self.config_generation() == fetch_generation
    }

    fn should_update_quota(&self, keyspace_id: u32, new_quota_us: u64) -> bool {
        if self.update_threshold <= 0.0 {
            return true;
        }
        match self.last_quota_us.get(&keyspace_id) {
            Some(old_quota) => {
                let old = *old_quota as f64;
                let base = old.max(1.0);
                ((new_quota_us as f64 - old).abs() / base) > self.update_threshold
            }
            None => true,
        }
    }
}

fn parse_keyspace_state(state: &str) -> Option<KeyspaceState> {
    match state.to_ascii_uppercase().as_str() {
        "ENABLED" => Some(KeyspaceState::Enabled),
        "DISABLED" => Some(KeyspaceState::Disabled),
        "ARCHIVED" => Some(KeyspaceState::Archived),
        "TOMBSTONE" => Some(KeyspaceState::Tombstone),
        _ => None,
    }
}

fn clamp_gauge_value(value: u64) -> i64 {
    if value > i64::MAX as u64 {
        i64::MAX
    } else {
        value as i64
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use pd_client::PdClient;
    use security::{GetSecurityManager, SecurityManager};

    use super::*;
    use crate::CpuThrottleConfig;

    // Minimal no-op PdClient: the generation-guard tests drive apply/remove
    // directly and never issue real PD requests, so all trait methods can use
    // their defaults.
    struct NoopPdClient;

    impl GetSecurityManager for NoopPdClient {
        fn get_security_mgr(&self) -> Arc<SecurityManager> {
            Arc::new(SecurityManager::default())
        }
    }

    impl PdClient for NoopPdClient {}

    fn service_limit_config() -> CpuThrottleConfig {
        let mut config = CpuThrottleConfig::default();
        config.enforce_global_limit = true;
        config.enable_pd_service_limit = true;
        config.tenant_max_cpu_ratio = 0.0;
        config.enable_burst = false;
        config.enable_fair_allocation = false;
        config.enable_dynamic_adjustment = false;
        config.low_priority_dag_global_enabled = false;
        config.pd_override_config = true;
        config
    }

    fn build_fetcher() -> (Arc<CpuThrottleManager>, Arc<ServiceLimitFetcher>) {
        let config = service_limit_config();
        let manager = Arc::new(CpuThrottleManager::new(
            config.clone(),
            100_000,
            HashMap::new(),
        ));
        manager.update_pd_service_limit_enabled(true);
        let fetcher =
            ServiceLimitFetcher::new_for_test(Arc::new(NoopPdClient), manager.clone(), config);
        (manager, fetcher)
    }

    #[test]
    fn test_apply_quotas_skipped_when_generation_changed_mid_fetch() {
        let (manager, fetcher) = build_fetcher();
        assert!(fetcher.can_apply_service_limit());

        // Snapshot the generation a "fetch" started under, then simulate a
        // concurrent reconfigure landing before the results are applied.
        let fetch_generation = fetcher.config_generation();
        fetcher.reconfigure(&service_limit_config());
        assert_ne!(fetcher.config_generation(), fetch_generation);

        let mut service_limits = HashMap::new();
        service_limits.insert(7u32, 100.0);
        fetcher.apply_quotas(service_limits, fetch_generation);

        // Stale results must not have created a quota.
        assert_eq!(manager.keyspace_available(7), None);
    }

    #[test]
    fn test_apply_quotas_applied_when_generation_unchanged() {
        let (manager, fetcher) = build_fetcher();
        let fetch_generation = fetcher.config_generation();

        let mut service_limits = HashMap::new();
        service_limits.insert(7u32, 100.0);
        fetcher.apply_quotas(service_limits, fetch_generation);

        // The service-limit share is a per-interval refill rate. With the
        // default 10ms refill and 100ms burst window, a 100,000us refill rate
        // produces a 1,000,000us capacity, not a second 10x scale-up.
        assert_eq!(fetcher.last_quota_us.get(&7).map(|v| *v), Some(100_000));
        assert_eq!(manager.keyspace_available(7), Some(1_000_000));
    }

    #[test]
    fn test_remove_keyspace_skipped_when_generation_changed_mid_fetch() {
        let (manager, fetcher) = build_fetcher();
        // Seed a keyspace quota so a removal would be observable.
        let fetch_generation = fetcher.config_generation();
        let mut service_limits = HashMap::new();
        service_limits.insert(7u32, 100.0);
        fetcher.apply_quotas(service_limits, fetch_generation);
        assert!(manager.keyspace_available(7).is_some());

        // A new fetch round starts, then a reconfigure lands before its
        // archived/tombstone/404 removal would be applied.
        let stale_generation = fetcher.config_generation();
        fetcher.reconfigure(&service_limit_config());
        fetcher.remove_keyspace(7, stale_generation);

        // The stale removal must be ignored.
        assert!(manager.keyspace_available(7).is_some());

        // A removal under the current generation does take effect.
        let current_generation = fetcher.config_generation();
        fetcher.remove_keyspace(7, current_generation);
        assert_eq!(manager.keyspace_available(7), None);
    }

    #[test]
    fn test_stale_not_found_does_not_leak_counter_into_current_generation() {
        let config = service_limit_config();
        let manager = Arc::new(CpuThrottleManager::new(
            config.clone(),
            100_000,
            HashMap::new(),
        ));
        manager.update_pd_service_limit_enabled(true);
        // consecutive_404 = 2 so a single leaked stale count would be enough to
        // trip removal on the first current-generation 404.
        let mut cfg = config.clone();
        cfg.pd_service_limit_consecutive_404 = 2;
        // Large TTL so removal can only come from the consecutive count, not the
        // time-based path.
        cfg.pd_service_limit_ttl_ms = 60_000;
        let fetcher =
            ServiceLimitFetcher::new_for_test(Arc::new(NoopPdClient), manager.clone(), cfg);

        // Seed a quota so a removal would be observable.
        let g0 = fetcher.config_generation();
        let mut service_limits = HashMap::new();
        service_limits.insert(7u32, 100.0);
        fetcher.apply_quotas(service_limits, g0);
        assert!(manager.keyspace_available(7).is_some());

        // A stale-generation 404 must be dropped entirely (no counter bump).
        let stale = fetcher.config_generation();
        fetcher.reconfigure(&service_limit_config());
        assert!(!fetcher.on_not_found(7, stale));

        // First current-generation 404 must NOT remove yet (count should be 1,
        // not 2). If the stale 404 had leaked its count, this would remove.
        let current = fetcher.config_generation();
        assert!(!fetcher.on_not_found(7, current));
        assert!(manager.keyspace_available(7).is_some());

        // Second current-generation 404 reaches the threshold and removes.
        assert!(fetcher.on_not_found(7, current));
        assert_eq!(manager.keyspace_available(7), None);
    }
}