// Copyright 2026 PingCAP, Inc.

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

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

use crate::metrics::{
    ENGINE_IA_IO_THROTTLE_DURATION, ENGINE_IA_IO_THROUGHPUT_BYTES, ENGINE_IA_IO_THROUGHPUT_CAP,
};

pub(crate) const IA_IO_RATE_LIMITER_REFILL_PERIOD_DEF: Duration = Duration::from_millis(100);

#[must_use]
pub(crate) struct IaIoCharge {
    limiter: Arc<Limiter>,
    kind: IaIoType,
    bytes: u64,
    committed: bool,
}

#[derive(Clone, Copy)]
enum IaIoType {
    RemoteRead,
    LocalWrite,
}

impl IaIoCharge {
    fn new(limiter: Arc<Limiter>, kind: IaIoType, bytes: u64) -> Self {
        Self {
            limiter,
            kind,
            bytes,
            committed: false,
        }
    }

    pub(crate) fn commit(mut self) {
        self.committed = true;
        match self.kind {
            IaIoType::RemoteRead => ENGINE_IA_IO_THROUGHPUT_BYTES.remote_read.inc_by(self.bytes),
            IaIoType::LocalWrite => ENGINE_IA_IO_THROUGHPUT_BYTES.local_write.inc_by(self.bytes),
        }
    }
}

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

#[derive(Clone)]
pub(crate) struct IaIoRateLimiter {
    remote_read: Arc<Limiter>,
    local_write: Arc<Limiter>,
}

impl IaIoRateLimiter {
    pub(crate) fn unlimited() -> Self {
        Self::new(
            ReadableSize(0),
            ReadableSize(0),
            IA_IO_RATE_LIMITER_REFILL_PERIOD_DEF,
        )
    }

    pub(crate) fn new(
        remote_read_max_throughput: ReadableSize,
        local_write_max_throughput: ReadableSize,
        refill_period: Duration,
    ) -> Self {
        let limiter = Self {
            remote_read: Self::new_limiter(remote_read_max_throughput, refill_period),
            local_write: Self::new_limiter(local_write_max_throughput, refill_period),
        };
        limiter.record_remote_read_cap(remote_read_max_throughput);
        limiter.record_local_write_cap(local_write_max_throughput);
        limiter
    }

    pub(crate) fn update(
        &self,
        remote_read_max_throughput: ReadableSize,
        local_write_max_throughput: ReadableSize,
    ) {
        self.update_remote_read(remote_read_max_throughput);
        self.update_local_write(local_write_max_throughput);
    }

    pub(crate) fn update_remote_read(&self, remote_read_max_throughput: ReadableSize) {
        self.record_remote_read_cap(remote_read_max_throughput);
        Self::update_limiter("remote read", &self.remote_read, remote_read_max_throughput);
    }

    pub(crate) fn update_local_write(&self, local_write_max_throughput: ReadableSize) {
        self.record_local_write_cap(local_write_max_throughput);
        Self::update_limiter("local write", &self.local_write, local_write_max_throughput);
    }

    pub(crate) async fn throttle_remote_read(&self, bytes: u64) -> IaIoCharge {
        let limiter = self.remote_read.clone();
        let duration = limiter.consume_duration(bytes as usize);
        let charge = IaIoCharge::new(limiter, IaIoType::RemoteRead, bytes);
        if !duration.is_zero() {
            tokio::time::sleep(duration).await;
            ENGINE_IA_IO_THROTTLE_DURATION
                .remote_read
                .observe(duration.as_secs_f64());
        }
        charge
    }

    pub(crate) fn throttle_local_write(&self, bytes: u64) -> IaIoCharge {
        let limiter = self.local_write.clone();
        let duration = limiter.consume_duration(bytes as usize);
        let charge = IaIoCharge::new(limiter, IaIoType::LocalWrite, bytes);
        if !duration.is_zero() {
            std::thread::sleep(duration);
            ENGINE_IA_IO_THROTTLE_DURATION
                .local_write
                .observe(duration.as_secs_f64());
        }
        charge
    }

    fn new_limiter(max_throughput: ReadableSize, refill_period: Duration) -> Arc<Limiter> {
        assert!(
            !refill_period.is_zero(),
            "IA IO rate limiter refill period must be positive"
        );
        Arc::new(
            <Limiter>::builder(Self::speed_limit(max_throughput))
                .refill(refill_period)
                .build(),
        )
    }

    fn update_limiter(name: &str, limiter: &Limiter, max_throughput: ReadableSize) {
        let speed_limit = Self::speed_limit(max_throughput);
        let old_speed_limit = limiter.speed_limit();
        if speed_limit != old_speed_limit {
            info!(
                "update IA IO rate limiter";
                "name" => name,
                "old_speed_limit" => old_speed_limit,
                "speed_limit" => speed_limit,
            );
            limiter.set_speed_limit(speed_limit);
        }
    }

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

