use std::time::Instant;

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

impl TokenBucket {
    pub(super) fn new(max_bytes_per_sec: u64, burst_bytes: u64) -> Self {
        Self {
            capacity: burst_bytes as f64,
            tokens: burst_bytes as f64,
            refill_per_sec: max_bytes_per_sec as f64,
            last_refill: Instant::now(),
        }
    }

    pub(super) fn try_consume(&mut self, amount: u64) -> bool {
        self.refill();
        let amount_f = amount as f64;
        if self.tokens < amount_f {
            return false;
        }
        self.tokens -= amount_f;
        true
    }

    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        self.last_refill = now;

        let added = elapsed * self.refill_per_sec;
        self.tokens = (self.tokens + added).min(self.capacity);
    }
}
