// Copyright 2025 PingCAP, Inc.

use std::{
    sync::{Arc, Mutex},
    time::Instant,
};

use tikv_util::{info, sys::SysQuota, warn};

use crate::{
    ACTIVE_KEYSPACE_READ_BYTES, Config, CpuThrottleConfig, CpuType, ReadLimiter, ResourceEvent,
    ResourceType, Severity, SeverityThreshold, TotalCpuMonitor, TransferLeaderLimiter,
    metrics::{CPU_USAGE_MONITOR_GLOBAL, CPU_USAGE_MONITOR_KEYSPACE, keyspace_label},
    usage::{KeyspaceInstantUsages, Usages},
};

#[async_trait::async_trait]
pub trait ResourceSubscriber: Send + Sync {
    fn on_event(&mut self, event: &ResourceEvent);
}

pub struct ReadSubscriberData {
    pub(crate) config: Config,
    pub(crate) max_cpu_time: f64,
    pub(crate) max_read_bytes: u64,
    pub(crate) severity_threshold: SeverityThreshold,
    pub(crate) cpu_severity: Severity,
    pub(crate) global_cpu_ratio: f64,
    pub(crate) usages: KeyspaceInstantUsages,
    pub(crate) read_limiter: ReadLimiter,
    pub(crate) cpu_throttle_manager: Option<Arc<crate::CpuThrottleManager>>,
    /// CPU quota captured when this subscriber is initialized. Cache it so the
    /// periodic usage path does not repeatedly read and parse cgroup files.
    cpu_cores: f64,
    /// The `max_read_cpu_ratio` the global CPU bucket was sized from at
    /// startup. The fair-gate charge scale is always relative to this value,
    /// never to a later config's ratio, because the bucket is never resized.
    /// The non-fair legacy path continues to use online config updates for its
    /// usage denominator without resizing that bucket.
    startup_max_read_cpu_ratio: f64,
}

impl ReadSubscriberData {
    pub fn new(read_limiter: ReadLimiter, config: Config) -> Self {
        let mut subscriber = ReadSubscriberData {
            config: config.clone(),
            max_cpu_time: Default::default(),
            max_read_bytes: Default::default(),
            severity_threshold: Default::default(),
            cpu_severity: Default::default(),
            global_cpu_ratio: 0.0,
            usages: Default::default(),
            read_limiter,
            cpu_throttle_manager: None,
            cpu_cores: SysQuota::cpu_cores_quota().max(1.0),
            startup_max_read_cpu_ratio: config.max_read_cpu_ratio,
        };
        subscriber.update_config(&config);
        subscriber
    }

    fn max_cpu_time_for_ratio(&self, config: &Config, ratio: f64) -> f64 {
        let window_secs = config.window_size.as_secs_f64();
        self.cpu_cores * ratio * window_secs
    }

    fn effective_max_read_cpu_ratio(&self, config: &Config) -> f64 {
        if self.read_limiter.is_legacy_read_limiter_allowed() {
            return config.max_read_cpu_ratio;
        }
        self.cpu_throttle_manager
            .as_ref()
            .map_or(config.max_read_cpu_ratio, |manager| {
                manager.enforced_max_read_cpu_ratio()
            })
    }

    fn refresh_max_cpu_time(&mut self) {
        let ratio = self.effective_max_read_cpu_ratio(&self.config);
        self.max_cpu_time = self.max_cpu_time_for_ratio(&self.config, ratio);
    }

    pub fn set_cpu_throttle_manager(&mut self, manager: Arc<crate::CpuThrottleManager>) {
        manager.initialize_fair_max_read_cpu_ratio(self.startup_max_read_cpu_ratio);
        self.cpu_throttle_manager = Some(manager);
        self.refresh_max_cpu_time();
    }

