// Copyright 2026 PingCAP, Inc.

use std::{
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
    time::{Duration, Instant},
};

use futures_util::future::BoxFuture;
use pin_project::pin_project;
use tikv_util::{cpu_observe::CpuObserveFuture, info, warn};

use crate::{CpuTokenHandle, ThrottleError};

/// Error type for CPU token checking
#[derive(Debug)]
pub enum CpuTokenError {
    /// Request timed out waiting for CPU tokens
    Timeout,
}

impl std::fmt::Display for CpuTokenError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CpuTokenError::Timeout => write!(f, "Request timeout waiting for CPU tokens"),
        }
    }
}

impl std::error::Error for CpuTokenError {}

/// State machine for CPU token checking
enum CheckState {
    /// Normal execution - polling the wrapped future
    Executing {
        last_check_time: Instant,
        last_cpu_check: u64,
    },

    /// Allocating more tokens - waiting for allocate_more_with_wait()
    AllocatingTokens {
        allocation_future: BoxFuture<'static, Result<u64, ThrottleError>>,
    },

    /// Request timed out waiting for tokens
    TimedOut,

    /// Completed successfully
    Completed,
}

/// Future wrapper that performs periodic CPU token checking during execution
///
/// This future:
/// 1. Wraps an inner future being executed
/// 2. Tracks CPU time consumption using CpuObserveFuture
/// 3. Periodically checks if CPU usage exceeds threshold
/// 4. Allocates additional tokens via allocate_more_with_wait() when needed
/// 5. Handles timeout by terminating the request
///
/// # Type Parameters
/// - `F`: The wrapped future type
///
/// # Example
/// ```ignore
/// let check_future = CpuTokenCheckFuture::new(
///     my_async_task,
///     token_handle,
///     Duration::from_millis(2),  // check every 2ms
///     0.5,                        // allocate when 50% consumed
///     4000,                       // allocate 4000us each time
/// );
///
/// match check_future.await {
///     Ok(result) => println!("Task completed: {:?}", result),
///     Err(CpuTokenError::Timeout) => println!("Timed out waiting for tokens"),
/// }
/// ```
#[pin_project]
pub struct CpuTokenCheckFuture<F> {
    /// The wrapped future being executed
    #[pin]
    observed: CpuObserveFuture<F>,

    /// CPU token handle for allocation
    token_handle: Arc<CpuTokenHandle>,

    /// Current state of the state machine
    state: CheckState,

    /// Check interval - how often to check CPU usage
    check_interval: Duration,

    /// Threshold percentage (0.0 - 1.0) - when to allocate more tokens
    /// e.g., 0.5 means allocate when usage reaches 50% of allocated tokens
    additional_allocation_threshold: f64,

    /// Amount to allocate each time in microseconds
    per_allocation_us: u64,
}

impl<F> CpuTokenCheckFuture<F> {
    /// Create a new CpuTokenCheckFuture
    ///
    /// # Arguments
    /// * `inner` - The future to wrap and execute
    /// * `token_handle` - Handle for CPU token allocation
    /// * `check_interval` - How often to check CPU usage
    /// * `additional_allocation_threshold` - Threshold percentage (0.0 - 1.0)
    /// * `per_allocation_us` - Microseconds to allocate each time
    pub fn new(
        inner: F,
        token_handle: Arc<CpuTokenHandle>,
        check_interval: Duration,
        additional_allocation_threshold: f64,
        per_allocation_us: u64,
    ) -> Self {
        Self {
            observed: CpuObserveFuture::new(inner),
            token_handle,
            state: CheckState::Executing {
                last_check_time: Instant::now(),
                last_cpu_check: 0,
            },
            check_interval,
            additional_allocation_threshold,
            per_allocation_us,
        }
    }
}

impl<F: Future> Future for CpuTokenCheckFuture<F> {
    type Output = Result<F::Output, CpuTokenError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();

