package main

import (
	"fmt"
	"time"
)

// RateLimiter holds the state for our token bucket.
type RateLimiter struct {
	bucket chan struct{}
	rate   time.Duration
}

// NewRateLimiter creates a new rate limiter with a given capacity and fill rate.
func NewRateLimiter(capacity int, fillRate time.Duration) *RateLimiter {
	l := &RateLimiter{
		bucket: make(chan struct{}, capacity),
		rate:   fillRate,
	}

	// Fill the bucket initially
	for i := 0; i < capacity; i++ {
		l.bucket <- struct{}{}
	}

	// Start a goroutine to continuously refill the bucket
	go func() {
		ticker := time.NewTicker(fillRate)
		defer ticker.Stop()
		for range ticker.C {
			select {
			case l.bucket <- struct{}{}:
				// Token added successfully
			default:
				// Bucket is full, drop token
			}
		}
	}()

	return l
}

// Allow checks if a request is allowed by consuming a token.
func (l *RateLimiter) Allow() bool {
	select {
	case <-l.bucket:
		return true
	default:
		return false
	}
}

// Wait blocks until a token is available.
func (l *RateLimiter) Wait() {
	<-l.bucket
}

func main() {
	// Allow 2 requests per second, with a burst capacity of 5
	rateLimiter := NewRateLimiter(5, 500*time.Millisecond) // 5 tokens, refill 1 token every 500ms

	fmt.Println("Testing rate limiter (Allow):")
	for i := 0; i < 10; i++ {
		if rateLimiter.Allow() {
			fmt.Printf("[%s] Request %d: ALLOWED
", time.Now().Format("15:04:05.000"), i+1)
		} else {
			fmt.Printf("[%s] Request %d: DENIED
", time.Now().Format("15:04:05.000"), i+1)
		}
		time.Sleep(200 * time.Millisecond)
	}

	fmt.Println("
Testing rate limiter (Wait):")
	for i := 0; i < 10; i++ {
		start := time.Now()
		rateLimiter.Wait() // This will block if no token is available
		duration := time.Since(start)
		fmt.Printf("[%s] Request %d: PROCESSED after waiting %s
", time.Now().Format("15:04:05.000"), i+1, duration.Round(time.Millisecond))
		time.Sleep(100 * time.Millisecond) // Simulate some quick processing after getting a token
	}
}