package ratelimiter

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

// Limiter is a thread-safe token-bucket rate limiter.
type Limiter struct {
	mu       sync.Mutex
	rate     float64
	burst    float64
	tokens   float64
	lastFill time.Time
}

// NewLimiter creates a limiter that refills rate tokens per second,
// up to a maximum bucket size of burst.
func NewLimiter(rate float64, burst int) *Limiter {
	if rate <= 0 {
		panic("rate must be positive")
	}
	if burst <= 0 {
		panic("burst must be positive")
	}

	return &Limiter{
		rate:     rate,
		burst:    float64(burst),
		tokens:   float64(burst),
		lastFill: time.Now(),
	}
}

// Allow reports whether one token is available now.
func (l *Limiter) Allow() bool {
	return l.AllowN(1)
}

// AllowN reports whether n tokens are available now.
// If available, the tokens are consumed.
func (l *Limiter) AllowN(n int) bool {
	if n <= 0 {
		return true
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	l.refill(time.Now())

	need := float64(n)
	if need > l.tokens {
		return false
	}

	l.tokens -= need
	return true
}

// Wait blocks until one token is available or ctx is canceled.
func (l *Limiter) Wait(ctx context.Context) error {
	return l.WaitN(ctx, 1)
}

// WaitN blocks until n tokens are available or ctx is canceled.
func (l *Limiter) WaitN(ctx context.Context, n int) error {
	if n <= 0 {
		return nil
	}

	for {
		wait := l.reserveDelay(n)
		if wait == 0 {
			return nil
		}

		timer := time.NewTimer(wait)
		select {
		case <-timer.C:
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return ctx.Err()
		}
	}
}

func (l *Limiter) reserveDelay(n int) time.Duration {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.refill(time.Now())

	need := float64(n)
	if need <= l.tokens {
		l.tokens -= need
		return 0
	}

	missing := need - l.tokens
	return time.Duration((missing / l.rate) * float64(time.Second))
}

func (l *Limiter) refill(now time.Time) {
	if now.Before(l.lastFill) {
		l.lastFill = now
		return
	}

	elapsed := now.Sub(l.lastFill).Seconds()
	l.tokens += elapsed * l.rate
	if l.tokens > l.burst {
		l.tokens = l.burst
	}
	l.lastFill = now
}

func TestLimiterAllowsBurstThenRejects(t *testing.T) {
	limiter := NewLimiter(100, 2)

	if !limiter.Allow() || !limiter.Allow() {
		t.Fatal("expected initial burst tokens to be available")
	}

	if limiter.Allow() {
		t.Fatal("expected third immediate request to be rejected")
	}
}

func TestLimiterRefillsOverTime(t *testing.T) {
	limiter := NewLimiter(10, 1)

	if !limiter.Allow() {
		t.Fatal("expected first token to be available")
	}

	time.Sleep(120 * time.Millisecond)

	if !limiter.Allow() {
		t.Fatal("expected token to refill after waiting")
	}
}