//! Token bucket rate limiter for smooth packet pacing.
//!
//! Provides algorithm-agnostic rate limiting that works with any congestion control
//! algorithm (LEDBAT, AIMD, etc.). The token bucket ensures smooth packet pacing
//! without bursts, and supports dynamic rate updates from the congestion controller.

use parking_lot::Mutex;
use std::time::{Duration, Instant};

/// Token bucket for rate limiting packet transmission.
///
/// Uses a reserve + consume pattern to prevent TOCTOU (time-of-check-time-of-use) races:
/// 1. Reserve tokens (may require waiting)
/// 2. Wait if needed
/// 3. Consume the reservation
///
/// Thread-safe via internal mutex.
pub struct TokenBucket {
    state: Mutex<BucketState>,
}

struct BucketState {
    /// Maximum tokens (burst capacity in bytes)
    capacity: usize,
    /// Current available tokens (bytes)
    tokens: usize,
    /// Fractional tokens (prevents precision loss at high rates)
    fractional_tokens: f64,
    /// Refill rate (bytes/second)
    rate: usize,
    /// Last refill timestamp
    last_refill: Instant,
}

impl TokenBucket {
    /// Create a new token bucket.
    ///
    /// # Arguments
    /// * `capacity` - Maximum burst capacity (bytes)
    /// * `rate` - Refill rate (bytes/second)
    ///
    /// # Example
    /// ```ignore
    /// use freenet::transport::token_bucket::TokenBucket;
    /// let bucket = TokenBucket::new(
    ///     10_000,      // 10 KB burst
    ///     1_000_000,   // 1 MB/s rate
    /// );
    /// ```
    pub fn new(capacity: usize, rate: usize) -> Self {
        Self {
            state: Mutex::new(BucketState {
                capacity,
                tokens: capacity, // Start full
                fractional_tokens: 0.0,
                rate,
                last_refill: Instant::now(),
            }),
        }
    }

    /// Reserve tokens for transmission.
    ///
    /// Immediately deducts tokens and returns wait time if needed.
    /// Caller should wait before transmitting if wait time > 0.
    ///
    /// # Arguments
    /// * `bytes` - Number of bytes to send
    ///
    /// # Returns
    /// Duration to wait before transmitting
    ///
    /// # Example
    /// ```ignore
    /// use freenet::transport::token_bucket::TokenBucket;
    /// use std::time::Duration;
    /// tokio_test::block_on(async {
    ///     let bucket = TokenBucket::new(1000, 1_000_000);
    ///
    ///     let wait = bucket.reserve(500);
    ///     if wait > Duration::ZERO {
    ///         tokio::time::sleep(wait).await;
    ///     }
    ///     // Tokens already deducted, can transmit now
    /// });
    /// ```
    pub fn reserve(&self, bytes: usize) -> Duration {
        let mut state = self.state.lock();
        state.refill();

        // Calculate wait time BEFORE deducting
        let wait_time = if state.tokens >= bytes {
            // Sufficient tokens available
            Duration::ZERO
        } else if state.rate == 0 {
            // Rate is 0 (not yet initialized or set to 0) - no rate limiting
            Duration::ZERO
        } else {
            // Need to wait for more tokens
            let deficit = bytes - state.tokens;
            let wait_secs = deficit as f64 / state.rate as f64;
            Duration::from_secs_f64(wait_secs)
        };

        // Deduct tokens (may go to 0 if we had a deficit)
        state.tokens = state.tokens.saturating_sub(bytes);

        wait_time
    }

    /// Update the refill rate dynamically.
    ///
    /// Called by the congestion controller when the rate changes.
    ///
    /// # Arguments
    /// * `new_rate` - New refill rate (bytes/second)
    pub fn set_rate(&self, new_rate: usize) {
        let mut state = self.state.lock();
        // Enforce minimum rate of 1KB/s to prevent division by zero
        state.rate = new_rate.max(1024);
    }

    /// Get current rate (bytes/second)
    pub fn rate(&self) -> usize {
        self.state.lock().rate
    }

    /// Get current available tokens (bytes)
    #[cfg(test)]
    pub fn available_tokens(&self) -> usize {
        let mut state = self.state.lock();
        state.refill();
        state.tokens
    }
}