    pub fn update_config(&mut self, new_config: &Config) {
        let window_secs = new_config.window_size.as_secs_f64();
        self.max_read_bytes = (new_config.max_read_bytes_per_core_per_second.0 as f64
            * self.cpu_cores
            * window_secs) as u64;
        self.severity_threshold = SeverityThreshold::from_config(new_config);
        if self.config.get_read_enabled() != new_config.get_read_enabled() {
            self.read_limiter
                .update_enabled(new_config.get_read_enabled());
        }
        if self.config.dry_run != new_config.dry_run {
            if let (true, Some(manager)) = (
                self.config.dry_run && !new_config.dry_run,
                self.cpu_throttle_manager.as_ref(),
            ) {
                manager.reset_fair_global_cpu_after_dry_run();
            }
            self.read_limiter.update_dry_run(new_config.dry_run);
        }
        if self.config.limiter_timeout != new_config.limiter_timeout {
            self.read_limiter
                .update_timeout(new_config.limiter_timeout.0);
        }
        if self.config.limiter_stats_interval != new_config.limiter_stats_interval {
            self.read_limiter
                .update_stats_interval(new_config.limiter_stats_interval.0);
        }
        if self.config.max_read_wait_time != new_config.max_read_wait_time {
            self.read_limiter
                .update_max_wait_time(new_config.max_read_wait_time.0);
        }
        if let Some(ref manager) = self.cpu_throttle_manager {
            let old_cpu = &self.config.cpu_throttle;
            let new_cpu = &new_config.cpu_throttle;
            let old_tenant_cap_enabled = old_cpu.tenant_max_cpu_ratio > 0.0;
            let new_tenant_cap_enabled = new_cpu.tenant_max_cpu_ratio > 0.0;

            if self.config.debug != new_config.debug {
                manager.update_debug(new_config.debug);
            }

            if old_tenant_cap_enabled && !new_tenant_cap_enabled {
                manager.update_tenant_max_cpu_ratio(0.0);
            }

            if old_cpu.tenant_cap_idle_ttl != new_cpu.tenant_cap_idle_ttl {
                manager.update_tenant_cap_idle_ttl(new_cpu.tenant_cap_idle_ttl.as_millis());
            }

            if old_cpu.enable_dynamic_adjustment != new_cpu.enable_dynamic_adjustment
                || old_cpu.high_watermark != new_cpu.high_watermark
                || old_cpu.low_watermark != new_cpu.low_watermark
                || old_tenant_cap_enabled != new_tenant_cap_enabled
            {
                manager.update_dynamic_adjustment(
                    new_cpu.enable_dynamic_adjustment,
                    new_cpu.high_watermark,
                    new_cpu.low_watermark,
                );
            }

            if old_cpu.low_priority_dag_global_enabled != new_cpu.low_priority_dag_global_enabled
                || old_cpu.low_priority_dag_global_ratio != new_cpu.low_priority_dag_global_ratio
                || old_tenant_cap_enabled != new_tenant_cap_enabled
            {
                manager.update_low_priority_dag_settings(
                    new_cpu.low_priority_dag_global_enabled,
                    new_cpu.low_priority_dag_global_ratio,
                );
            }

            if old_cpu.enable_fair_allocation != new_cpu.enable_fair_allocation
                || old_cpu.fair_allocation_threshold != new_cpu.fair_allocation_threshold
                || old_tenant_cap_enabled != new_tenant_cap_enabled
            {
                manager.update_fair_allocation(
                    new_cpu.enable_fair_allocation,
                    new_cpu.fair_allocation_threshold,
                );
            }

            if old_cpu.enforce_global_limit != new_cpu.enforce_global_limit
                || old_tenant_cap_enabled != new_tenant_cap_enabled
            {
                manager.update_enforce_global_limit(new_cpu.enforce_global_limit);
            }

            if old_cpu.enable_pd_service_limit != new_cpu.enable_pd_service_limit
                || old_tenant_cap_enabled != new_tenant_cap_enabled
                || old_cpu.pd_service_limit_fetch_interval
                    != new_cpu.pd_service_limit_fetch_interval
                || old_cpu.pd_service_limit_request_timeout
                    != new_cpu.pd_service_limit_request_timeout
                || old_cpu.pd_override_config != new_cpu.pd_override_config
                || old_cpu.keyspace_cpu_quotas != new_cpu.keyspace_cpu_quotas
            {
                // Push the latest PD config into the resident fetcher *before*
                // flipping the enabled flag, so the first fetch after a
                // cap -> PD service_limit transition uses fresh config rather
                // than startup-time values.
                self.read_limiter
                    .reconfigure_service_limit_fetcher(&new_config.to_cpu_throttle_config());
                manager.update_pd_service_limit_enabled(
                    new_cpu.enable_pd_service_limit && !new_tenant_cap_enabled,
                );
            }

            if old_cpu.enable_burst != new_cpu.enable_burst
                || old_cpu.burst_threshold != new_cpu.burst_threshold
                || old_cpu.keyspace_burst_enabled != new_cpu.keyspace_burst_enabled
                || old_tenant_cap_enabled != new_tenant_cap_enabled
            {
                let keyspace_burst_enabled = CpuThrottleConfig::parse_keyspace_burst_enabled(
                    &new_cpu.keyspace_burst_enabled,
                );
                manager.update_burst_settings(
                    new_cpu.enable_burst,
                    new_cpu.burst_threshold,
                    keyspace_burst_enabled,
                );
            }

            if (new_tenant_cap_enabled || !old_tenant_cap_enabled)
                && old_cpu.tenant_max_cpu_ratio != new_cpu.tenant_max_cpu_ratio
            {
                manager.update_tenant_max_cpu_ratio(new_cpu.tenant_max_cpu_ratio);
            }

            if old_cpu.keyspace_cpu_quotas != new_cpu.keyspace_cpu_quotas {
                let refill_interval_ms = manager.refill_interval_ms();
                let quotas = CpuThrottleConfig::parse_keyspace_quotas(
                    &new_cpu.keyspace_cpu_quotas,
                    refill_interval_ms,
                );
                manager.update_config_keyspace_quotas(quotas);
            }

            if old_cpu.refill_interval != new_cpu.refill_interval {
                warn!(
                    "[CPU throttle] refill-interval change requires restart to take effect";
                    "old" => ?old_cpu.refill_interval,
                    "new" => ?new_cpu.refill_interval,
                );
            }
            if old_cpu.burst_window != new_cpu.burst_window {
                warn!(
                    "[CPU throttle] burst-window change requires restart to take effect";
                    "old" => ?old_cpu.burst_window,
                    "new" => ?new_cpu.burst_window,
                );
            }
        }
        let ratio = self.effective_max_read_cpu_ratio(new_config);
        self.max_cpu_time = self.max_cpu_time_for_ratio(new_config, ratio);
        self.config = new_config.clone()
    }

