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

/// A simple token-bucket rate limiter.
///
/// The bucket starts full. Each successful acquisition removes tokens, and
/// tokens are replenished 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 and a refill rate of
    /// `refill_tokens` every `refill_period`.
    pub fn new(capacity: u32, refill_tokens: u32, refill_period: Duration) -> Self {
        assert!(capacity > 0, "capacity must be greater than zero");
        assert!(refill_tokens > 0, "refill_tokens must be greater than zero");
        assert!(!refill_period.is_zero(), "refill_period must be non-zero");

        Self {
            capacity: capacity as f64,
            tokens: capacity as f64,
            refill_per_second: refill_tokens as f64 / refill_period.as_secs_f64(),
            last_refill: Instant::now(),
        }
    }

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

    /// Attempts to acquire `n` tokens immediately.
    pub fn try_acquire_n(&mut self, n: u32) -> bool {
        self.try_acquire_n_at(n, Instant::now())
    }

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

    /// Returns how long until `n` tokens can be acquired.
    pub fn wait_time_for(&mut self, n: u32) -> Option<Duration> {
        self.wait_time_for_at(n, Instant::now())
    }

    fn try_acquire_n_at(&mut self, n: u32, now: Instant) -> bool {
        if n == 0 {
            return true;
        }

        let requested = n as f64;
        if requested > self.capacity {
            return false;
        }

        self.refill_at(now);

        if self.tokens >= requested {
            self.tokens -= requested;
            true
        } else {
            false
        }
    }

    fn wait_time_for_at(&mut self, n: u32, now: Instant) -> Option<Duration> {
        if n == 0 {
            return Some(Duration::ZERO);
        }

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

        self.refill_at(now);

        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_at(&mut self, now: Instant) {
        if now <= self.last_refill {
            return;
        }

        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        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_tokens_until_empty() {
        let mut bucket = TokenBucket::new(2, 1, Duration::from_secs(1));

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

    #[test]
    fn refills_over_time_without_exceeding_capacity() {
        let mut bucket = TokenBucket::new(3, 1, Duration::from_secs(1));
        let start = Instant::now();

        assert!(bucket.try_acquire_n_at(3, start));
        assert!(!bucket.try_acquire_at_one(start));

        assert!(bucket.try_acquire_n_at(1, start + Duration::from_secs(1)));
        assert!(!bucket.try_acquire_n_at(3, start + Duration::from_secs(2)));

        bucket.refill_at(start + Duration::from_secs(10));
        assert_eq!(bucket.available(), 3);
    }

    trait TestAcquireOne {
        fn try_acquire_at_one(&mut self, at: Instant) -> bool;
    }

    impl TestAcquireOne for TokenBucket {
        fn try_acquire_at_one(&mut self, at: Instant) -> bool {
            self.try_acquire_n_at(1, at)
        }
    }
}