// Copyright 2025 PingCAP, Inc.

use std::{
    ops::{Add, Deref, Sub},
    sync::{
        Arc, RwLock,
        atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering, Ordering::Relaxed},
    },
    time::{Duration, Instant},
};

use dashmap::DashMap;
use tikv_util::{deadline::DeadlineError, error, info};

use crate::{
    ACTIVE_KEYSPACE_READ_BYTES, AtomicDuration, AtomicTime, Config, CpuThrottleConfig,
    CpuThrottleManager, ServiceLimitFetcher, ThrottleError, TimeUnit, metrics::keyspace_label,
};

const MIN_WAIT_TIME_INTERVAL: Duration = Duration::from_millis(1);

pub const MAX_WAIT_TIME: Duration = Duration::from_millis(50);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[allow(dead_code)]
pub enum Action {
    None,
    ScaleOut,
    Throttle,
    Reject,
}

#[derive(Clone)]
pub struct ReadLimiter {
    pub(crate) core: Arc<ReadLimiterCore>,
}

pub struct ReadLimiterCore {
    /// Whether the pre-fairness per-keyspace QPS/bytes limiter is allowed for
    /// this process. This is fixed at startup so an online resource-control
    /// update cannot reactivate the legacy path while read-pool fairness is on.
    legacy_read_limiter_allowed: bool,
    pub(crate) enabled: AtomicBool,
    pub(crate) dry_run: AtomicBool,
    startup_max_read_cpu_ratio: f64,
    pub(crate) timeout: AtomicDuration,
    pub(crate) stats_interval: AtomicDuration,
    pub(crate) max_wait_time: AtomicDuration,
    pub(crate) keyspace_limiters: DashMap<u32, (Instant, KeyspaceReadLimiter)>,
    pub(crate) cpu_throttle_manager: RwLock<Option<Arc<CpuThrottleManager>>>,
    pub(crate) service_limit_fetcher: RwLock<Option<Arc<ServiceLimitFetcher>>>,
}

impl Deref for ReadLimiter {
    type Target = ReadLimiterCore;

    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

impl ReadLimiterCore {
    pub fn new(config: Config) -> Self {
        Self::new_with_legacy_read_limiter(config, true)
    }

    pub fn new_with_legacy_read_limiter(config: Config, legacy_read_limiter_allowed: bool) -> Self {
        let timeout = config.limiter_timeout.0;
        let stats_interval = config.limiter_stats_interval.0;
        Self {
            legacy_read_limiter_allowed,
            enabled: AtomicBool::from(config.get_read_enabled()),
            dry_run: AtomicBool::from(config.dry_run),
            startup_max_read_cpu_ratio: config.max_read_cpu_ratio,
            timeout: AtomicDuration::new(timeout, TimeUnit::Millisecond),
            stats_interval: AtomicDuration::new(stats_interval, TimeUnit::Millisecond),
            // Preserve the pre-fairness startup behavior. The legacy config is
            // applied only by a later online update.
            max_wait_time: AtomicDuration::new(MAX_WAIT_TIME, TimeUnit::Millisecond),
            keyspace_limiters: DashMap::new(),
            cpu_throttle_manager: RwLock::new(None),
            service_limit_fetcher: RwLock::new(None),
        }
    }
}

impl ReadLimiter {
    pub fn new(config: Config) -> Self {
        Self::new_with_legacy_read_limiter(config, true)
    }

    /// Creates a read limiter with a startup-fixed gate for the legacy
    /// per-keyspace QPS/bytes scan limiter.
    pub fn new_with_legacy_read_limiter(config: Config, legacy_read_limiter_allowed: bool) -> Self {
        Self {
            core: Arc::new(ReadLimiterCore::new_with_legacy_read_limiter(
                config,
                legacy_read_limiter_allowed,
            )),
        }
    }

    pub(crate) fn update_enabled(&self, enabled: bool) {
        self.enabled.store(enabled, Relaxed);
        if !enabled {
            self.clear_all_limiter();
        }
    }

    pub fn is_legacy_read_limiter_enabled(&self) -> bool {
        self.legacy_read_limiter_allowed && self.enabled.load(Relaxed)
    }

    /// Returns the startup-fixed construction gate for the legacy read
    /// limiter, independent of its online enable switch.
    pub fn is_legacy_read_limiter_allowed(&self) -> bool {
        self.legacy_read_limiter_allowed
    }

    pub(crate) fn update_dry_run(&self, dry_run: bool) {
        self.dry_run.store(dry_run, Relaxed);
    }

