// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use std::{
    sync::{Arc, RwLock},
    time::Duration,
};

use tikv_util::{config::ReadableSize, time::Limiter};

use crate::metrics::{
    ENGINE_IA_TO_STANDARD_TRANSITION_REMOTE_READ_BYTES,
    ENGINE_IA_TO_STANDARD_TRANSITION_REMOTE_READ_THROTTLE_DURATION,
};

/// Per-store limiter for storage-class transition materialization.
///
/// The schema worker owns region admission because it is single threaded. This
/// limiter stores the dynamic region cap and paces full-file loads from DFS at
/// file granularity.
#[derive(Clone)]
pub struct TransitionLimiter {
    inner: Arc<TransitionLimiterInner>,
}

#[must_use]
pub struct TransitionRemoteReadCharge {
    inner: Arc<TransitionLimiterInner>,
    bytes: u64,
    committed: bool,
}

struct TransitionLimiterInner {
    max_regions: RwLock<usize>,
    remote_read: Limiter,
    wait_timeout: RwLock<Duration>,
}

impl TransitionRemoteReadCharge {
    /// Marks the charged remote read as completed.
    pub fn commit(mut self) {
        self.committed = true;
        ENGINE_IA_TO_STANDARD_TRANSITION_REMOTE_READ_BYTES.inc_by(self.bytes);
    }
}

impl Drop for TransitionRemoteReadCharge {
    fn drop(&mut self) {
        if !self.committed {
            self.inner.remote_read.unconsume(self.bytes as usize);
        }
    }
}

impl TransitionLimiter {
    /// Creates a limiter with a per-store remote-read cap, region cap, and
    /// per-region wait timeout.
    pub fn new(max_throughput: ReadableSize, max_regions: usize, wait_timeout: Duration) -> Self {
        assert!(max_regions > 0, "max regions must be positive");
        assert!(!wait_timeout.is_zero(), "wait timeout must be positive");
        Self {
            inner: Arc::new(TransitionLimiterInner {
                max_regions: RwLock::new(max_regions),
                remote_read: <Limiter>::builder(speed_limit(max_throughput))
                    .refill(Duration::from_millis(10))
                    .build(),
                wait_timeout: RwLock::new(wait_timeout),
            }),
        }
    }

    /// Throttles one remote file read by the file size known before reading.
    pub async fn throttle_remote_read(&self, bytes: u64) -> TransitionRemoteReadCharge {
        let duration = self.inner.remote_read.consume_duration(bytes as usize);
        if !duration.is_zero() {
            tokio::time::sleep(duration).await;
            ENGINE_IA_TO_STANDARD_TRANSITION_REMOTE_READ_THROTTLE_DURATION
                .inc_by(duration.as_secs_f64());
        }
        TransitionRemoteReadCharge {
            inner: self.inner.clone(),
            bytes,
            committed: false,
        }
    }

    /// Updates the per-store remote-read throughput cap. `0` means unlimited.
    pub fn update_max_throughput(&self, max_throughput: ReadableSize) {
        let speed_limit = speed_limit(max_throughput);
        let old_speed_limit = self.inner.remote_read.speed_limit();
        if speed_limit != old_speed_limit {
            info!(
                "update IA to Standard transition remote read limiter";
                "old_speed_limit" => old_speed_limit,
                "speed_limit" => speed_limit,
            );
            self.inner.remote_read.set_speed_limit(speed_limit);
        }
    }

    /// Updates the per-store transition region cap.
    ///
    /// Lowering the cap does not cancel existing regions. The schema worker
    /// reads this value before admitting a new transition region.
    pub fn update_max_regions(&self, max_regions: usize) {
        assert!(max_regions > 0, "max regions must be positive");
        let mut current = self.inner.max_regions.write().unwrap();
        let old_max_regions = *current;
        if old_max_regions != max_regions {
            *current = max_regions;
            info!(
                "update IA to Standard transition region limiter";
                "old_max_regions" => old_max_regions,
                "max_regions" => max_regions,
            );
        }
    }

    /// Updates the timeout used by subsequent transition regions.
    pub fn update_wait_timeout(&self, wait_timeout: Duration) {
        assert!(!wait_timeout.is_zero(), "wait timeout must be positive");
        let mut current = self.inner.wait_timeout.write().unwrap();
        if *current != wait_timeout {
            info!(
                "update IA to Standard transition wait timeout";
                "old_wait_timeout" => ?*current,
                "wait_timeout" => ?wait_timeout,
            );
            *current = wait_timeout;
        }
    }

    /// Returns the timeout to use for the next transition region.
    pub fn wait_timeout(&self) -> Duration {
        *self.inner.wait_timeout.read().unwrap()
    }

    /// Returns the region cap to use for schema-worker admission.
    pub fn max_regions(&self) -> usize {
        *self.inner.max_regions.read().unwrap()
    }

    #[cfg(test)]
    pub(crate) fn remote_read_speed_limit(&self) -> f64 {
        self.inner.remote_read.speed_limit()
    }

    #[cfg(test)]
    pub(crate) fn remote_read_total_bytes_consumed(&self) -> usize {
        self.inner.remote_read.total_bytes_consumed()
    }
}

fn speed_limit(max_throughput: ReadableSize) -> f64 {
    if max_throughput.0 == 0 {
        f64::INFINITY
    } else {
        max_throughput.0 as f64
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_transition_limiter_dynamic_updates() {
        let limiter = TransitionLimiter::new(ReadableSize::mb(100), 4, Duration::from_secs(60));

        limiter.update_max_throughput(ReadableSize(0));
        limiter.update_max_regions(2);
        limiter.update_wait_timeout(Duration::from_secs(120));

        assert_eq!(limiter.remote_read_speed_limit(), f64::INFINITY);
        assert_eq!(limiter.max_regions(), 2);
        assert_eq!(limiter.wait_timeout(), Duration::from_secs(120));
    }

    #[test]
    fn test_transition_remote_read_charge_unconsumes_until_committed() {
        let limiter = TransitionLimiter::new(ReadableSize(0), 4, Duration::from_secs(60));

        let charge = futures::executor::block_on(limiter.throttle_remote_read(1024));
        assert_eq!(limiter.remote_read_total_bytes_consumed(), 1024);
        drop(charge);
        assert_eq!(limiter.remote_read_total_bytes_consumed(), 0);

        let charge = futures::executor::block_on(limiter.throttle_remote_read(2048));
        charge.commit();
        assert_eq!(limiter.remote_read_total_bytes_consumed(), 2048);
    }
}