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

const EPSILON: f64 = 1e-9;

/// A simple token-bucket rate limiter.
///
/// The bucket starts full. Tokens refill over time up to `capacity`, and each
/// successful acquisition removes tokens from the bucket.
#[derive(Debug, Clone)]
pub struct TokenBucket {
    capacity: u64,
    tokens: f64,
    refill_per_second: f64,
    last_refill: Instant,
}

impl TokenBucket {
    /// Creates a bucket with `capacity` maximum tokens and a refill rate in
    /// tokens per second.
    pub fn new(capacity: u64, refill_per_second: f64) -> Self {
        assert!(capacity > 0, "capacity must be greater than zero");
        assert!(
            refill_per_second.is_finite() && refill_per_second > 0.0,
            "refill_per_second must be finite and greater than zero"
        );

        Self {
            capacity,
            tokens: capacity as f64,
            refill_per_second,
            last_refill: Instant::now(),
        }
    }

    /// Attempts to acquire `amount` tokens without blocking.
    pub fn try_acquire(&mut self, amount: u64) -> bool {
        assert!(amount > 0, "amount must be greater than zero");

        if amount > self.capacity {
            return false;
        }

        self.refill();

        let amount = amount as f64;
        if self.tokens + EPSILON >= amount {
            self.tokens = (self.tokens - amount).max(0.0);
            true
        } else {
            false
        }
    }

    /// Returns the current number of whole tokens available.
    pub fn available(&mut self) -> u64 {
        self.refill();
        self.tokens.floor() as u64
    }

    /// Returns the maximum number of tokens this bucket can hold.
    pub fn capacity(&self) -> u64 {
        self.capacity
    }

    /// Returns the refill rate in tokens per second.
    pub fn refill_per_second(&self) -> f64 {
        self.refill_per_second
    }

    /// Returns how long until `amount` tokens are expected to be available.
    pub fn time_until_available(&mut self, amount: u64) -> Option<Duration> {
        assert!(amount > 0, "amount must be greater than zero");

        if amount > self.capacity {
            return None;
        }

        self.refill();

        let needed = amount as f64 - self.tokens;
        if needed <= EPSILON {
            Some(Duration::ZERO)
        } else {
            Some(Duration::from_secs_f64(needed / self.refill_per_second))
        }
    }

    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill);

        if elapsed.is_zero() {
            return;
        }

        let added = elapsed.as_secs_f64() * self.refill_per_second;
        self.tokens = (self.tokens + added).min(self.capacity as f64);
        self.last_refill = now;
    }
}

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

    #[test]
    fn consumes_tokens_until_empty() {
        let mut bucket = TokenBucket::new(3, 10.0);

        assert!(bucket.try_acquire(1));
        assert!(bucket.try_acquire(2));
        assert!(!bucket.try_acquire(1));
    }

    #[test]
    fn refills_over_elapsed_time() {
        let mut bucket = TokenBucket::new(5, 2.0);

        assert!(bucket.try_acquire(5));
        assert_eq!(bucket.available(), 0);

        bucket.last_refill -= Duration::from_secs(2);

        assert_eq!(bucket.available(), 4);
        assert!(bucket.try_acquire(4));
        assert!(!bucket.try_acquire(1));
    }

    #[test]
    fn reports_wait_time_for_missing_tokens() {
        let mut bucket = TokenBucket::new(10, 5.0);

        assert!(bucket.try_acquire(10));

        let wait = bucket.time_until_available(3).unwrap();
        assert!(wait <= Duration::from_millis(600));
        assert!(wait > Duration::from_millis(500));
    }
}