    /// Returns whether resource control is running without enforcement.
    pub fn is_dry_run(&self) -> bool {
        self.dry_run.load(Relaxed)
    }

    pub(crate) fn update_timeout(&self, timeout: Duration) {
        self.timeout.store(timeout);
    }

    pub(crate) fn update_stats_interval(&self, stats_interval: Duration) {
        self.stats_interval.store(stats_interval);
    }

    pub(crate) fn update_max_wait_time(&self, max_wait_time: Duration) {
        self.max_wait_time.store(max_wait_time);
    }

    pub(crate) fn clear_all_limiter(&self) {
        if self.keyspace_limiters.is_empty() {
            return;
        }
        self.keyspace_limiters.iter().for_each(|keyspace_limiter_ref|{
            let keyspace_id = *keyspace_limiter_ref.key();
            let keyspace_label = keyspace_label(keyspace_id);
            let keyspace_str = keyspace_label.as_ref();
            let _ =  ACTIVE_KEYSPACE_READ_BYTES.remove_label_values(&[keyspace_str]).map_err(
                |err| error!("failed to remove active keyspace read bytes metric"; "keyspace_id" => keyspace_str, "err" => %err),
            );
        });
        self.keyspace_limiters.clear();
    }

    pub fn set_cpu_throttle_manager(&self, manager: Arc<CpuThrottleManager>) {
        info!("set cpu throttle manager");
        manager.initialize_fair_max_read_cpu_ratio(self.startup_max_read_cpu_ratio);
        let mut guard = self.cpu_throttle_manager.write().unwrap();
        *guard = Some(manager);
    }

    pub fn get_cpu_throttle_manager(&self) -> Option<Arc<CpuThrottleManager>> {
        self.cpu_throttle_manager.read().unwrap().clone()
    }

    pub fn is_cpu_throttle_enabled(&self) -> bool {
        self.cpu_throttle_manager.read().unwrap().is_some()
    }

    pub fn is_low_priority_dag_cpu_throttle_enabled(&self) -> bool {
        self.cpu_throttle_manager
            .read()
            .unwrap()
            .as_ref()
            .map(|manager| manager.is_low_priority_dag_cpu_throttle_enabled())
            .unwrap_or(false)
    }

    pub fn cpu_throttle_refill_interval(&self) -> Option<Duration> {
        self.cpu_throttle_manager
            .read()
            .unwrap()
            .as_ref()
            .map(|manager| Duration::from_millis(manager.refill_interval_ms().max(1)))
    }

    /// Returns the exact refill deadline relevant to a failed CPU allocation.
    pub fn cpu_throttle_allocation_retry_at(
        &self,
        keyspace_id: u32,
        error: &ThrottleError,
    ) -> Option<Instant> {
        self.cpu_throttle_manager
            .read()
            .unwrap()
            .as_ref()
            .map(|manager| manager.allocation_retry_at(keyspace_id, error))
    }

    pub fn set_service_limit_fetcher(&self, fetcher: Arc<ServiceLimitFetcher>) {
        let mut guard = self.service_limit_fetcher.write().unwrap();
        *guard = Some(fetcher);
    }

    /// Hot-update the resident service_limit fetcher's configuration, if one is
    /// installed. Used by config subscription so a fetcher that started dormant
    /// (e.g. boot in tenant cap mode) picks up the current PD config before it
    /// is allowed to apply service limits.
    pub fn reconfigure_service_limit_fetcher(&self, config: &CpuThrottleConfig) {
        if let Some(ref fetcher) = *self.service_limit_fetcher.read().unwrap() {
            fetcher.reconfigure(config);
        }
    }

    pub fn register_keyspace_for_service_limit(&self, keyspace_id: u32) {
        if self
            .cpu_throttle_manager
            .read()
            .unwrap()
            .as_ref()
            .map(|manager| manager.tenant_cap_enabled())
            .unwrap_or(false)
        {
            return;
        }
        if let Some(ref fetcher) = *self.service_limit_fetcher.read().unwrap() {
            fetcher.register_keyspace(keyspace_id);
        }
    }