    pub fn on_usages(&mut self, usages: &Usages) {
        // Fair-gate registration can happen after manager injection. Refresh
        // here so usage ratios always use the ratio currently in force rather
        // than a startup-only cached denominator.
        self.refresh_max_cpu_time();
        match usages {
            Usages::Global { usages } => {
                let cpu_time = usages.cpu_time();

                let read_pool_cpu_ratio = cpu_time / self.max_cpu_time;

                // Warn if CPU ratio exceeds 1.0 (CPU time exceeded max expected)
                if read_pool_cpu_ratio > 1.0 {
                    tikv_util::warn!(
                        "Global CPU ratio exceeded 1.0";
                        "cpu_time" => cpu_time,
                        "max_cpu_time" => self.max_cpu_time,
                        "ratio" => read_pool_cpu_ratio,
                        "resource_type" => ?usages.resource_type,
                    );
                }

                // Save the global CPU ratio for later use in Keyspace events
                self.global_cpu_ratio = read_pool_cpu_ratio;
                self.cpu_severity =
                    Severity::get_severity(read_pool_cpu_ratio, self.severity_threshold);
                info!(
                    "read subscriber on_usages Global {:?} {:?} cpu_time {} max_cpu_time {} read_pool_cpu_ratio {} cpu_severity {:?}",
                    usages.resource_type,
                    usages.total(),
                    cpu_time,
                    self.max_cpu_time,
                    read_pool_cpu_ratio,
                    self.cpu_severity,
                );

                // Update CPU usage gauge metric
                CPU_USAGE_MONITOR_GLOBAL.set(read_pool_cpu_ratio);

                // Update CPU throttle manager with global CPU ratio
                if let Some(ref manager) = self.cpu_throttle_manager {
                    // Only update global ratio, preserve existing keyspace ratios
                    // The keyspace ratios will be updated when we receive Keyspace usages
                    manager.update_global_ratio(read_pool_cpu_ratio);
                    manager.adjust_refill_rates();
                }
            }
            Usages::Keyspace { usages } => {
                if self.config.debug {
                    info!(
                        "read subscriber on_usages Keyspace {:?} {:?} cpu_severity {:?}",
                        usages.resource_type,
                        usages.total(),
                        self.cpu_severity
                    );
                }

                match usages.resource_type {
                    ResourceType::Cpu {
                        cpu_type: CpuType::UnifiedReadPool,
                    } => self.on_keyspace_cpu_usages(usages),
                    ResourceType::Read => self.on_keyspace_read_usages(usages),
                    _ => {}
                }
            }
        };
    }

    fn on_keyspace_cpu_usages(&mut self, usages: &KeyspaceInstantUsages) {
        if let Some(ref manager) = self.cpu_throttle_manager {
            let mut keyspace_ratios = std::collections::HashMap::new();

            for (keyspace_id, keyspace_usage) in &usages.usages {
                let keyspace_cpu_time = keyspace_usage.cpu_time();
                let keyspace_cpu_ratio = if self.max_cpu_time > 0.0 {
                    keyspace_cpu_time / self.max_cpu_time
                } else {
                    0.0
                };

                if keyspace_cpu_ratio > 1.0 {
                    warn!(
                        "Keyspace CPU ratio exceeded 1.0";
                        "keyspace_id" => keyspace_id,
                        "cpu_time" => keyspace_cpu_time,
                        "max_cpu_time" => self.max_cpu_time,
                        "ratio" => keyspace_cpu_ratio,
                    );
                }

                keyspace_ratios.insert(*keyspace_id, keyspace_cpu_ratio);

                let keyspace_label = keyspace_label(*keyspace_id);
                CPU_USAGE_MONITOR_KEYSPACE
                    .with_label_values(&[keyspace_label.as_ref()])
                    .set(keyspace_cpu_ratio);
            }

            manager.update_usage(self.global_cpu_ratio, keyspace_ratios);
            manager.adjust_refill_rates();
        }
    }

