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

#[derive(Debug, Clone)]
pub struct TokenBucket {
    capacity: f64,
    tokens: f64,
    refill_per_second: f64,
    last_refill: Instant,
}

impl TokenBucket {
    /// Creates a bucket that starts full.
    pub fn new(capacity: u32, refill_per_second: f64) -> Self {
        Self::new_at(capacity, refill_per_second, Instant::now())
    }

    /// Creates a bucket using an explicit clock value, useful for deterministic tests.
    pub fn new_at(capacity: u32, refill_per_second: f64, now: Instant) -> Self {
        assert!(capacity > 0, "capacity must be positive");
        assert!(
            refill_per_second.is_finite() && refill_per_second > 0.0,
            "refill rate must be finite and positive"
        );

        Self {
            capacity: capacity as f64,
            tokens: capacity as f64,
            refill_per_second,
            last_refill: now,
        }
    }

    pub fn capacity(&self) -> u32 {
        self.capacity as u32
    }

    pub fn refill_per_second(&self) -> f64 {
        self.refill_per_second
    }

    /// Returns the currently available whole tokens after applying elapsed refill.
    pub fn available(&mut self) -> u32 {
        self.available_at(Instant::now())
    }

    pub fn available_at(&mut self, now: Instant) -> u32 {
        self.refill(now);
        self.tokens.floor() as u32
    }

    pub fn try_acquire(&mut self) -> bool {
        self.try_acquire_n_at(1, Instant::now())
    }

    pub fn try_acquire_n(&mut self, amount: u32) -> bool {
        self.try_acquire_n_at(amount, Instant::now())
    }

    /// Attempts to spend `amount` tokens at the given time.
    pub fn try_acquire_n_at(&mut self, amount: u32, now: Instant) -> bool {
        assert!(amount > 0, "amount must be positive");

        self.refill(now);

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

    /// Returns how long until `amount` tokens are available, assuming no other spending.
    pub fn wait_time_for(&mut self, amount: u32) -> Option<Duration> {
        self.wait_time_for_at(amount, Instant::now())
    }

    pub fn wait_time_for_at(&mut self, amount: u32, now: Instant) -> Option<Duration> {
        assert!(amount > 0, "amount must be positive");

        if amount as f64 > self.capacity {
            return None;
        }

        self.refill(now);

        let missing = amount as f64 - self.tokens;
        if missing <= 0.0 {
            Some(Duration::ZERO)
        } else {
            Some(Duration::from_secs_f64(missing / self.refill_per_second))
        }
    }

    fn refill(&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 spends_until_empty() {
        let start = Instant::now();
        let mut bucket = TokenBucket::new_at(3, 1.0, start);

        assert!(bucket.try_acquire_n_at(2, start));
        assert_eq!(bucket.available_at(start), 1);
        assert!(bucket.try_acquire_n_at(1, start));
        assert!(!bucket.try_acquire_n_at(1, start));
    }

    #[test]
    fn refills_over_time_up_to_capacity() {
        let start = Instant::now();
        let mut bucket = TokenBucket::new_at(5, 2.0, start);

        assert!(bucket.try_acquire_n_at(5, start));
        assert_eq!(bucket.available_at(start + Duration::from_millis(500)), 1);
        assert_eq!(bucket.available_at(start + Duration::from_secs(10)), 5);
    }

    #[test]
    fn reports_wait_time() {
        let start = Instant::now();
        let mut bucket = TokenBucket::new_at(4, 2.0, start);

        assert!(bucket.try_acquire_n_at(4, start));
        assert_eq!(
            bucket.wait_time_for_at(3, start),
            Some(Duration::from_secs_f64(1.5))
        );
        assert_eq!(bucket.wait_time_for_at(5, start), None);
    }
}