    pub fn try_allocate_cpu(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
        request_deadline: Instant,
    ) -> Result<crate::CpuTokenHandle, ThrottleError> {
        let guard = self.cpu_throttle_manager.read().unwrap();
        if let Some(ref manager) = *guard {
            manager.try_allocate(keyspace_id, estimated_cpu_us, request_deadline)
        } else {
            // CPU throttling not enabled - return a no-op token handle
            Ok(crate::CpuTokenHandle {
                keyspace_id,
                allocated_us: std::sync::atomic::AtomicU64::new(estimated_cpu_us),
                keyspace_allocated_us: std::sync::atomic::AtomicU64::new(0),
                global_allocated_us: std::sync::atomic::AtomicU64::new(0),
                actual_used_us: std::sync::atomic::AtomicU64::new(0),
                global_bucket: Arc::new(crate::CpuTokenBucket::new(
                    u64::MAX,
                    u64::MAX,
                    Duration::from_millis(100),
                    0, // interval_allocation_limit (0 = disabled)
                )),
                keyspace_bucket: None,
                low_priority_bucket: None,
                keyspace_inflight_state: None,
                low_priority_allocated_us: std::sync::atomic::AtomicU64::new(0),
                low_priority_inflight_allocated_us: None,
                is_burst: std::sync::atomic::AtomicBool::new(false),
                request_deadline,
                skip_throttle: true,
                keyspace_burst_allowed: true,
                cpu_usage_monitor: None,
                config: None,
                adaptive_estimator: None,
                burst_usage_token: None,
                burst_usage_tracker: None,
                last_dynamic_burst_log: Some(Arc::new(RwLock::new(Instant::now()))),
                last_periodic_allocation_log: Some(Arc::new(RwLock::new(Instant::now()))),
            })
        }
    }

    pub fn try_allocate_low_priority_dag_cpu(
        &self,
        keyspace_id: u32,
        estimated_cpu_us: u64,
        request_deadline: Instant,
    ) -> Result<crate::CpuTokenHandle, ThrottleError> {
        let guard = self.cpu_throttle_manager.read().unwrap();
        if let Some(ref manager) = *guard {
            manager.try_allocate_low_priority_dag(keyspace_id, estimated_cpu_us, request_deadline)
        } else {
            self.try_allocate_cpu(keyspace_id, estimated_cpu_us, request_deadline)
        }
    }

    /// Get the estimated CPU time per request for a keyspace
    /// Returns the configured value if CPU throttling is enabled, otherwise
    /// returns default value
    pub fn get_estimated_cpu_per_request_us(&self, keyspace_id: u32) -> u64 {
        let guard = self.cpu_throttle_manager.read().unwrap();
        if let Some(ref manager) = *guard {
            manager.get_estimated_cpu_per_request_us(keyspace_id)
        } else {
            // Return default value when CPU throttling is not enabled
            1000 // 1ms default
        }
    }

    /// Get the estimated CPU time per low-priority DAG request for a keyspace.
    pub fn get_estimated_cpu_per_low_priority_dag_request_us(&self, keyspace_id: u32) -> u64 {
        let guard = self.cpu_throttle_manager.read().unwrap();
        if let Some(ref manager) = *guard {
            manager.get_estimated_cpu_per_low_priority_dag_request_us(keyspace_id)
        } else {
            1000
        }
    }

    pub(crate) fn update_limit(
        &self,
        keyspace_id: u32,
        req_speed_limit: Option<f64>,
        bytes_speed_limit: Option<f64>,
        instant: Instant,
    ) {
        if !self.is_legacy_read_limiter_enabled() {
            return;
        }
        if req_speed_limit.unwrap_or_default().is_infinite()
            && bytes_speed_limit.unwrap_or_default().is_infinite()
        {
            self.remove_limiter(keyspace_id);
            return;
        }
        self.keyspace_limiters
            .entry(keyspace_id)
            .and_modify(|(ts, keyspace_read_limiter)| {
                *ts = instant;
                keyspace_read_limiter.update_stats_interval(self.stats_interval.load());
                keyspace_read_limiter.update_max_wait_time(self.max_wait_time.load());
                keyspace_read_limiter.set_speed_limit(req_speed_limit, bytes_speed_limit);
            })
            .or_insert_with(|| {
                let keyspace_read_limiter = KeyspaceReadLimiter::new(
                    keyspace_id,
                    self.stats_interval.load(),
                    self.max_wait_time.load(),
                );
                keyspace_read_limiter.set_speed_limit(req_speed_limit, bytes_speed_limit);
                (instant, keyspace_read_limiter)
            });
    }

    pub fn get_limiter(&self, keyspace_id: u32) -> Option<KeyspaceReadLimiter> {
        if !self.is_legacy_read_limiter_enabled() {
            return None;
        }
        let timeout = self.timeout.load();
        if let Some((ts, keyspace_read_limiter)) =
            self.keyspace_limiters.get(&keyspace_id).map(|v| v.clone())
        {
            if ts.elapsed() < timeout {
                Some(keyspace_read_limiter.clone())
            } else {
                self.remove_limiter(keyspace_id);
                None
            }
        } else {
            None
        }
    }