    fn on_keyspace_read_usages(&mut self, usages: &KeyspaceInstantUsages) {
        if !self.read_limiter.is_legacy_read_limiter_allowed() {
            self.read_limiter.clear_all_limiter();
            return;
        }
        self.usages = usages.clone();

        if !self.config.get_read_enabled() || !self.config.enable_qps_bytes_limiter {
            self.read_limiter.clear_all_limiter();
            return;
        }

        if !self.cpu_severity.is_abnormal() || usages.total() == 0 {
            self.read_limiter.clear_all_limiter();
            return;
        }

        let bytes_usage = usages.bytes();
        let heavy_usage = bytes_usage / 100;
        let mut keyspace_usages = usages
            .usages
            .iter()
            .map(|(keyspace_id, usages)| (*keyspace_id, usages.bytes(), usages.qps()))
            .collect::<Vec<(u32, u64, usize)>>();
        keyspace_usages.retain(|(_, bytes, _)| *bytes > heavy_usage);
        keyspace_usages.sort_by(|a, b| b.1.cmp(&a.1).then(b.2.cmp(&a.2)));
        let max_read_bytes = self.max_read_bytes.max(bytes_usage) as f64;
        let active_quota_read_bytes = (max_read_bytes * self.config.active_quota_ratio) as u64;
        let min_read_bytes_quota = (max_read_bytes * self.config.min_quota_ratio) as u64;
        let mut active_bytes = 0u64;
        let mut active_pos = 0usize;
        for (keyspace_id, bytes, _) in keyspace_usages.iter() {
            let keyspace_label = keyspace_label(*keyspace_id);
            ACTIVE_KEYSPACE_READ_BYTES
                .with_label_values(&[keyspace_label.as_ref()])
                .set(*bytes as i64);
            active_bytes += *bytes;
            active_pos += 1;
            if active_bytes > active_quota_read_bytes {
                break;
            }
        }
        if self.config.debug {
            info!(
                "read subscriber active read quota, active_bytes:{} active_pos:{}",
                active_bytes, active_pos
            );
        }
        let now = Instant::now();
        let factor = match self.cpu_severity {
            Severity::Normal => {
                return;
            }
            Severity::Stressed => {
                let factor = 1.1 * self.config.severity_stressed_factor;
                if keyspace_usages.is_empty() {
                    return;
                }
                let (keyspace_id, _, qps) = keyspace_usages[0];
                self.read_limiter.update_limit(
                    keyspace_id,
                    Some(qps as f64 * factor),
                    Some(f64::INFINITY),
                    now,
                );
                info!(
                    "read subscriber active read quota, keyspace_id: {}, qps: {}",
                    keyspace_id, qps
                );
                return;
            }
            Severity::Critical => 0.9 * self.config.severity_critical_factor,
            Severity::Exhausted => 0.5 * self.config.severity_exhausted_factor,
        };

        if active_pos == 0 {
            info!(
                "read subscriber: no active keyspaces found for quota distribution";
                "cpu_severity" => ?self.cpu_severity,
            );
            return;
        }

        let keyspace_read_bytes_quota = active_bytes / active_pos as u64;
        let limited_keyspace_read_bytes_quota = (keyspace_read_bytes_quota as f64 * factor) as u64;
        let window_secs = self.config.window_size.as_secs_f64();
        for (idx, (keyspace_id, bytes, qps)) in keyspace_usages.iter().enumerate() {
            let (qps, target_bytes) = if idx < active_pos
                && *bytes > limited_keyspace_read_bytes_quota
            {
                let mut target_bytes = (*bytes as f64 * self.config.smoothing_factor
                    + keyspace_read_bytes_quota as f64 * (1.0 - self.config.smoothing_factor))
                    / window_secs;
                let qps = *qps as f64 * factor;
                target_bytes = target_bytes.max(min_read_bytes_quota as f64 / window_secs) * factor;
                info!(
                    "read subscriber active read quota, keyspace_id: {}, qps: {}, target_bytes: {}",
                    keyspace_id, qps, target_bytes
                );
                (qps, target_bytes)
            } else {
                (f64::INFINITY, f64::INFINITY)
            };
            self.read_limiter
                .update_limit(*keyspace_id, Some(qps), Some(target_bytes), now);
        }
    }
}

#[derive(Clone)]
pub struct ReadSubscriber {
    pub(crate) data: Arc<Mutex<ReadSubscriberData>>,
}

impl ReadSubscriber {
    pub fn new(read_limiter: ReadLimiter, config: Config) -> Self {
        let mut subscriber = ReadSubscriber {
            data: Arc::new(Mutex::new(ReadSubscriberData::new(
                read_limiter,
                config.clone(),
            ))),
        };
        subscriber.update_config(&config);
        subscriber
    }

    pub fn update_config(&mut self, config: &Config) {
        let mut data = self.data.lock().unwrap();
        data.update_config(config);
    }

    pub fn on_usages(&mut self, usages: &Usages) {
        let mut data = self.data.lock().unwrap();
        data.on_usages(usages);
    }

    pub fn set_cpu_throttle_manager(&mut self, manager: Arc<crate::CpuThrottleManager>) {
        let mut data = self.data.lock().unwrap();
        data.set_cpu_throttle_manager(manager);
    }
}

#[async_trait::async_trait]
impl ResourceSubscriber for ReadSubscriber {
    fn on_event(&mut self, event: &ResourceEvent) {
        match event {
            ResourceEvent::DispatchUsages(usages) => {
                self.on_usages(usages);
            }
            ResourceEvent::UpdateConfig(config) => self.update_config(config),
            _ => {}
        }
    }
}

pub struct TotalCpuSubscriberData {
    cpu_cores: f64,
    monitor: Arc<TotalCpuMonitor>,
}

impl TotalCpuSubscriberData {
    pub fn new(monitor: Arc<TotalCpuMonitor>, config: Config) -> Self {
        let mut subscriber = Self {
            cpu_cores: SysQuota::cpu_cores_quota().max(1.0),
            monitor,
        };
        subscriber.update_config(&config);
        subscriber
    }

    fn update_config(&mut self, config: &Config) {
        self.monitor
            .set_max_staleness(config.stats_interval.0.saturating_mul(2));
    }

    fn on_usages(&mut self, usages: &Usages) {
        let Usages::Global { usages } = usages else {
            return;
        };
        let interval = usages.usages.duration().as_secs_f64();
        if interval <= 0.0 {
            self.monitor.clear();
            return;
        }
        let total_cpu_ratio = usages.cpu_time() / (self.cpu_cores * interval);
        self.monitor.set_ratio(total_cpu_ratio);
    }
}