        loop {
            match this.state {
                CheckState::Executing {
                    last_check_time,
                    last_cpu_check,
                } => {
                    // 1. Poll the inner future
                    match this.observed.as_mut().poll(cx) {
                        Poll::Ready((cpu_time, result)) => {
                            // Record final CPU usage
                            this.token_handle
                                .record_actual_usage(cpu_time.as_micros() as u64);
                            *this.state = CheckState::Completed;
                            return Poll::Ready(Ok(result));
                        }
                        Poll::Pending => {
                            // 2. Check if enough time has passed for a periodic check
                            let now = Instant::now();
                            if now.duration_since(*last_check_time) >= *this.check_interval {
                                // 3. Get current CPU usage
                                let current_cpu_us =
                                    this.observed.total_duration().as_micros() as u64;
                                let allocated = this.token_handle.allocated();
                                let threshold = (allocated as f64
                                    * *this.additional_allocation_threshold)
                                    as u64;

                                // 4. Check if we need more tokens
                                if current_cpu_us > threshold {
                                    // Create allocation future
                                    let handle = this.token_handle.clone();
                                    let per_allocation_us = *this.per_allocation_us;
                                    let allocation_future: BoxFuture<
                                        'static,
                                        Result<u64, ThrottleError>,
                                    > = Box::pin(async move {
                                        handle.allocate_more_with_wait(per_allocation_us).await
                                    });

                                    *this.state =
                                        CheckState::AllocatingTokens { allocation_future };
                                    continue; // Loop to handle AllocatingTokens state
                                }

                                *last_check_time = now;
                                *last_cpu_check = current_cpu_us;
                            }

                            return Poll::Pending;
                        }
                    }
                }

                CheckState::AllocatingTokens { allocation_future } => {
                    // Poll the allocation future
                    match allocation_future.as_mut().poll(cx) {
                        Poll::Ready(Ok(allocated_us)) => {
                            // Allocation succeeded, return to executing
                            // Rate-limit log: always print in debug mode, otherwise every 5 seconds
                            let should_log =
                                if this.token_handle.config.as_ref().is_some_and(|c| c.debug) {
                                    true
                                } else {
                                    this.token_handle
                                        .last_periodic_allocation_log
                                        .as_ref()
                                        .is_some_and(|last_log| {
                                            let mut last_log = last_log.write().unwrap();
                                            if last_log.elapsed() >= Duration::from_secs(5) {
                                                *last_log = Instant::now();
                                                true
                                            } else {
                                                false
                                            }
                                        })
                                };

                            if should_log {
                                info!(
                                    "CPU token periodic allocation succeeded";
                                    "keyspace_id" => this.token_handle.keyspace_id,
                                    "allocated_us" => allocated_us,
                                    "total_allocated" => this.token_handle.allocated(),
                                );
                            }

                            *this.state = CheckState::Executing {
                                last_check_time: Instant::now(),
                                last_cpu_check: this.observed.total_duration().as_micros() as u64,
                            };

                            // Wake the inner future to continue
                            cx.waker().wake_by_ref();
                            continue;
                        }
                        Poll::Ready(Err(ThrottleError::RequestTimeout)) => {
                            let actual_cpu_us = this.observed.total_duration().as_micros() as u64;
                            this.token_handle.record_actual_usage(actual_cpu_us);
                            // Timeout - terminate request
                            warn!(
                                "CPU token allocation timeout";
                                "keyspace_id" => this.token_handle.keyspace_id,
                                "allocated" => this.token_handle.allocated(),
                                "current_cpu_us" => actual_cpu_us,
                            );
                            *this.state = CheckState::TimedOut;
                            return Poll::Ready(Err(CpuTokenError::Timeout));
                        }
                        Poll::Ready(Err(e)) => {
                            // Other errors - log and continue execution
                            warn!(
                                "CPU token allocation failed, continuing execution";
                                "keyspace_id" => this.token_handle.keyspace_id,
                                "error" => ?e,
                            );
                            *this.state = CheckState::Executing {
                                last_check_time: Instant::now(),
                                last_cpu_check: this.observed.total_duration().as_micros() as u64,
                            };
                            cx.waker().wake_by_ref();
                            continue;
                        }
                        Poll::Pending => {
                            return Poll::Pending;
                        }
                    }
                }

                CheckState::TimedOut => {
                    return Poll::Ready(Err(CpuTokenError::Timeout));
                }

                CheckState::Completed => {
                    panic!("CpuTokenCheckFuture polled after completion");
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{RwLock, atomic::AtomicU64};

    use super::*;
    use crate::CpuTokenBucket;

    #[tokio::test]
    async fn test_cpu_token_check_future_success() {
        let global_bucket = Arc::new(CpuTokenBucket::new(
            100000,
            100000,
            Duration::from_millis(100),
            0, // interval_allocation_limit (0 = disabled for test)
        ));
        let keyspace_bucket = Arc::new(CpuTokenBucket::new(
            50000,
            50000,
            Duration::from_millis(100),
            0, // interval_allocation_limit (0 = disabled for test)
        ));

        // Allocate initial tokens
        assert!(global_bucket.try_allocate(10000));
        assert!(keyspace_bucket.try_allocate(10000));

        let handle = Arc::new(CpuTokenHandle {
            keyspace_id: 1,
            allocated_us: AtomicU64::new(10000),
            keyspace_allocated_us: AtomicU64::new(10000),
            global_allocated_us: AtomicU64::new(10000),
            actual_used_us: AtomicU64::new(0),
            global_bucket: global_bucket.clone(),
            keyspace_bucket: Some(keyspace_bucket.clone()),
            low_priority_bucket: None,
            keyspace_inflight_state: None,
            low_priority_allocated_us: AtomicU64::new(0),
            low_priority_inflight_allocated_us: None,
            is_burst: std::sync::atomic::AtomicBool::new(false),
            request_deadline: Instant::now() + Duration::from_secs(5),
            skip_throttle: false,
            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()))),
        });

        let check_future = CpuTokenCheckFuture::new(
            async {
                // Simulate some work
                tokio::time::sleep(Duration::from_millis(1)).await;
                42
            },
            handle.clone(),
            Duration::from_millis(10), // check interval
            0.5,                       // threshold
            4000,                      // per allocation
        );

        match check_future.await {
            Ok(value) => {
                assert_eq!(value, 42);
            }
            Err(CpuTokenError::Timeout) => {
                panic!("Should not timeout");
            }
        }
    }