    fn remove_limiter(&self, keyspace_id: u32) {
        self.keyspace_limiters.remove(&keyspace_id);
        let keyspace_label = keyspace_label(keyspace_id);
        let keyspace_str = keyspace_label.as_ref();
        let _ =  ACTIVE_KEYSPACE_READ_BYTES.remove_label_values(&[keyspace_str]).map_err(
            |err| error!("failed to remove active keyspace read bytes metric"; "keyspace_id" => keyspace_str, "err" => %err),
        );
    }
}

#[derive(Clone)]
pub struct KeyspaceReadLimiter {
    core: Arc<KeyspaceReadLimiterCore>,
}

pub struct KeyspaceReadLimiterCore {
    keyspace_id: u32,
    req_limiter: tikv_util::time::Limiter,
    bytes_limiter: tikv_util::time::Limiter,
    waiting_cnt: AtomicI64,
    req_allowed_time: AtomicTime, // Requests after this time do not need to wait.
    bytes_allowed_time: AtomicTime, // Bytes after this time do not need to wait.
    last_time: AtomicTime,
    stats_interval: AtomicDuration,
    max_wait_time: AtomicDuration,
}

impl Deref for KeyspaceReadLimiter {
    type Target = KeyspaceReadLimiterCore;

    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

impl KeyspaceReadLimiterCore {
    pub fn new(keyspace_id: u32, stats_interval: Duration, max_wait_time: Duration) -> Self {
        let req_limiter = <tikv_util::time::Limiter>::builder(f64::INFINITY).build();
        let bytes_limiter = <tikv_util::time::Limiter>::builder(f64::INFINITY).build();
        let start = Instant::now();
        Self {
            keyspace_id,
            req_limiter,
            bytes_limiter,
            waiting_cnt: AtomicI64::default(),
            req_allowed_time: AtomicTime::new(start, TimeUnit::Microsecond),
            bytes_allowed_time: AtomicTime::new(start, TimeUnit::Microsecond),
            last_time: AtomicTime::new(start, TimeUnit::Millisecond),
            stats_interval: AtomicDuration::new(stats_interval, TimeUnit::Millisecond),
            max_wait_time: AtomicDuration::new(max_wait_time, TimeUnit::Millisecond),
        }
    }
}

impl Default for KeyspaceReadLimiter {
    fn default() -> Self {
        Self::new(0, Duration::from_secs(1), MAX_WAIT_TIME)
    }
}

impl KeyspaceReadLimiter {
    pub fn new(keyspace_id: u32, stats_interval: Duration, max_wait_time: Duration) -> Self {
        Self {
            core: Arc::new(KeyspaceReadLimiterCore::new(
                keyspace_id,
                stats_interval,
                max_wait_time,
            )),
        }
    }

    fn update_allowed_time<F>(&self, allowed_time: &AtomicTime, update: F) -> Duration
    where
        F: Fn(Instant) -> (Option<Instant>, Duration),
    {
        let mut current_time = allowed_time.load();
        loop {
            let (new_allowed_time, dur) = update(current_time);
            let Some(new_allowed_time) = new_allowed_time else {
                return dur;
            };
            match allowed_time.compare_exchange(current_time, new_allowed_time) {
                Some(current) => {
                    current_time = current;
                }
                None => {
                    return dur;
                }
            }
        }
    }

    fn max_allowed_time(&self) -> Instant {
        let req_allowed = self.req_allowed_time.load();
        let bytes_allowed = self.bytes_allowed_time.load();
        if req_allowed > bytes_allowed {
            req_allowed
        } else {
            bytes_allowed
        }
    }

    fn update_max_allowed_time<F>(&self, update: F) -> Duration
    where
        F: Fn(Instant) -> (Option<Instant>, Duration),
    {
        loop {
            let req_allowed = self.req_allowed_time.load();
            let bytes_allowed = self.bytes_allowed_time.load();
            let (max_allowed, update_req, update_bytes) = match req_allowed.cmp(&bytes_allowed) {
                std::cmp::Ordering::Greater => (req_allowed, true, false),
                std::cmp::Ordering::Less => (bytes_allowed, false, true),
                std::cmp::Ordering::Equal => (req_allowed, true, true),
            };
            let (new_allowed_time, dur) = update(max_allowed);
            let Some(new_allowed_time) = new_allowed_time else {
                return dur;
            };
            let mut updated = true;
            if update_req
                && self
                    .req_allowed_time
                    .compare_exchange(req_allowed, new_allowed_time)
                    .is_some()
            {
                updated = false;
            }
            if update_bytes
                && self
                    .bytes_allowed_time
                    .compare_exchange(bytes_allowed, new_allowed_time)
                    .is_some()
            {
                updated = false;
            }
            if updated {
                return dur;
            }
        }
    }

