package ratelimit

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

// Limiter is a goroutine-safe token-bucket rate limiter.
// Tokens refill continuously at rate tokens per second, up to burst capacity.
type Limiter struct {
	mu     sync.Mutex
	rate   float64
	burst  float64
	tokens float64
	last   time.Time
	now    func() time.Time
}

// NewLimiter creates a limiter with the given token rate and burst size.
// A new limiter starts full, allowing an initial 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),
		last:   time.Now(),
		now:    time.Now,
	}
}

// Allow reports whether one token can be consumed immediately.
func (l *Limiter) Allow() bool {
	return l.AllowN(1)
}

// AllowN reports whether n tokens can be consumed immediately.
func (l *Limiter) AllowN(n int) bool {
	if n <= 0 {
		return true
	}

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

	l.refillLocked()

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

	l.tokens -= need
	return true
}

// Tokens returns the current number of available tokens.
func (l *Limiter) Tokens() float64 {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.refillLocked()
	return l.tokens
}

func (l *Limiter) refillLocked() {
	now := l.now()
	elapsed := now.Sub(l.last).Seconds()
	if elapsed <= 0 {
		return
	}

	l.tokens += elapsed * l.rate
	if l.tokens > l.burst {
		l.tokens = l.burst
	}
	l.last = now
}

func TestLimiterAllowsInitialBurst(t *testing.T) {
	limiter := NewLimiter(1, 2)

	if !limiter.Allow() {
		t.Fatal("first token should be allowed")
	}
	if !limiter.Allow() {
		t.Fatal("second token should be allowed")
	}
	if limiter.Allow() {
		t.Fatal("third token should be denied")
	}
}

func TestLimiterRefillsOverTime(t *testing.T) {
	start := time.Unix(0, 0)
	limiter := NewLimiter(2, 2)
	limiter.now = func() time.Time { return start }
	limiter.last = start

	if !limiter.AllowN(2) {
		t.Fatal("initial burst should be allowed")
	}

	start = start.Add(500 * time.Millisecond)
	if !limiter.Allow() {
		t.Fatal("one token should refill after half a second at two tokens per second")
	}
	if limiter.Allow() {
		t.Fatal("only one token should have refilled")
	}
}