rate_limiters 0.1.13

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, UNIX_EPOCH};

pub use crate::traits::{RateLimiter, RateLimiterShared};

#[derive(Debug, Clone, Copy)]
pub struct TokenBucketConfig {
    pub capacity: u32,
    pub refill_rate: u32,
}

pub struct TokenBucket {
    capacity: u32,
    tokens: u32,
    refill_rate: u32,
    last_refill: Instant,
}

impl TokenBucket {
    pub fn new(config: TokenBucketConfig) -> Self {
        Self {
            capacity: config.capacity,
            tokens: config.capacity,
            refill_rate: config.refill_rate,
            last_refill: Instant::now(),
        }
    }
}

impl RateLimiter for TokenBucket {
    fn refresh(&mut self) {
        if self.refill_rate == 0 {
            return;
        }
        let elapsed = Instant::now().duration_since(self.last_refill);
        let new_tokens = (elapsed.as_secs_f64() * self.refill_rate as f64).floor() as u32;

        if new_tokens > 0 {
            self.tokens = std::cmp::min(self.capacity, self.tokens + new_tokens);
            // Advance last_refill by exactly the time the new tokens represent
            // so any sub-token fraction is preserved (prevents drift across calls).
            self.last_refill +=
                Duration::from_secs_f64(new_tokens as f64 / self.refill_rate as f64);
        }
    }

    fn try_acquire(&mut self, tokens: u32) -> bool {
        self.refresh();
        if self.tokens >= tokens {
            self.tokens -= tokens;
            true
        } else {
            false
        }
    }

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

    fn get_remaining(&self) -> u32 {
        self.tokens
    }

    fn get_used(&self) -> u32 {
        self.capacity - self.tokens
    }

    fn get_reset(&self) -> u64 {
        let now = std::time::SystemTime::now();
        if self.refill_rate == 0 {
            return u64::MAX;
        }
        let refill_secs = (self.capacity - self.tokens) as f64 / self.refill_rate as f64;
        let reset_time = now + std::time::Duration::from_secs_f64(refill_secs);
        reset_time.duration_since(UNIX_EPOCH).unwrap().as_secs()
    }
}

pub struct TokenBucketShared {
    inner: Arc<Mutex<TokenBucket>>,
}

impl TokenBucketShared {
    pub fn new(config: TokenBucketConfig) -> Self {
        Self {
            inner: Arc::new(Mutex::new(TokenBucket::new(config))),
        }
    }
}

impl RateLimiterShared for TokenBucketShared {
    fn refresh(&self) {
        let mut bucket = self.inner.lock().unwrap();
        bucket.refresh()
    }

    fn try_acquire(&self, tokens: u32) -> bool {
        let mut bucket = self.inner.lock().unwrap();
        bucket.try_acquire(tokens)
    }

    fn get_limit(&self) -> u32 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_limit()
    }

    fn get_remaining(&self) -> u32 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_remaining()
    }

    fn get_used(&self) -> u32 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_used()
    }

    fn get_reset(&self) -> u64 {
        let bucket = self.inner.lock().unwrap();
        bucket.get_reset()
    }
}