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

/// Token bucket rate limiter supporting sustained and burst rates.
pub struct TokenBucket {
    sustained_rate: u32,
    burst_capacity: u32,
    inner: Mutex<TokenBucketInner>,
}

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

impl TokenBucket {
    pub fn new(sustained_per_min: u32, burst: u32) -> Self {
        Self {
            sustained_rate: sustained_per_min,
            burst_capacity: burst,
            inner: Mutex::new(TokenBucketInner {
                tokens: burst as f64,
                last_refill: Instant::now(),
            }),
        }
    }

    /// Try to acquire one token. Returns true if successful.
    pub fn try_acquire(&self) -> bool {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let now = Instant::now();
        let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
        let refill = elapsed * (self.sustained_rate as f64 / 60.0);
        inner.tokens = (inner.tokens + refill).min(self.burst_capacity as f64);
        inner.last_refill = now;

        if inner.tokens >= 1.0 {
            inner.tokens -= 1.0;
            true
        } else {
            false
        }
    }

    /// Current number of available tokens.
    pub fn available(&self) -> f64 {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.tokens
    }

    /// Sustained rate (tokens per minute).
    pub fn sustained_rate(&self) -> u32 {
        self.sustained_rate
    }
}

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

    #[test]
    fn token_bucket_allows_burst() {
        let bucket = TokenBucket::new(10, 5);
        for _ in 0..5 {
            assert!(bucket.try_acquire());
        }
        // 6th should fail (burst exhausted)
        assert!(!bucket.try_acquire());
    }
}
