package main

import (
    "sync"
    "time"
)

// RateLimiter implements a per-user token bucket rate limiter.
type RateLimiter struct {
    mu    sync.Mutex
    users map[string]*bucket
    rate  int // requests per minute
}

type bucket struct {
    tokens     float64
    lastRefill time.Time
}

// NewRateLimiter creates a rate limiter with the given rate (requests/minute).
func NewRateLimiter(ratePerMinute int) *RateLimiter {
    if ratePerMinute < 0 {
        ratePerMinute = 0
    }
    return &RateLimiter{
        users: make(map[string]*bucket),
        rate:  ratePerMinute,
    }
}

// Allow checks if a request from the given user ID should be allowed.
func (rl *RateLimiter) Allow(userID string) bool {
    // 0 = unlimited.
    if rl.rate == 0 {
        return true
    }
    rl.mu.Lock()
    defer rl.mu.Unlock()

    b, exists := rl.users[userID]
    now := time.Now()

    if !exists {
        b = &bucket{tokens: 1, lastRefill: now}
        rl.users[userID] = b
    }

    elapsed := now.Sub(b.lastRefill).Seconds()
    refill := elapsed * float64(rl.rate) / 60.0
    b.tokens += refill
    if b.tokens > 1 {
        b.tokens = 1
    }
    b.lastRefill = now

    if b.tokens >= 1 {
        b.tokens--
        return true
    }
    return false
}

// Reset clears all rate limiter state.
func (rl *RateLimiter) Reset() {
    rl.mu.Lock()
    defer rl.mu.Unlock()
    rl.users = make(map[string]*bucket)
}