    fn record_remote_read_cap(&self, max_throughput: ReadableSize) {
        ENGINE_IA_IO_THROUGHPUT_CAP
            .remote_read
            .set(max_throughput.0 as i64);
    }

    fn record_local_write_cap(&self, max_throughput: ReadableSize) {
        ENGINE_IA_IO_THROUGHPUT_CAP
            .local_write
            .set(max_throughput.0 as i64);
    }

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

    #[cfg(test)]
    pub(crate) fn local_write_speed_limit(&self) -> f64 {
        self.local_write.speed_limit()
    }

    #[cfg(any(test, feature = "testexport"))]
    pub(crate) fn remote_read_total_bytes_consumed(&self) -> usize {
        self.remote_read.total_bytes_consumed()
    }

    #[cfg(test)]
    pub(crate) fn local_write_total_bytes_consumed(&self) -> usize {
        self.local_write.total_bytes_consumed()
    }
}

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

    #[test]
    fn test_ia_io_rate_limiter_records_metrics() {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        let limiter = IaIoRateLimiter::new(
            ReadableSize::gb(1),
            ReadableSize::gb(1),
            Duration::from_millis(10),
        );
        let bytes = 11 * 1024 * 1024;

        assert_eq!(
            ENGINE_IA_IO_THROUGHPUT_CAP.remote_read.get(),
            ReadableSize::gb(1).0 as i64
        );
        assert_eq!(
            ENGINE_IA_IO_THROUGHPUT_CAP.local_write.get(),
            ReadableSize::gb(1).0 as i64
        );

        let remote_read_bytes = ENGINE_IA_IO_THROUGHPUT_BYTES.remote_read.get();
        let local_write_bytes = ENGINE_IA_IO_THROUGHPUT_BYTES.local_write.get();
        let remote_read_throttle_count = ENGINE_IA_IO_THROTTLE_DURATION
            .remote_read
            .get_sample_count();
        let local_write_throttle_count = ENGINE_IA_IO_THROTTLE_DURATION
            .local_write
            .get_sample_count();

        runtime
            .block_on(limiter.throttle_remote_read(bytes))
            .commit();
        limiter.throttle_local_write(bytes).commit();

        assert_eq!(
            ENGINE_IA_IO_THROUGHPUT_BYTES.remote_read.get(),
            remote_read_bytes + bytes
        );
        assert_eq!(
            ENGINE_IA_IO_THROUGHPUT_BYTES.local_write.get(),
            local_write_bytes + bytes
        );
        assert_eq!(
            ENGINE_IA_IO_THROTTLE_DURATION
                .remote_read
                .get_sample_count(),
            remote_read_throttle_count + 1
        );
        assert_eq!(
            ENGINE_IA_IO_THROTTLE_DURATION
                .local_write
                .get_sample_count(),
            local_write_throttle_count + 1
        );

        limiter.update_remote_read(ReadableSize::mb(512));
        limiter.update_local_write(ReadableSize(0));
        assert_eq!(
            ENGINE_IA_IO_THROUGHPUT_CAP.remote_read.get(),
            ReadableSize::mb(512).0 as i64
        );
        assert_eq!(ENGINE_IA_IO_THROUGHPUT_CAP.local_write.get(), 0);
    }

    #[test]
    fn test_ia_io_rate_limiter_uncommitted_bytes_are_refunded() {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        let limiter = IaIoRateLimiter::new(
            ReadableSize::gb(1),
            ReadableSize::gb(1),
            IA_IO_RATE_LIMITER_REFILL_PERIOD_DEF,
        );
        let bytes = 1024;
        let remote_read_bytes = ENGINE_IA_IO_THROUGHPUT_BYTES.remote_read.get();
        let local_write_bytes = ENGINE_IA_IO_THROUGHPUT_BYTES.local_write.get();

        runtime.block_on(async {
            let charge = limiter.throttle_remote_read(bytes).await;
            drop(charge);
        });
        let charge = limiter.throttle_local_write(bytes);
        drop(charge);

        assert_eq!(limiter.remote_read_total_bytes_consumed(), 0);
        assert_eq!(limiter.local_write_total_bytes_consumed(), 0);
        assert_eq!(
            ENGINE_IA_IO_THROUGHPUT_BYTES.remote_read.get(),
            remote_read_bytes
        );
        assert_eq!(
            ENGINE_IA_IO_THROUGHPUT_BYTES.local_write.get(),
            local_write_bytes
        );

        let throttled_limiter = IaIoRateLimiter::new(
            ReadableSize::kb(1),
            ReadableSize::gb(1),
            Duration::from_millis(10),
        );
        runtime.block_on(async {
            let result = tokio::time::timeout(
                Duration::from_millis(10),
                throttled_limiter.throttle_remote_read(1024),
            )
            .await;
            assert!(result.is_err());
        });
        assert_eq!(throttled_limiter.remote_read_total_bytes_consumed(), 0);
    }
}