Implement a Token Bucket Rate Limiter

package ratelimit

import (
    "sync"
    "time"
)

type TokenBucket struct {
    capacity     float64       // Maximum tokens
    tokens       float64       // Current tokens
    refillRate   float64       // Tokens per second
    lastRefill   time.Time     // Last refill time
    mu           sync.Mutex
}

func NewTokenBucket(capacity, refillRate float64) *TokenBucket {
    return &TokenBucket{
        capacity:   capacity,
        tokens:     capacity, // Start full
        refillRate: refillRate,
        lastRefill: time.Now(),
    }
}

func (tb *TokenBucket) Allow() bool {
    return tb.AllowN(1)
}

func (tb *TokenBucket) AllowN(n float64) bool {
    tb.mu.Lock()
    defer tb.mu.Unlock()
    
    tb.refill()
    
    if tb.tokens >= n {
        tb.tokens -= n
        return true
    }
    return false
}

func (tb *TokenBucket) refill() {
    now := time.Now()
    elapsed := now.Sub(tb.lastRefill).Seconds()
    tb.lastRefill = now
    
    tb.tokens += elapsed * tb.refillRate
    if tb.tokens > tb.capacity {
        tb.tokens = tb.capacity
    }
}

// Wait blocks until a token is available
func (tb *TokenBucket) Wait() {
    for {
        if tb.Allow() {
            return
        }
        time.Sleep(time.Millisecond * 10)
    }
}

type RateLimiter struct {
    buckets map[string]*TokenBucket
    mu      sync.RWMutex
    
    capacity   float64
    refillRate float64
}

func NewRateLimiter(capacity, refillRate float64) *RateLimiter {
    return &RateLimiter{
        buckets:    make(map[string]*TokenBucket),
        capacity:   capacity,
        refillRate: refillRate,
    }
}

func (rl *RateLimiter) Allow(userID string) bool {
    rl.mu.RLock()
    bucket, exists := rl.buckets[userID]
    rl.mu.RUnlock()
    
    if !exists {
        rl.mu.Lock()
        // Double-check after acquiring write lock
        if bucket, exists = rl.buckets[userID]; !exists {
            bucket = NewTokenBucket(rl.capacity, rl.refillRate)
            rl.buckets[userID] = bucket
        }
        rl.mu.Unlock()
    }
    
    return bucket.Allow()
}

func RateLimitMiddleware(limiter *RateLimiter) gin.HandlerFunc {
    return func(c *gin.Context) {
        userID := c.GetHeader("X-User-ID")
        if userID == "" {
            userID = c.ClientIP()
        }
        
        if !limiter.Allow(userID) {
            c.JSON(429, gin.H{
                "error": "Rate limit exceeded",
                "retry_after": "1s",
            })
            c.Abort()
            return
        }
        
        c.Next()
    }
}