impl BucketState {
    /// Refill tokens based on elapsed time.
    ///
    /// Uses fractional token tracking to prevent precision loss at high rates.
    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill);

        // Avoid precision issues with very small intervals
        if elapsed < Duration::from_millis(1) {
            return;
        }

        // Calculate new tokens with fractional precision
        let new_tokens_f64 = self.rate as f64 * elapsed.as_secs_f64();

        // Add to fractional accumulator
        self.fractional_tokens += new_tokens_f64;

        // Convert whole tokens
        let new_tokens_whole = self.fractional_tokens.floor() as usize;
        if new_tokens_whole > 0 {
            // Add tokens, capped at capacity
            self.tokens = (self.tokens + new_tokens_whole).min(self.capacity);

            // Keep fractional part for next refill
            self.fractional_tokens -= new_tokens_whole as f64;

            self.last_refill = now;
        }
    }
}

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

    #[test]
    fn test_token_bucket_creation() {
        let bucket = TokenBucket::new(10_000, 1_000_000);
        assert_eq!(bucket.available_tokens(), 10_000);
        assert_eq!(bucket.rate(), 1_000_000);
    }

    #[test]
    fn test_token_bucket_immediate_tokens() {
        let bucket = TokenBucket::new(10_000, 1_000_000);

        // Reserve less than available - should be immediate
        let wait = bucket.reserve(5_000);
        assert_eq!(wait, Duration::ZERO);

        // Tokens already deducted
        assert_eq!(bucket.available_tokens(), 5_000);
    }

    #[test]
    fn test_token_bucket_requires_wait() {
        let bucket = TokenBucket::new(1_000, 1_000_000); // 1KB burst, 1MB/s rate

        // Reserve more than available
        let wait = bucket.reserve(5_000);

        // Should need to wait for 4KB at 1MB/s
        // wait ≈ 4000 / 1_000_000 = 0.004s = 4ms
        assert!(wait > Duration::ZERO);
        assert!(wait >= Duration::from_millis(3));
        assert!(wait <= Duration::from_millis(5));
    }

    #[tokio::test]
    async fn test_token_bucket_rate_limiting_accuracy() {
        let bucket = TokenBucket::new(1_000, 1_000_000); // 1KB burst, 1MB/s rate

        let start = Instant::now();
        let mut total_sent = 0;

        // Send 100 KB in 1KB chunks
        while total_sent < 100_000 {
            let wait = bucket.reserve(1000);
            if wait > Duration::ZERO {
                tokio::time::sleep(wait).await;
            }
            // Tokens already deducted
            total_sent += 1000;
        }

        let elapsed = start.elapsed();

        // Should take ~100ms (100KB at 1MB/s)
        // Allow wide tolerance for timing variance (tokio sleep isn't precise,
        // especially under system load)
        assert!(
            elapsed >= Duration::from_millis(50),
            "Too fast: {:?}",
            elapsed
        );
        assert!(
            elapsed <= Duration::from_millis(300),
            "Too slow: {:?}",
            elapsed
        );
    }

    #[test]
    fn test_token_bucket_dynamic_rate_update() {
        let bucket = TokenBucket::new(10_000, 1_000_000);

        assert_eq!(bucket.rate(), 1_000_000);

        // Update rate
        bucket.set_rate(5_000_000);
        assert_eq!(bucket.rate(), 5_000_000);
    }

    #[tokio::test]
    async fn test_token_bucket_refill() {
        let bucket = TokenBucket::new(10_000, 10_000_000); // 10KB capacity, 10MB/s rate

        // Consume all tokens
        let wait = bucket.reserve(10_000);
        assert_eq!(wait, Duration::ZERO);
        assert_eq!(bucket.available_tokens(), 0);

        // Wait 10ms - should refill 100KB at 10MB/s, but capped at capacity (10KB)
        tokio::time::sleep(Duration::from_millis(10)).await;

        // Should be back to capacity
        let available = bucket.available_tokens();
        assert!(available >= 9_000, "Only refilled to {}", available); // Allow some timing variance
        assert!(available <= 10_000);
    }

    #[test]
    fn test_token_bucket_deducts_immediately() {
        let bucket = TokenBucket::new(5_000, 1_000_000);

        // Reserve more than available
        let wait = bucket.reserve(10_000);
        assert!(wait > Duration::ZERO);

        // Tokens should be depleted (deducted immediately)
        let available = bucket.available_tokens();
        assert_eq!(available, 0);
    }

    #[tokio::test]
    async fn test_token_bucket_concurrent_streams() {
        // Simulate multiple streams sharing a token bucket
        let bucket = std::sync::Arc::new(TokenBucket::new(10_000, 1_000_000));

        let start = Instant::now();
        let mut handles = vec![];

        // Spawn 5 tasks, each sending 20KB
        for _ in 0..5 {
            let bucket_clone = bucket.clone();
            let handle = tokio::spawn(async move {
                let mut sent = 0;
                while sent < 20_000 {
                    let wait = bucket_clone.reserve(1000);
                    if wait > Duration::ZERO {
                        tokio::time::sleep(wait).await;
                    }
                    // Tokens already deducted
                    sent += 1000;
                }
            });
            handles.push(handle);
        }

        // Wait for all tasks
        for handle in handles {
            handle.await.unwrap();
        }

        let elapsed = start.elapsed();

        // Total: 100KB at 1MB/s should take ~100ms
        // Allow very wide tolerance for concurrent scheduling variance
        // Lowered from 40ms to 35ms to account for timing precision on fast systems
        assert!(
            elapsed >= Duration::from_millis(35),
            "Too fast: {:?}",
            elapsed
        );
        assert!(
            elapsed <= Duration::from_millis(250),
            "Too slow: {:?}",
            elapsed
        );
    }
}
