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

/// A simple token-bucket rate limiter.
///
/// The bucket starts full. Each successful acquire removes tokens, and tokens
/// are refilled over time up to `capacity`.
#[derive(Debug, Clone)]
pub struct TokenBucket {
    capacity: f64,
    tokens: f64,
    refill_per_second: f64,
    last_refill: Instant,
}

impl TokenBucket {
    /// Creates a bucket with `capacity` tokens, refilled at `refill_per_second`.
    pub fn new(capacity: u32, 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 rate must be finite and greater than zero"
        );

        let capacity = capacity as f64;

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

    /// Creates a bucket with a custom initial token count.
    pub fn with_initial_tokens(
        capacity: u32,
        refill_per_second: f64,
        initial_tokens: u32,
    ) -> Self {
        let mut bucket = Self::new(capacity, refill_per_second);
        bucket.tokens = (initial_tokens as f64).min(bucket.capacity);
        bucket
    }

    /// Attempts to acquire one token.
    pub fn try_acquire(&mut self) -> bool {
        self.try_acquire_n(1)
    }

    /// Attempts to acquire `count` tokens.
    pub fn try_acquire_n(&mut self, count: u32) -> bool {
        assert!(count > 0, "token count must be greater than zero");

        self.refill();

        let requested = count as f64;
        if self.tokens >= requested {
            self.tokens -= requested;
            true
        } else {
            false
        }
    }

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

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

        self.refill();

        let requested = count as f64;
        if requested > self.capacity {
            return None;
        }

        if self.tokens >= requested {
            Some(Duration::ZERO)
        } else {
            let missing = requested - self.tokens;
            Some(Duration::from_secs_f64(missing / self.refill_per_second))
        }
    }

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

        if elapsed > 0.0 {
            self.tokens = (self.tokens + elapsed * self.refill_per_second).min(self.capacity);
            self.last_refill = now;
        }
    }
}

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

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

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

    #[test]
    fn refills_over_elapsed_time() {
        let mut bucket = TokenBucket::with_initial_tokens(10, 4.0, 0);

        bucket.last_refill = bucket
            .last_refill
            .checked_sub(Duration::from_millis(500))
            .unwrap();

        assert_eq!(bucket.available(), 2);
        assert!(bucket.try_acquire_n(2));
        assert!(!bucket.try_acquire());
    }
}