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

/// A simple token-bucket rate limiter.
///
/// Tokens refill continuously over time up to `capacity`. Calling `try_acquire`
/// consumes tokens when enough are available and returns `false` otherwise.
#[derive(Debug, Clone)]
pub struct TokenBucket {
    capacity: u64,
    tokens: f64,
    refill_per_second: f64,
    last_refill: Instant,
}

impl TokenBucket {
    /// Creates a bucket with `capacity` tokens and a refill rate in tokens/sec.
    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 consume one token.
    pub fn try_acquire_one(&mut self) -> bool {
        self.try_acquire(1)
    }

    /// Attempts to consume `amount` tokens.
    pub fn try_acquire(&mut self, amount: u64) -> bool {
        self.try_acquire_at(amount, Instant::now())
    }

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

    /// Returns the bucket's maximum token capacity.
    pub fn capacity(&self) -> u64 {
        self.capacity
    }

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

    fn try_acquire_at(&mut self, amount: u64, now: Instant) -> bool {
        if amount == 0 {
            return true;
        }

        if amount > self.capacity {
            return false;
        }

        self.refill(now);

        if self.tokens >= amount as f64 {
            self.tokens -= amount as f64;
            true
        } else {
            false
        }
    }

    fn refill(&mut self, now: Instant) {
        let Some(elapsed) = now.checked_duration_since(self.last_refill) else {
            return;
        };

        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 starts_full_and_rejects_when_empty() {
        let mut bucket = TokenBucket::new(3, 1.0);
        let start = bucket.last_refill;

        assert!(bucket.try_acquire_at(1, start));
        assert!(bucket.try_acquire_at(2, start));
        assert!(!bucket.try_acquire_at(1, start));
    }

    #[test]
    fn refills_over_time_without_exceeding_capacity() {
        let mut bucket = TokenBucket::new(5, 2.0);
        let start = bucket.last_refill;

        assert!(bucket.try_acquire_at(5, start));
        assert!(!bucket.try_acquire_at(1, start));

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

        assert!(bucket.try_acquire_at(5, start + Duration::from_secs(10)));
        assert!(!bucket.try_acquire_at(1, start + Duration::from_secs(10)));
    }
}