#[derive(Clone)]
pub struct TotalCpuSubscriber {
    pub(crate) data: Arc<Mutex<TotalCpuSubscriberData>>,
}

impl TotalCpuSubscriber {
    pub fn new(monitor: Arc<TotalCpuMonitor>, config: Config) -> Self {
        Self {
            data: Arc::new(Mutex::new(TotalCpuSubscriberData::new(monitor, config))),
        }
    }
}

#[async_trait::async_trait]
impl ResourceSubscriber for TotalCpuSubscriber {
    fn on_event(&mut self, event: &ResourceEvent) {
        match event {
            ResourceEvent::DispatchUsages(usages) => self.data.lock().unwrap().on_usages(usages),
            ResourceEvent::UpdateConfig(config) => {
                self.data.lock().unwrap().update_config(config);
            }
            _ => {}
        }
    }
}

pub struct TransferLeaderSubscriberData {
    pub(crate) config: Config,
    pub(crate) max_transfer_leader: u64,
    pub(crate) severity_threshold: SeverityThreshold,
    #[allow(dead_code)]
    pub(crate) transfer_leader_severity: Severity,
    pub(crate) usages: KeyspaceInstantUsages,
    pub(crate) transfer_leader_limiter: TransferLeaderLimiter,
}

impl TransferLeaderSubscriberData {
    pub fn new(transfer_leader_limiter: TransferLeaderLimiter, config: Config) -> Self {
        let mut subscriber = TransferLeaderSubscriberData {
            config: config.clone(),
            max_transfer_leader: Default::default(),
            severity_threshold: Default::default(),
            transfer_leader_severity: Default::default(),
            usages: Default::default(),
            transfer_leader_limiter,
        };
        subscriber.update_config(&config);
        subscriber
    }

    pub fn update_config(&mut self, new_config: &Config) {
        let window_secs = new_config.window_size.as_secs_f64();
        self.max_transfer_leader =
            (new_config.max_read_bytes_per_core_per_second.0 as f64 * window_secs) as u64;
        self.severity_threshold = SeverityThreshold::from_config(new_config);
        if self.config.get_transfer_leader_enabled() != new_config.get_transfer_leader_enabled() {
            self.transfer_leader_limiter
                .update_enabled(new_config.get_transfer_leader_enabled());
        }
        if self.config.limiter_stats_interval != new_config.limiter_stats_interval {
            self.transfer_leader_limiter
                .update_stats_interval(new_config.limiter_stats_interval.0);
        }
        if self.config.max_transfer_leader_per_second != new_config.max_transfer_leader_per_second {
            self.transfer_leader_limiter
                .update_limit(new_config.max_transfer_leader_per_second);
        }
        self.config = new_config.clone()
    }

    pub fn on_usages(&mut self, usages: &Usages) {
        // TODO: Collect metrics affected by transfer leader to dynamically adjust the
        // transfer leader speed.
        match usages {
            Usages::Global { usages } => {
                info!(
                    "transfer leader subscriber on_usages Global {:?} {:?}",
                    usages.resource_type,
                    usages.total(),
                );
            }
            Usages::Keyspace { usages } => {
                self.usages = usages.clone();
                info!(
                    "transfer leader subscriber on_usages Keyspace {:?} {:?}",
                    usages.resource_type,
                    usages.total(),
                );
            }
        };
        self.transfer_leader_limiter
            .update_limit(self.max_transfer_leader);
    }
}

#[derive(Clone)]
pub struct TransferLeaderSubscriber {
    pub(crate) data: Arc<Mutex<TransferLeaderSubscriberData>>,
}

impl TransferLeaderSubscriber {
    pub fn new(transfer_leader_limiter: TransferLeaderLimiter, config: Config) -> Self {
        let mut subscriber = TransferLeaderSubscriber {
            data: Arc::new(Mutex::new(TransferLeaderSubscriberData::new(
                transfer_leader_limiter,
                config.clone(),
            ))),
        };
        subscriber.update_config(&config);
        subscriber
    }

    pub fn update_config(&mut self, config: &Config) {
        let mut data = self.data.lock().unwrap();
        data.update_config(config);
    }

    pub fn on_usages(&mut self, usages: &Usages) {
        let mut data = self.data.lock().unwrap();
        data.on_usages(usages);
    }
}