    #[tokio::test]
    async fn test_cpu_token_check_future_timeout() {
        let global_bucket = Arc::new(CpuTokenBucket::new(
            1000,
            1000,
            Duration::from_millis(100),
            0,
        ));
        let keyspace_bucket =
            Arc::new(CpuTokenBucket::new(500, 500, Duration::from_millis(100), 0));

        // Allocate initial tokens (leave very little)
        assert!(global_bucket.try_allocate(500));
        assert!(keyspace_bucket.try_allocate(250));

        let handle = Arc::new(CpuTokenHandle {
            keyspace_id: 1,
            allocated_us: AtomicU64::new(500),
            keyspace_allocated_us: AtomicU64::new(250),
            global_allocated_us: AtomicU64::new(500),
            actual_used_us: AtomicU64::new(0),
            global_bucket: global_bucket.clone(),
            keyspace_bucket: Some(keyspace_bucket.clone()),
            low_priority_bucket: None,
            keyspace_inflight_state: None,
            low_priority_allocated_us: AtomicU64::new(0),
            low_priority_inflight_allocated_us: None,
            is_burst: std::sync::atomic::AtomicBool::new(false),
            request_deadline: Instant::now() + Duration::from_millis(100), // Short deadline
            skip_throttle: false,
            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()))),
        });

        let check_future = CpuTokenCheckFuture::new(
            async {
                // Long-running task
                tokio::time::sleep(Duration::from_secs(10)).await;
                42
            },
            handle.clone(),
            Duration::from_micros(10), // Very frequent checks
            0.1,                       // Low threshold - trigger allocation quickly
            10000,                     // Large allocation - will fail
        );

        match check_future.await {
            Ok(_) => {
                // May succeed if task completes quickly
            }
            Err(CpuTokenError::Timeout) => {
                // Expected: timeout waiting for tokens
            }
        }
    }
}