package tokenbucket

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

// Limiter is a concurrency-safe token-bucket rate limiter.
type Limiter struct {
	mu       sync.Mutex
	rate     float64
	capacity float64
	tokens   float64
	last     time.Time
	now      func() time.Time
}

// New creates a full bucket that refills at rate tokens per second.
func New(rate float64, capacity int) *Limiter {
	if rate <= 0 {
		panic("tokenbucket: rate must be positive")
	}
	if capacity <= 0 {
		panic("tokenbucket: capacity must be positive")
	}

	now := time.Now
	return &Limiter{
		rate:     rate,
		capacity: float64(capacity),
		tokens:   float64(capacity),
		last:     now(),
		now:      now,
	}
}

// Allow reports whether one token is available and consumes it if so.
func (l *Limiter) Allow() bool {
	return l.AllowN(1)
}

// AllowN reports whether n tokens are available and consumes them if so.
func (l *Limiter) AllowN(n int) bool {
	if n <= 0 {
		return n == 0
	}

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

	now := l.now()
	if elapsed := now.Sub(l.last).Seconds(); elapsed > 0 {
		l.tokens += elapsed * l.rate
		if l.tokens > l.capacity {
			l.tokens = l.capacity
		}
		l.last = now
	}

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

	l.tokens -= float64(n)
	return true
}

func newTestLimiter(rate float64, capacity int, now func() time.Time) *Limiter {
	return &Limiter{
		rate:     rate,
		capacity: float64(capacity),
		tokens:   float64(capacity),
		last:     now(),
		now:      now,
	}
}

func TestLimiterExhaustsBucket(t *testing.T) {
	current := time.Unix(0, 0)
	limiter := newTestLimiter(1, 2, func() time.Time { return current })

	if !limiter.Allow() || !limiter.Allow() {
		t.Fatal("expected initial tokens to be available")
	}
	if limiter.Allow() {
		t.Fatal("expected exhausted bucket to reject request")
	}
}

func TestLimiterRefillsOverTime(t *testing.T) {
	current := time.Unix(0, 0)
	limiter := newTestLimiter(2, 2, func() time.Time { return current })

	if !limiter.AllowN(2) {
		t.Fatal("expected full bucket to allow two tokens")
	}

	current = current.Add(500 * time.Millisecond)
	if !limiter.Allow() {
		t.Fatal("expected one token after half a second")
	}
	if limiter.Allow() {
		t.Fatal("expected refilled token to be consumed")
	}
}