    fn update_recovery_time(&self, allowed_time: &AtomicTime, recovery_duration: Duration) {
        // consume_duration returns the current recovery horizon, not an incremental
        // debt. Convert it to an absolute time once so CAS retries do not
        // re-anchor the same consumption to a later clock sample.
        let now = Instant::now();
        let recovery_time = now.add(recovery_duration);
        let update = |allowed_time: Instant| {
            if recovery_time <= allowed_time || recovery_time <= now {
                return (None, Duration::default());
            }
            let base_time = if allowed_time > now {
                allowed_time
            } else {
                now
            };
            (
                Some(recovery_time),
                recovery_time.saturating_duration_since(base_time),
            )
        };
        self.update_allowed_time(allowed_time, update);
    }

    pub fn take_wait_time(&self) -> Duration {
        let now = Instant::now();
        let update = |allowed_time: Instant| {
            let dur = allowed_time.duration_since(now);
            if dur < MIN_WAIT_TIME_INTERVAL {
                return (None, dur);
            }
            let new_allowed_time = now;
            (Some(new_allowed_time), dur)
        };
        let req_dur = self.update_allowed_time(&self.req_allowed_time, update);
        let bytes_dur = self.update_allowed_time(&self.bytes_allowed_time, update);
        req_dur.max(bytes_dur)
    }

    pub fn wait_time(&self) -> Duration {
        let update = |allowed_time: Instant| {
            let now = Instant::now();
            let dur = allowed_time.saturating_duration_since(now);
            if dur < MIN_WAIT_TIME_INTERVAL {
                return (None, dur);
            };
            let new_allowed_time = allowed_time.sub(MIN_WAIT_TIME_INTERVAL);
            (Some(new_allowed_time), MIN_WAIT_TIME_INTERVAL)
        };
        self.update_max_allowed_time(update)
    }

    pub async fn wait(&self) -> Duration {
        self.waiting_cnt.fetch_add(1, Relaxed);
        let mut wait_time = Duration::default();
        loop {
            let dur = self.wait_time();
            if dur.is_zero() {
                break;
            }
            tokio::time::sleep(dur).await;
            wait_time = wait_time.add(dur);
            let avg_wait_time = Duration::from_secs_f64(
                self.max_allowed_time()
                    .duration_since(Instant::now())
                    .as_secs_f64()
                    / self.waiting_cnt.load(Relaxed) as f64,
            );
            if wait_time >= avg_wait_time.min(self.max_wait_time.load()) {
                break;
            }
        }
        self.waiting_cnt.fetch_add(-1, Relaxed);
        wait_time
    }

    /// Waits until allowed time or returns an error if the deadline is
    /// exceeded.
    pub async fn wait_with_deadline(&self, deadline: Instant) -> Result<Duration, DeadlineError> {
        self.waiting_cnt.fetch_add(1, Relaxed);
        let mut wait_time = Duration::default();
        loop {
            let dur = self
                .max_allowed_time()
                .saturating_duration_since(Instant::now());
            if dur.is_zero() {
                break;
            }
            let remaining = deadline
                .checked_duration_since(Instant::now())
                .ok_or(DeadlineError)?;
            if remaining < dur {
                self.waiting_cnt.fetch_add(-1, Relaxed);
                return Err(DeadlineError);
            }
            if dur >= MIN_WAIT_TIME_INTERVAL {
                let update = |allowed_time: Instant| {
                    let new_allowed_time = allowed_time.sub(MIN_WAIT_TIME_INTERVAL);
                    (Some(new_allowed_time), MIN_WAIT_TIME_INTERVAL)
                };
                self.update_max_allowed_time(update);
                tokio::time::sleep(MIN_WAIT_TIME_INTERVAL).await;
                wait_time = wait_time.add(MIN_WAIT_TIME_INTERVAL);
            } else {
                tokio::time::sleep(dur).await;
                wait_time = wait_time.add(dur);
            }
            let avg_wait_time = Duration::from_secs_f64(
                self.max_allowed_time()
                    .duration_since(Instant::now())
                    .as_secs_f64()
                    / self.waiting_cnt.load(Relaxed) as f64,
            );
            if wait_time >= avg_wait_time.min(self.max_wait_time.load()) {
                break;
            }
        }
        self.waiting_cnt.fetch_add(-1, Relaxed);
        Ok(wait_time)
    }

