use std::sync::Mutex;
use std::time::Instant;

/// A thread-safe token-bucket rate limiter.
pub struct TokenBucket {
    capacity: f64,
    refill_rate: f64,
    state: Mutex<State>,
}

struct State {
    tokens: f64,
    last_refill: Instant,
}

impl TokenBucket {
    /// Creates a full bucket that refills at `tokens_per_second`.
    ///
    /// Panics if `capacity` is zero or the refill rate is not finite and positive.
    pub fn new(capacity: u64, tokens_per_second: f64) -> Self {
        assert!(capacity > 0, "capacity must be positive");
        assert!(
            tokens_per_second.is_finite() && tokens_per_second > 0.0,
            "refill rate must be finite and positive"
        );

        Self {
            capacity: capacity as f64,
            refill_rate: tokens_per_second,
            state: Mutex::new(State {
                tokens: capacity as f64,
                last_refill: Instant::now(),
            }),
        }
    }

    /// Attempts to consume `tokens`, returning `false` without consuming
    /// anything when the bucket does not contain enough.
    pub fn try_acquire(&self, tokens: u64) -> bool {
        self.try_acquire_at(tokens, Instant::now())
    }

    fn try_acquire_at(&self, tokens: u64, now: Instant) -> bool {
        let requested = tokens as f64;
        if requested > self.capacity {
            return false;
        }

        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        self.refill(&mut state, now);

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

    fn refill(&self, state: &mut State, now: Instant) {
        if let Some(elapsed) = now.checked_duration_since(state.last_refill) {
            state.tokens = (state.tokens + elapsed.as_secs_f64() * self.refill_rate)
                .min(self.capacity);
            state.last_refill = now;
        }
    }
}

fn main() {}

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

    #[test]
    fn rejects_requests_that_would_overdraw_the_bucket() {
        let bucket = TokenBucket::new(3, 1.0);

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

    #[test]
    fn refills_over_time_without_exceeding_capacity() {
        let bucket = TokenBucket::new(2, 2.0);
        let start = {
            let state = bucket.state.lock().unwrap();
            state.last_refill
        };

        assert!(bucket.try_acquire_at(2, start));
        assert!(!bucket.try_acquire_at(1, start + Duration::from_millis(499)));
        assert!(bucket.try_acquire_at(1, start + Duration::from_millis(500)));

        // Extra elapsed time fills the bucket only to its capacity.
        assert!(bucket.try_acquire_at(2, start + Duration::from_secs(10)));
        assert!(!bucket.try_acquire_at(1, start + Duration::from_secs(10)));
    }
}