// Copyright 2026 PingCAP, Inc.

use std::{
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
    time::{Duration, Instant},
};

use async_trait::async_trait;
use resource_control::{KeyspaceReadLimiter, REQUEST_WAIT_HISTOGRAM_VEC};
use tidb_query_common::Result;
use tidb_query_executors::interface::ScanLimiter;

/// Rate limiter for DAG scan batches driven by keyspace read limits.
pub struct ScanRateLimiter {
    limiter: KeyspaceReadLimiter,
    deadline: Instant,
    dry_run: bool,
    debug: bool,
    keyspace_id: u32,
    region_id: u64,
    task_id: u64,
    request_consumed: AtomicBool,
    used: AtomicBool,
    wait_nanos: AtomicU64,
}

impl ScanRateLimiter {
    /// Creates a scan limiter bound to a request deadline and keyspace scope.
    pub fn new(
        limiter: KeyspaceReadLimiter,
        deadline: Instant,
        dry_run: bool,
        debug: bool,
        keyspace_id: u32,
        region_id: u64,
        task_id: u64,
    ) -> Self {
        Self {
            limiter,
            deadline,
            dry_run,
            debug,
            keyspace_id,
            region_id,
            task_id,
            request_consumed: AtomicBool::new(false),
            used: AtomicBool::new(false),
            wait_nanos: AtomicU64::new(0),
        }
    }

    fn record_wait(&self, wait_time: Duration) {
        let wait_nanos = wait_time.as_nanos().min(u64::MAX as u128) as u64;
        self.wait_nanos.fetch_add(wait_nanos, Ordering::Relaxed);
    }
}

#[async_trait]
impl ScanLimiter for ScanRateLimiter {
    async fn on_batch(&self, batch_bytes: usize) -> Result<()> {
        self.used.store(true, Ordering::Relaxed);
        if !self.request_consumed.swap(true, Ordering::AcqRel) {
            self.limiter.consume_request();
        }
        if batch_bytes > 0 {
            self.limiter.consume_bytes(batch_bytes);
        }
        let wait_time = if self.dry_run {
            self.limiter.take_wait_time()
        } else {
            self.limiter.wait_with_deadline(self.deadline).await?
        };
        if !wait_time.is_zero() {
            self.record_wait(wait_time);
        }
        Ok(())
    }
}

impl Drop for ScanRateLimiter {
    fn drop(&mut self) {
        if !self.used.load(Ordering::Relaxed) {
            return;
        }
        let wait_time = Duration::from_nanos(self.wait_nanos.load(Ordering::Relaxed));
        REQUEST_WAIT_HISTOGRAM_VEC
            .with_label_values(&["read"])
            .observe(wait_time.as_secs_f64());
        if self.debug && !wait_time.is_zero() {
            info!(
                "resource control keyspace {}, region {}, task {}, mock {} sleep {:?} for DAG request",
                self.keyspace_id, self.region_id, self.task_id, self.dry_run, wait_time
            );
        }
    }
}

/// Shared scan limiter handle.
pub type ScanLimiterRef = Arc<dyn ScanLimiter>;

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

    #[tokio::test]
    async fn test_scan_limiter_consumes_one_request_and_all_batch_bytes() {
        let limiter = KeyspaceReadLimiter::default();
        let scan_limiter = ScanRateLimiter::new(
            limiter.clone(),
            Instant::now() + Duration::from_secs(1),
            false,
            false,
            1,
            2,
            3,
        );

        scan_limiter.on_batch(11).await.unwrap();
        scan_limiter.on_batch(31).await.unwrap();

        assert_eq!(limiter.total_consumed(), (1, 42));
    }
}