    /// Consumes one request unit for QPS limiting.
    pub fn consume_request(&self) {
        let dur = self.req_limiter.consume_duration(1);
        self.update_recovery_time(&self.req_allowed_time, dur);
    }

    /// Consumes bytes for throughput limiting.
    pub fn consume_bytes(&self, bytes: usize) {
        let dur = self.bytes_limiter.consume_duration(bytes);
        self.update_recovery_time(&self.bytes_allowed_time, dur);
    }

    /// Returns one request unit to the QPS limiter.
    pub fn unconsume_request(&self) {
        self.req_limiter.unconsume(1);
    }

    /// Returns bytes to the throughput limiter.
    pub fn unconsume_bytes(&self, bytes: usize) {
        self.bytes_limiter.unconsume(bytes);
    }

    pub fn is_unlimited(&self) -> bool {
        self.req_limiter.speed_limit() == f64::INFINITY
            && self.bytes_limiter.speed_limit() == f64::INFINITY
    }

    pub fn total_consumed(&self) -> (usize /* req */, usize /* bytes */) {
        (
            self.req_limiter.total_bytes_consumed(),
            self.bytes_limiter.total_bytes_consumed(),
        )
    }

    pub fn speed_limit(&self) -> (f64 /* req */, f64 /* bytes */) {
        (
            self.req_limiter.speed_limit(),
            self.bytes_limiter.speed_limit(),
        )
    }

    pub fn set_speed_limit(&self, req_speed_limit: Option<f64>, bytes_speed_limit: Option<f64>) {
        if let Some(req_speed_limit) = req_speed_limit {
            if req_speed_limit > 0.0 {
                self.req_limiter.set_speed_limit(req_speed_limit);
            } else {
                self.req_limiter.set_speed_limit(f64::INFINITY);
            }
        }
        if let Some(bytes_speed_limit) = bytes_speed_limit {
            if bytes_speed_limit > 0.0 {
                self.bytes_limiter.set_speed_limit(bytes_speed_limit);
            } else {
                self.bytes_limiter.set_speed_limit(f64::INFINITY);
            }
        }
        self.update_statistics();
    }

    fn update_statistics(&self) {
        let last_time = self.last_time.load();
        let dur = last_time.elapsed();
        if dur < self.stats_interval.load() {
            return;
        }
        self.last_time.store(Instant::now());
        let total_requests = self.req_limiter.total_bytes_consumed() as f64;
        let total_bytes = self.bytes_limiter.total_bytes_consumed() as f64;
        self.req_limiter.reset_statistics();
        self.bytes_limiter.reset_statistics();
        if total_requests == 0.0 {
            return;
        }
        let qps = total_requests / dur.as_secs_f64();
        let bytes_rate = total_bytes / dur.as_secs_f64();
        let bytes_per_req = total_bytes / total_requests;
        info!("resource control update_statistics dur {:?}", dur;
            "keyspace_id" => self.keyspace_id,
            "total_requests" => total_requests,
            "total_bytes" => total_bytes,
            "qps" => qps,
            "bytes_rate" => bytes_rate,
            "bytes_per_req" => bytes_per_req,
        );
    }

    fn update_stats_interval(&self, stats_interval: Duration) {
        self.stats_interval.store(stats_interval);
    }

    fn update_max_wait_time(&self, max_wait_time: Duration) {
        self.max_wait_time.store(max_wait_time);
    }
}

#[derive(Clone)]
pub struct TransferLeaderLimiter {
    pub(crate) core: Arc<TransferLeaderLimiterCore>,
}

pub struct TransferLeaderLimiterCore {
    pub(crate) enabled: AtomicBool,
    rate_limiter: UniformLimiter,
    last_time: AtomicTime,
    stats_interval: AtomicDuration,
    last_rate: AtomicU64,
}

impl Deref for TransferLeaderLimiter {
    type Target = TransferLeaderLimiterCore;

    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

impl TransferLeaderLimiterCore {
    pub fn new(config: Config) -> Self {
        let rate_limiter = UniformLimiter::new(config.max_transfer_leader_per_second);
        Self {
            enabled: AtomicBool::from(config.get_transfer_leader_enabled()),
            last_time: AtomicTime::new(Instant::now(), TimeUnit::Millisecond),
            stats_interval: AtomicDuration::new(
                config.limiter_stats_interval.0,
                TimeUnit::Millisecond,
            ),
            rate_limiter,
            last_rate: AtomicU64::default(),
        }
    }
}

impl TransferLeaderLimiter {
    pub fn new(config: Config) -> Self {
        Self {
            core: Arc::new(TransferLeaderLimiterCore::new(config)),
        }
    }