#[async_trait::async_trait]
impl ResourceSubscriber for TransferLeaderSubscriber {
    fn on_event(&mut self, event: &ResourceEvent) {
        match event {
            ResourceEvent::DispatchUsages(usages) => {
                self.on_usages(usages);
            }
            ResourceEvent::UpdateConfig(config) => self.update_config(config),
            _ => {}
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, sync::Arc, time::Duration};

    use tikv_util::config::ReadableDuration;

    use super::*;
    use crate::{
        CpuThrottleManager, ResourceValue,
        metrics::CPU_THROTTLE_KEYSPACE_REFILL_RATE,
        usage::{GlobalInstantUsages, InstantUsage, InstantUsages},
    };

    fn build_keyspace_usages(
        resource_type: ResourceType,
        keyspace_id: u32,
        usage: ResourceValue,
    ) -> Usages {
        Usages::Keyspace {
            usages: KeyspaceInstantUsages {
                resource_type,
                usages: vec![(
                    keyspace_id,
                    InstantUsages {
                        usages: vec![InstantUsage {
                            instant: Instant::now(),
                            usage,
                        }],
                    },
                )],
            },
        }
    }

    fn build_global_cpu_usages(
        cpu_type: CpuType,
        first: (Instant, f64),
        second: (Instant, f64),
    ) -> Usages {
        Usages::Global {
            usages: GlobalInstantUsages {
                resource_type: ResourceType::Cpu { cpu_type },
                usages: InstantUsages {
                    usages: vec![
                        InstantUsage {
                            instant: first.0,
                            usage: ResourceValue::CpuTime { value: first.1 },
                        },
                        InstantUsage {
                            instant: second.0,
                            usage: ResourceValue::CpuTime { value: second.1 },
                        },
                    ],
                },
            },
        }
    }

    #[test]
    fn test_total_cpu_subscriber_updates_monitor_ratio() {
        let monitor = Arc::new(TotalCpuMonitor::default());
        let mut subscriber = TotalCpuSubscriber::new(monitor.clone(), Config::default());
        let start = Instant::now();
        let end = start + Duration::from_secs(2);
        let cpu_cores = SysQuota::cpu_cores_quota().max(1.0);
        let expected_ratio = 0.5;
        let cpu_delta = cpu_cores * 2.0 * expected_ratio;
        subscriber.on_event(&ResourceEvent::DispatchUsages(build_global_cpu_usages(
            CpuType::Total,
            (start, 10.0),
            (end, 10.0 + cpu_delta),
        )));
        let ratio = monitor.ratio().unwrap();
        assert!(
            (ratio - expected_ratio).abs() < 1e-6,
            "expected ratio {expected_ratio}, got {ratio}"
        );
    }

    #[test]
    fn test_total_cpu_monitor_returns_none_when_ratio_is_stale() {
        let monitor = TotalCpuMonitor::default();
        monitor.set_max_staleness(Duration::ZERO);
        monitor.set_ratio(0.42);
        std::thread::sleep(Duration::from_millis(1));
        assert!(monitor.ratio().is_none());
    }

    #[test]
    fn test_read_usages_do_not_clear_keyspace_cpu_ratio() {
        let keyspace_id = 991001;
        let keyspace_label = keyspace_label(keyspace_id);
        let keyspace_label = keyspace_label.as_ref();
        CPU_USAGE_MONITOR_KEYSPACE
            .with_label_values(&[keyspace_label])
            .set(0.0);
        CPU_THROTTLE_KEYSPACE_REFILL_RATE
            .with_label_values(&[keyspace_label])
            .set(0);

        let mut config = Config::default();
        config.window_size = ReadableDuration(Duration::from_secs(1));
        config.max_read_cpu_ratio = 1.0;
        config.enable_qps_bytes_limiter = false;

        let read_limiter = ReadLimiter::new(config.clone());
        let mut data = ReadSubscriberData::new(read_limiter, config);

        let mut cpu_throttle_config = CpuThrottleConfig::default();
        cpu_throttle_config.enable_dynamic_adjustment = true;
        cpu_throttle_config.high_watermark = 0.8;
        cpu_throttle_config.low_watermark = 0.5;

        let mut keyspace_cpu_us = HashMap::new();
        keyspace_cpu_us.insert(keyspace_id, 100_000);
        let manager = Arc::new(CpuThrottleManager::new(
            cpu_throttle_config,
            200_000,
            keyspace_cpu_us,
        ));
        data.set_cpu_throttle_manager(manager);

        let expected_ratio = 0.45;
        let cpu_time = data.max_cpu_time * expected_ratio;
        let cpu_usages = build_keyspace_usages(
            ResourceType::Cpu {
                cpu_type: CpuType::UnifiedReadPool,
            },
            keyspace_id,
            ResourceValue::CpuTime { value: cpu_time },
        );
        data.on_usages(&cpu_usages);

        let ratio_after_cpu = CPU_USAGE_MONITOR_KEYSPACE
            .with_label_values(&[keyspace_label])
            .get();
        assert!(
            (ratio_after_cpu - expected_ratio).abs() < 1e-6,
            "expected ratio {expected_ratio}, got {ratio_after_cpu}"
        );
        let refill_after_cpu = CPU_THROTTLE_KEYSPACE_REFILL_RATE
            .with_label_values(&[keyspace_label])
            .get();
        assert!(
            refill_after_cpu > 0 && refill_after_cpu < 100_000,
            "expected refill rate to be adjusted below base rate, got {refill_after_cpu}"
        );

        let read_usages = build_keyspace_usages(
            ResourceType::Read,
            keyspace_id,
            ResourceValue::Bytes { value: 4096 },
        );
        data.on_usages(&read_usages);

        let ratio_after_read = CPU_USAGE_MONITOR_KEYSPACE
            .with_label_values(&[keyspace_label])
            .get();
        assert!(
            (ratio_after_read - expected_ratio).abs() < 1e-6,
            "expected read usages not to clear cpu ratio, got {ratio_after_read}"
        );
        assert_eq!(data.usages.bytes(), 4096);
        assert_eq!(
            CPU_THROTTLE_KEYSPACE_REFILL_RATE
                .with_label_values(&[keyspace_label])
                .get(),
            refill_after_cpu
        );
    }

    #[test]
    fn test_legacy_read_budget_tracks_online_max_read_cpu_ratio_with_cpu_throttle() {
        let config = Config {
            enabled: true,
            max_read_cpu_ratio: 0.6,
            cpu_throttle: crate::config::CpuThrottleSettings {
                enabled: true,
                ..Default::default()
            },
            ..Default::default()
        };
        let read_limiter = ReadLimiter::new(config.clone());
        let manager = Arc::new(CpuThrottleManager::new(
            config.to_cpu_throttle_config(),
            config.global_cpu_us_per_interval(manager_refill_interval_ms(&config)),
            std::collections::HashMap::new(),
        ));
        read_limiter.set_cpu_throttle_manager(manager.clone());
        let mut data = ReadSubscriberData::new(read_limiter, config.clone());
        data.set_cpu_throttle_manager(manager);
        let initial_max_cpu_time = data.max_cpu_time;

        let updated = Config {
            max_read_cpu_ratio: 0.3,
            ..config
        };
        data.update_config(&updated);

        assert!((data.max_cpu_time - initial_max_cpu_time / 2.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_legacy_max_wait_time_applies_only_after_online_update() {
        let config = Config {
            max_read_wait_time: ReadableDuration::secs(2),
            ..Config::default()
        };
        let read_limiter = ReadLimiter::new(config.clone());
        let mut data = ReadSubscriberData::new(read_limiter.clone(), config.clone());
        assert_eq!(
            read_limiter.max_wait_time.load(),
            crate::limiter::MAX_WAIT_TIME
        );

        let updated = Config {
            max_read_wait_time: ReadableDuration::secs(3),
            ..config
        };
        data.update_config(&updated);

        assert_eq!(read_limiter.max_wait_time.load(), Duration::from_secs(3));
    }

    #[test]
    fn test_tenant_cap_hot_update_rewrites_config_quota_before_removal() {
        let keyspace_id = 991002;
        let mut old_config = Config::default();
        old_config.enabled = true;
        old_config.cpu_throttle.enabled = true;
        old_config.cpu_throttle.enforce_global_limit = true;
        old_config.cpu_throttle.enable_pd_service_limit = false;
        old_config.cpu_throttle.keyspace_cpu_quotas = format!("{keyspace_id}:1000");

        let read_limiter = ReadLimiter::new(old_config.clone());
        let mut data = ReadSubscriberData::new(read_limiter, old_config.clone());
        let manager = Arc::new(CpuThrottleManager::new(
            old_config.to_cpu_throttle_config(),
            200_000,
            old_config.to_cpu_throttle_config().keyspace_cpu_us(),
        ));
        data.set_cpu_throttle_manager(manager.clone());

        let mut new_config = old_config;
        new_config.cpu_throttle.keyspace_cpu_quotas.clear();
        new_config.cpu_throttle.tenant_max_cpu_ratio = 0.4;

        data.update_config(&new_config);

        let expected_quota = (SysQuota::cpu_cores_quota().max(1.0) * 0.4 * 100.0 * 1000.0) as u64;
        assert_eq!(
            manager.keyspace_available(keyspace_id),
            Some(expected_quota)
        );
    }

    fn manager_refill_interval_ms(config: &Config) -> u64 {
        config.to_cpu_throttle_config().refill_interval_ms
    }

    fn fair_ratio_test_data() -> (Arc<CpuThrottleManager>, ReadSubscriberData) {
        let cpu_throttle = crate::config::CpuThrottleSettings {
            enabled: true,
            ..Default::default()
        };
        let config = Config {
            enabled: true,
            max_read_cpu_ratio: 0.6,
            cpu_throttle,
            ..Config::default()
        };
        let refill_interval_ms = manager_refill_interval_ms(&config);
        let read_limiter = ReadLimiter::new_with_legacy_read_limiter(config.clone(), false);
        let manager = Arc::new(CpuThrottleManager::new(
            config.to_cpu_throttle_config(),
            config.global_cpu_us_per_interval(refill_interval_ms),
            std::collections::HashMap::new(),
        ));
        read_limiter.set_cpu_throttle_manager(manager.clone());
        let mut data = ReadSubscriberData::new(read_limiter, config.clone());
        data.set_cpu_throttle_manager(manager.clone());
        (manager, data)
    }

    #[test]
    fn test_usage_denominator_tracks_fair_max_read_cpu_ratio() {
        let (manager, mut data) = fair_ratio_test_data();
        let startup_max_cpu_time = data.max_cpu_time;
        manager.register_fair_global_cpu_gate();

        manager.update_fair_max_read_cpu_ratio(0.3);
        let start = Instant::now();
        let usages = build_global_cpu_usages(
            CpuType::UnifiedReadPool,
            (start, 0.0),
            (start + Duration::from_secs(1), 0.0),
        );
        data.on_usages(&usages);
        assert_eq!(manager.fair_charge_scale(), 2.0);
        assert!((data.max_cpu_time - startup_max_cpu_time / 2.0).abs() < f64::EPSILON);

        manager.update_fair_max_read_cpu_ratio(0.0);
        data.on_usages(&usages);
        assert_eq!(manager.fair_charge_scale(), 1.0);
        assert!((data.max_cpu_time - startup_max_cpu_time).abs() < f64::EPSILON);
    }

    #[test]
    fn test_usage_denominator_refreshes_after_late_fair_gate_registration() {
        let (manager, mut data) = fair_ratio_test_data();
        let startup_max_cpu_time = data.max_cpu_time;

        manager.update_fair_max_read_cpu_ratio(0.2);
        manager.register_fair_global_cpu_gate();
        let start = Instant::now();
        data.on_usages(&build_global_cpu_usages(
            CpuType::UnifiedReadPool,
            (start, 0.0),
            (start + Duration::from_secs(1), 0.0),
        ));

        let expected_max_cpu_time = startup_max_cpu_time / 3.0;
        assert!((data.max_cpu_time - expected_max_cpu_time).abs() / expected_max_cpu_time < 1e-12);
    }

    #[test]
    fn test_read_usages_clear_inactive_keyspace_limits() {
        let active_keyspace_id = 991011;
        let inactive_keyspace_id = 991012;

        let mut config = Config::default();
        config.enabled = true;
        config.read_enabled = true;
        config.window_size = ReadableDuration(Duration::from_secs(1));

        let read_limiter = ReadLimiter::new(config.clone());
        let mut data = ReadSubscriberData::new(read_limiter.clone(), config);
        data.cpu_severity = Severity::Critical;
        data.max_read_bytes = 100;

        let initial_read_usages = Usages::Keyspace {
            usages: KeyspaceInstantUsages {
                resource_type: ResourceType::Read,
                usages: vec![
                    (
                        active_keyspace_id,
                        InstantUsages {
                            usages: vec![InstantUsage {
                                instant: Instant::now(),
                                usage: ResourceValue::Bytes { value: 81 },
                            }],
                        },
                    ),
                    (
                        inactive_keyspace_id,
                        InstantUsages {
                            usages: vec![InstantUsage {
                                instant: Instant::now(),
                                usage: ResourceValue::Bytes { value: 79 },
                            }],
                        },
                    ),
                ],
            },
        };

        data.on_usages(&initial_read_usages);

        assert!(read_limiter.get_limiter(active_keyspace_id).is_some());
        assert!(read_limiter.get_limiter(inactive_keyspace_id).is_some());

        let read_usages = Usages::Keyspace {
            usages: KeyspaceInstantUsages {
                resource_type: ResourceType::Read,
                usages: vec![
                    (
                        active_keyspace_id,
                        InstantUsages {
                            usages: vec![InstantUsage {
                                instant: Instant::now(),
                                usage: ResourceValue::Bytes { value: 81 },
                            }],
                        },
                    ),
                    (
                        inactive_keyspace_id,
                        InstantUsages {
                            usages: vec![InstantUsage {
                                instant: Instant::now(),
                                usage: ResourceValue::Bytes { value: 19 },
                            }],
                        },
                    ),
                ],
            },
        };

        data.on_usages(&read_usages);

        assert!(read_limiter.get_limiter(active_keyspace_id).is_some());
        assert!(read_limiter.get_limiter(inactive_keyspace_id).is_none());
    }

    #[test]
    fn test_read_disabled_keeps_legacy_limiter_inactive() {
        let keyspace_id = 991013;
        let config = Config {
            enabled: true,
            read_enabled: false,
            enable_qps_bytes_limiter: true,
            window_size: ReadableDuration(Duration::from_secs(1)),
            ..Config::default()
        };
        let read_limiter = ReadLimiter::new(config.clone());
        let mut data = ReadSubscriberData::new(read_limiter.clone(), config);
        data.cpu_severity = Severity::Critical;
        data.max_read_bytes = 100;

        data.on_usages(&build_keyspace_usages(
            ResourceType::Read,
            keyspace_id,
            ResourceValue::Bytes { value: 81 },
        ));

        assert!(read_limiter.keyspace_limiters.is_empty());
    }

    #[test]
    fn test_fairness_gate_keeps_legacy_read_usage_state_and_metrics_empty() {
        let keyspace_id = 991014;
        let keyspace_label = keyspace_label(keyspace_id);
        let _ = ACTIVE_KEYSPACE_READ_BYTES.remove_label_values(&[keyspace_label.as_ref()]);
        let config = Config {
            enabled: true,
            read_enabled: true,
            enable_qps_bytes_limiter: true,
            window_size: ReadableDuration(Duration::from_secs(1)),
            ..Config::default()
        };
        let read_limiter = ReadLimiter::new_with_legacy_read_limiter(config.clone(), false);
        let mut data = ReadSubscriberData::new(read_limiter.clone(), config);
        data.cpu_severity = Severity::Critical;
        data.max_read_bytes = 100;

        data.on_usages(&build_keyspace_usages(
            ResourceType::Read,
            keyspace_id,
            ResourceValue::Bytes { value: 81 },
        ));

        assert!(data.usages.usages.is_empty());
        assert!(read_limiter.keyspace_limiters.is_empty());
        assert!(
            ACTIVE_KEYSPACE_READ_BYTES
                .remove_label_values(&[keyspace_label.as_ref()])
                .is_err()
        );
    }
}
