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

use std::sync::{
    Arc,
    atomic::{AtomicU64, Ordering},
};

use async_trait::async_trait;
use resource_control::{CpuTokenHandle, ThrottleError};
use tidb_query_common::{Result, error::StorageError};
use tidb_query_executors::interface::CoprRuntimeBudget;

use crate::coprocessor::Error;

/// DAG batch-level CPU budget backed by the request's CPU token handle.
pub struct CpuTokenRuntimeBudget {
    token_handle: Arc<CpuTokenHandle>,
    consumed_cpu_us: AtomicU64,
}

impl CpuTokenRuntimeBudget {
    pub fn new(token_handle: Arc<CpuTokenHandle>) -> Self {
        Self {
            token_handle,
            consumed_cpu_us: AtomicU64::new(0),
        }
    }
}

#[async_trait]
impl CoprRuntimeBudget for CpuTokenRuntimeBudget {
    async fn on_batch(&self, cpu_delta_us: u64) -> Result<()> {
        let consumed_cpu_us = self
            .consumed_cpu_us
            .fetch_add(cpu_delta_us, Ordering::AcqRel)
            + cpu_delta_us;
        self.token_handle.record_actual_usage(consumed_cpu_us);
        let per_allocation_us = self
            .token_handle
            .get_runtime_config()
            .map(|(_, _, per_allocation_us)| per_allocation_us)
            .unwrap_or(0);
        let target_cpu_us = consumed_cpu_us.saturating_add(per_allocation_us);

        loop {
            let allocated_us = self.token_handle.allocated();
            if target_cpu_us <= allocated_us {
                return Ok(());
            }

            let additional_needed_us = target_cpu_us - allocated_us;
            let allocation_us = runtime_allocation_us(
                additional_needed_us,
                per_allocation_us,
                self.token_handle.max_single_runtime_allocation_us(),
            );

            match self
                .token_handle
                .allocate_more_with_wait(allocation_us)
                .await
            {
                Ok(_) => {}
                Err(ThrottleError::RequestTimeout) => {
                    warn!(
                        "[CPU throttle] Runtime CPU token allocation timeout at DAG batch boundary";
                        "keyspace_id" => self.token_handle.keyspace_id(),
                        "allocated_us" => self.token_handle.allocated(),
                        "consumed_cpu_us" => consumed_cpu_us,
                        "additional_needed_us" => additional_needed_us,
                        "allocation_us" => allocation_us,
                    );
                    return Err(cpu_throttled(
                        "[CPU throttle] Request timeout waiting for runtime CPU tokens",
                    ));
                }
                Err(err) => {
                    warn!(
                        "[CPU throttle] Runtime CPU token allocation failed at DAG batch boundary";
                        "keyspace_id" => self.token_handle.keyspace_id(),
                        "allocated_us" => self.token_handle.allocated(),
                        "consumed_cpu_us" => consumed_cpu_us,
                        "additional_needed_us" => additional_needed_us,
                        "allocation_us" => allocation_us,
                        "error" => ?err,
                    );
                    return Err(cpu_throttled(format!(
                        "[CPU throttle] Failed to allocate runtime CPU tokens: {:?}",
                        err
                    )));
                }
            }
        }
    }
}

fn runtime_allocation_us(
    additional_needed_us: u64,
    per_allocation_us: u64,
    max_allocation_us: u64,
) -> u64 {
    let allocation_us = if per_allocation_us > 0 {
        additional_needed_us.max(per_allocation_us)
    } else {
        additional_needed_us
    };
    allocation_us.min(max_allocation_us.max(1))
}

fn cpu_throttled(message: impl Into<String>) -> tidb_query_common::Error {
    StorageError::from(Error::CpuThrottled(message.into())).into()
}

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

    #[test]
    fn test_runtime_allocation_covers_deficit_in_one_chunk() {
        assert_eq!(runtime_allocation_us(8_000, 2_000, 20_000), 8_000);
    }

    #[test]
    fn test_runtime_allocation_preserves_min_chunk() {
        assert_eq!(runtime_allocation_us(500, 2_000, 20_000), 2_000);
    }

    #[test]
    fn test_runtime_allocation_caps_single_bucket_attempt() {
        assert_eq!(runtime_allocation_us(30_000, 2_000, 10_000), 10_000);
    }
}