    pub(crate) fn update_enabled(&self, enabled: bool) {
        self.enabled.store(enabled, Relaxed);
    }

    pub(crate) fn get_enabled(&self) -> bool {
        self.enabled.load(Relaxed)
    }

    pub(crate) fn update_stats_interval(&self, stats_interval: Duration) {
        self.stats_interval.store(stats_interval);
    }

    pub(crate) fn update_limit(&self, speed_limit: u64) {
        self.rate_limiter.set_speed_limit(speed_limit);
        self.update_statistics();
    }

    pub fn allow(&self) -> bool {
        if !self.get_enabled() {
            return true;
        }
        let allow = self.rate_limiter.allow(1);

        self.update_statistics(); // TODO: Remove this line after implementing dynamic speed adjustment.
        allow
    }

    fn update_statistics(&self) {
        let last_time = self.last_time.load();
        let dur = last_time.elapsed();
        if dur < self.stats_interval.load() {
            return;
        }
        self.last_time.store(Instant::now());
        let transfer_leader = self.rate_limiter.total_bytes_consumed() as f64;
        self.rate_limiter.reset_statistics();
        let rate = transfer_leader / dur.as_secs_f64();
        info!("resource control update_statistics dur {:?}", dur;
            "transfer_leader" => transfer_leader,
            "rate" => rate,
        );
        self.last_rate.store(rate as u64, Relaxed)
    }

    pub fn rate(&self) -> u64 {
        self.last_rate.load(Relaxed)
    }
}

/// Lock-free, non-blocking, uniform rate limiter.
/// Ensures smooth execution by rejecting requests if they arrive too early.
///
/// - rate_bytes_per_sec: throughput in bytes/sec
pub struct UniformLimiter {
    rate_bytes_per_sec: AtomicU64,
    allowed_time: AtomicTime,
    total_bytes_consumed: AtomicU64,
}

impl UniformLimiter {
    pub fn new(rate_bytes_per_sec: u64) -> Self {
        assert!(rate_bytes_per_sec > 0);
        let base = Instant::now();
        Self {
            rate_bytes_per_sec: AtomicU64::new(rate_bytes_per_sec),
            allowed_time: AtomicTime::new(base, TimeUnit::Nanosecond),
            total_bytes_consumed: AtomicU64::default(),
        }
    }

    fn bytes_to_dur(rate_bytes_per_sec: u64, bytes: u64) -> Duration {
        let secs: f64 = if rate_bytes_per_sec == 0 || bytes == 0 {
            0.0
        } else {
            bytes as f64 / rate_bytes_per_sec as f64
        };
        Duration::from_secs_f64(secs)
    }

    /// Dynamically changes the speed limit. The new limit applies to all clones
    /// of this instance.
    pub fn set_speed_limit(&self, speed_limit: u64) {
        debug_assert!(speed_limit > 0, "speed limit must be positive");
        self.rate_bytes_per_sec.store(speed_limit, Relaxed);
    }

    /// Non-blocking check. Returns true if request can proceed immediately,
    /// false if it arrives too early (would break uniform pacing).
    pub fn allow(&self, bytes: u64) -> bool {
        if bytes == 0 {
            return true;
        }
        self.total_bytes_consumed
            .fetch_add(bytes, Ordering::Relaxed);
        let now = Instant::now();
        let duration = UniformLimiter::bytes_to_dur(self.rate_bytes_per_sec.load(Relaxed), bytes);
        let mut allowed = self.allowed_time.load();
        loop {
            if now < allowed {
                // Too early, must reject to maintain smoothness
                return false;
            }
            let scheduled_end = allowed.add(duration).max(now.sub(duration));
            match self.allowed_time.compare_exchange(allowed, scheduled_end) {
                Some(current) => {
                    allowed = current;
                }
                None => {
                    return true;
                }
            }
        }
    }

    /// Obtains the total number of bytes consumed by this limiter so far.
    ///
    /// If more than `usize::MAX` bytes have been consumed, the count will wrap
    /// around.
    pub fn total_bytes_consumed(&self) -> usize {
        self.total_bytes_consumed.load(Ordering::Relaxed) as usize
    }

    /// Resets the total number of bytes consumed to 0.
    pub fn reset_statistics(&self) {
        self.total_bytes_consumed.store(0, Ordering::Relaxed);
    }
}

#[cfg(test)]
mod tests {
    use std::{thread::sleep, time::Instant};

    use super::*;

