//! Token-bucket rate limiter.
//!
//! Enforces a maximum number of requests per time window using the
//! token bucket algorithm. Tokens refill continuously; each request
//! consumes one token. When the bucket is empty, callers can either
//! poll with `tryAcquire` or block with `acquire`.

const std = @import("std");

/// Maximum tokens (requests) in the bucket
max_tokens: u32,
/// Current available tokens
tokens: f64,
/// Tokens added per nanosecond
refill_rate_per_ns: f64,
/// Last time tokens were refilled
last_refill: i128,

const RateLimiter = @This();

/// Create a rate limiter.
/// `max_per_window` is the max requests allowed in `window_ns` nanoseconds.
pub fn init(max_per_window: u32, window_ns: u64) RateLimiter {
    return .{
        .max_tokens = max_per_window,
        .tokens = @floatFromInt(max_per_window),
        .refill_rate_per_ns = @as(f64, @floatFromInt(max_per_window)) / @as(f64, @floatFromInt(window_ns)),
        .last_refill = std.time.nanoTimestamp(),
    };
}

/// Convenience: N requests per minute
pub fn perMinute(n: u32) RateLimiter {
    return init(n, std.time.ns_per_min);
}

/// Convenience: N requests per day
pub fn perDay(n: u32) RateLimiter {
    return init(n, std.time.ns_per_day);
}

/// Try to acquire a token. Returns true if granted, false if rate-limited.
/// Caller should sleep and retry if false.
pub fn tryAcquire(self: *RateLimiter) bool {
    self.refill();
    if (self.tokens >= 1.0) {
        self.tokens -= 1.0;
        return true;
    }
    return false;
}

/// Acquire a token, blocking (sleeping) until one is available.
pub fn acquire(self: *RateLimiter) void {
    while (!self.tryAcquire()) {
        // Sleep for the time needed to generate 1 token
        const wait_ns: u64 = @intFromFloat(1.0 / self.refill_rate_per_ns);
        std.Thread.sleep(wait_ns);
    }
}

/// Returns estimated wait time in nanoseconds until a token is available.
/// Returns 0 if a token is available now.
pub fn estimateWaitNs(self: *RateLimiter) u64 {
    self.refill();
    if (self.tokens >= 1.0) return 0;
    const deficit = 1.0 - self.tokens;
    return @intFromFloat(deficit / self.refill_rate_per_ns);
}

fn refill(self: *RateLimiter) void {
    const now = std.time.nanoTimestamp();
    const elapsed = now - self.last_refill;
    if (elapsed <= 0) return;

    const new_tokens = @as(f64, @floatFromInt(elapsed)) * self.refill_rate_per_ns;
    self.tokens = @min(self.tokens + new_tokens, @as(f64, @floatFromInt(self.max_tokens)));
    self.last_refill = now;
}

test "rate limiter basic" {
    var rl = RateLimiter.perMinute(60);
    // Should have full bucket initially
    try std.testing.expect(rl.tryAcquire());
}

test "rate limiter exhaustion" {
    var rl = RateLimiter.init(2, std.time.ns_per_s);
    try std.testing.expect(rl.tryAcquire());
    try std.testing.expect(rl.tryAcquire());
    // Bucket should be empty now
    try std.testing.expect(!rl.tryAcquire());
}