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

/// A thread-safe token-bucket rate limiter.
pub struct TokenBucket {
    capacity: f64,
    refill_rate: f64,
    state: Mutex<State>,
}

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

impl TokenBucket {
    /// Creates a full bucket with the given capacity and refill rate per second.
    pub fn new(capacity: u64, refill_rate: f64) -> Self {
        assert!(capacity > 0, "capacity must be positive");
        assert!(
            refill_rate.is_finite() && refill_rate >= 0.0,
            "refill rate must be finite and non-negative"
        );

        let capacity = capacity as f64;
        Self {
            capacity,
            refill_rate,
            state: Mutex::new(State {
                tokens: capacity,
                last_refill: Instant::now(),
            }),
        }
    }

    /// Attempts to consume `amount` tokens without blocking.
    pub fn try_acquire(&self, amount: u64) -> bool {
        if amount == 0 {
            return true;
        }

        let amount = amount as f64;
        if amount > self.capacity {
            return false;
        }

        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        self.refill(&mut state, Instant::now());

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

    /// Returns the currently available whole-token count.
    pub fn available_tokens(&self) -> u64 {
        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
        self.refill(&mut state, Instant::now());
        state.tokens.floor() as u64
    }

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

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

    #[test]
    fn enforces_capacity() {
        let bucket = TokenBucket::new(3, 0.0);

        assert!(bucket.try_acquire(2));
        assert!(!bucket.try_acquire(2));
        assert!(bucket.try_acquire(1));
        assert_eq!(bucket.available_tokens(), 0);
    }

    #[test]
    fn refills_without_exceeding_capacity() {
        let bucket = TokenBucket::new(5, 2.0);
        assert!(bucket.try_acquire(5));

        let mut state = bucket.state.lock().unwrap();
        let later = state.last_refill + Duration::from_secs(3);
        bucket.refill(&mut state, later);

        assert_eq!(state.tokens, 5.0);
    }
}