use std::time::Instant;

/// A token-bucket rate limiter.
///
/// Tokens are replenished continuously at `refill_rate` tokens per second,
/// up to `capacity`. Callers may consume one or more tokens at a time.
#[derive(Debug)]
pub struct TokenBucket {
    capacity: f64,
    tokens: f64,
    refill_rate: f64,
    last_refill: Instant,
}

impl TokenBucket {
    /// Creates a full bucket.
    ///
    /// Both `capacity` and `refill_rate` must be finite and positive.
    pub fn new(capacity: f64, refill_rate: f64) -> Result<Self, &'static str> {
        if !capacity.is_finite() || capacity <= 0.0 {
            return Err("capacity must be finite and positive");
        }
        if !refill_rate.is_finite() || refill_rate <= 0.0 {
            return Err("refill rate must be finite and positive");
        }

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

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

    /// Returns the currently available tokens after applying any refill.
    pub fn available_tokens(&mut self) -> f64 {
        self.refill(Instant::now());
        self.tokens
    }

    fn try_acquire_at(&mut self, amount: f64, now: Instant) -> bool {
        if !amount.is_finite() || amount <= 0.0 || amount > self.capacity {
            return false;
        }

        self.refill(now);

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

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

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

    #[test]
    fn rejects_requests_when_empty() {
        let mut bucket = TokenBucket::new(2.0, 1.0).unwrap();
        let now = bucket.last_refill;

        assert!(bucket.try_acquire_at(2.0, now));
        assert!(!bucket.try_acquire_at(1.0, now));
    }

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

        assert!(bucket.try_acquire_at(3.0, start));
        assert!(bucket.try_acquire_at(
            2.0,
            start + Duration::from_secs(1)
        ));
        assert!(!bucket.try_acquire_at(
            1.0,
            start + Duration::from_secs(1)
        ));

        bucket.refill(start + Duration::from_secs(10));
        assert_eq!(bucket.tokens, 3.0);
    }
}