    #[test]
    fn test_uniform_limiter() {
        let limiter = UniformLimiter::new(1000); // ~1000 bytes/sec

        // First request should pass.
        assert!(limiter.allow(100));

        // Immediate second request likely rejected (too soon).
        assert!(!limiter.allow(100));

        // After enough time passes, allow again.
        sleep(Duration::from_millis(100));
        assert!(limiter.allow(100));

        sleep(Duration::from_millis(100));
        assert!(limiter.allow(200));

        // Request rejected due to insufficient wait time.
        sleep(Duration::from_millis(100));
        assert!(!limiter.allow(100));

        // After enough time passes, allow again.
        sleep(Duration::from_millis(100));
        assert!(limiter.allow(100));

        assert!(!limiter.allow(100));
        sleep(Duration::from_millis(100));

        for i in 10..30 {
            assert!(limiter.allow(i));
            sleep(Duration::from_millis(i));
        }
    }

    #[test]
    fn test_keyspace_read_limiter_consume_separate() {
        let limiter = KeyspaceReadLimiter::new(42, Duration::from_secs(1), MAX_WAIT_TIME);
        limiter.set_speed_limit(Some(1.0), Some(10.0));
        limiter.consume_request();
        limiter.consume_bytes(100);
        assert_eq!(limiter.total_consumed(), (1, 100));
    }

    #[test]
    fn test_read_limiter_starts_with_legacy_max_wait_time() {
        let config = Config {
            max_read_wait_time: tikv_util::config::ReadableDuration::secs(2),
            ..Config::default()
        };
        let limiter = ReadLimiter::new(config);

        assert_eq!(limiter.max_wait_time.load(), MAX_WAIT_TIME);
    }

    #[test]
    fn test_fairness_gate_cannot_be_bypassed_by_legacy_read_limiter_update() {
        let config = Config {
            enabled: true,
            read_enabled: true,
            ..Config::default()
        };
        let limiter = ReadLimiter::new_with_legacy_read_limiter(config, false);

        limiter.update_limit(42, Some(1.0), Some(1024.0), Instant::now());
        assert!(!limiter.is_legacy_read_limiter_enabled());
        assert!(limiter.get_limiter(42).is_none());
        assert!(limiter.keyspace_limiters.is_empty());

        // Model a later online update trying to turn the legacy switch on.
        limiter.update_enabled(true);
        limiter.update_limit(42, Some(1.0), Some(1024.0), Instant::now());
        assert!(!limiter.is_legacy_read_limiter_enabled());
        assert!(limiter.get_limiter(42).is_none());
        assert!(limiter.keyspace_limiters.is_empty());
    }

    #[test]
    fn test_keyspace_read_limiter_wait_time_max() {
        let req_limiter = KeyspaceReadLimiter::new(1, Duration::from_secs(1), MAX_WAIT_TIME);
        req_limiter.set_speed_limit(Some(1.0), Some(10.0));
        req_limiter.consume_request();
        let req_wait = req_limiter.take_wait_time();

        let bytes_limiter = KeyspaceReadLimiter::new(2, Duration::from_secs(1), MAX_WAIT_TIME);
        bytes_limiter.set_speed_limit(Some(1.0), Some(10.0));
        bytes_limiter.consume_bytes(100);
        let bytes_wait = bytes_limiter.take_wait_time();

        let both_limiter = KeyspaceReadLimiter::new(3, Duration::from_secs(1), MAX_WAIT_TIME);
        both_limiter.set_speed_limit(Some(1.0), Some(10.0));
        both_limiter.consume_request();
        both_limiter.consume_bytes(100);
        let both_wait = both_limiter.take_wait_time();

        let tolerance = Duration::from_millis(1);
        assert!(both_wait + tolerance >= req_wait);
        assert!(both_wait + tolerance >= bytes_wait);
    }

    #[test]
    fn test_keyspace_read_limiter_does_not_accumulate_recovery_duration() {
        let limiter = KeyspaceReadLimiter::new(1, Duration::from_secs(1), MAX_WAIT_TIME);
        limiter.set_speed_limit(Some(100.0), Some(f64::INFINITY));

        for _ in 0..20 {
            limiter.consume_request();
        }

        let wait_time = limiter.take_wait_time();
        assert!(
            wait_time < Duration::from_millis(500),
            "wait_time should be bounded by the limiter recovery horizon, got {:?}",
            wait_time
        );
    }

    #[tokio::test]
    async fn test_keyspace_read_limiter_wait_with_deadline() {
        let limiter = KeyspaceReadLimiter::new(1, Duration::from_secs(1), MAX_WAIT_TIME);
        limiter.set_speed_limit(Some(1.0), Some(10.0));
        limiter.consume_request();
        let deadline = Instant::now();
        limiter.wait_with_deadline(deadline).await.unwrap_err();
    }
}