package ratelimit

import (
	"context"
	"sync"
	"time"
)

// RateLimiter implements a token bucket rate limiter
type RateLimiter struct {
	requestsPerSecond int
	burstSize         int
	windowSize        time.Duration
	tokens            int
	lastRefill        time.Time
	mutex             sync.Mutex
}

// NewRateLimiter creates a new rate limiter
func NewRateLimiter(requestsPerSecond, burstSize int, windowSize time.Duration) *RateLimiter {
	return &RateLimiter{
		requestsPerSecond: requestsPerSecond,
		burstSize:         burstSize,
		tokens:            burstSize,
		lastRefill:        time.Now(),
		windowSize:        windowSize,
	}
}

// Allow checks if a request is allowed under the rate limit
func (rl *RateLimiter) Allow() bool {
	rl.mutex.Lock()
	defer rl.mutex.Unlock()

	now := time.Now()

	// Calculate tokens to add based on time elapsed
	elapsed := now.Sub(rl.lastRefill)
	tokensToAdd := int(elapsed.Seconds() * float64(rl.requestsPerSecond))

	if tokensToAdd > 0 {
		rl.tokens += tokensToAdd
		if rl.tokens > rl.burstSize {
			rl.tokens = rl.burstSize
		}
		rl.lastRefill = now
	}

	// Check if we have tokens available
	if rl.tokens > 0 {
		rl.tokens--
		return true
	}

	return false
}

// AllowWithContext checks if a request is allowed with context cancellation
func (rl *RateLimiter) AllowWithContext(ctx context.Context) bool {
	select {
	case <-ctx.Done():
		return false
	default:
		return rl.Allow()
	}
}

// Wait blocks until a request is allowed
func (rl *RateLimiter) Wait(ctx context.Context) error {
	for {
		if rl.Allow() {
			return nil
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(rl.windowSize / time.Duration(rl.requestsPerSecond)):
			// Wait for next token
		}
	}
}

// GetStats returns current rate limiter statistics
func (rl *RateLimiter) GetStats() (tokens int, lastRefill time.Time) {
	rl.mutex.Lock()
	defer rl.mutex.Unlock()
	return rl.tokens, rl.lastRefill
}