package main

import (
	"fmt"
	"time"
)

// RateLimiter uses a token bucket algorithm.
type RateLimiter struct {
	tokens chan struct{} // The bucket of tokens
	burst  int           // Maximum burst size
	rate   time.Duration // Time between token refills
}

// NewRateLimiter creates a new rate limiter.
// 'burst' is the maximum number of requests allowed instantly.
// 'rate' is how often a new token is added.
func NewRateLimiter(burst int, rate time.Duration) *RateLimiter {
	rl := &RateLimiter{
		tokens: make(chan struct{}, burst),
		burst:  burst,
		rate:   rate,
	}

	// Fill the bucket initially
	for i := 0; i < burst; i++ {
		rl.tokens <- struct{}{}
	}

	// Refill tokens periodically
	go func() {
		ticker := time.NewTicker(rate)
		defer ticker.Stop()
		for range ticker.C {
			select {
			case rl.tokens <- struct{}{}:
				// Token added
			default:
				// Bucket is full, drop token
			}
		}
	}()

	return rl
}

// Allow checks if a request is allowed.
// It blocks if no tokens are available until one is, or returns immediately.
func (rl *RateLimiter) Allow() {
	<-rl.tokens // Block until a token is available
}

func main() {
	// Allow 5 requests instantly, then 1 request per second
	limiter := NewRateLimiter(5, time.Second/1)

	fmt.Println("Starting rate limiter test (5 burst, 1 req/sec after burst)...")

	for i := 1; i <= 15; i++ {
		limiter.Allow() // This will block if no tokens are available
		fmt.Printf("Request %d processed at %s
", i, time.Now().Format("15:04:05.000"))
		if i == 5 {
			fmt.Println("--- Burst limit hit, subsequent requests will be rate-limited ---")
		}
		if i == 10 {
			time.Sleep(3 * time.Second) // Simulate a pause, bucket should refill
			fmt.Println("--- Paused for 3 seconds, bucket should have refilled ---")
		}
	}
}