package ratelimiter

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

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

// NewTokenBucket creates a limiter that refills at tokensPerSecond up to capacity.
func NewTokenBucket(tokensPerSecond float64, capacity int) *TokenBucket {
	return newTokenBucket(tokensPerSecond, capacity, time.Now)
}

func newTokenBucket(tokensPerSecond float64, capacity int, now func() time.Time) *TokenBucket {
	if tokensPerSecond <= 0 {
		panic("tokensPerSecond must be positive")
	}
	if capacity <= 0 {
		panic("capacity must be positive")
	}

	t := now()
	return &TokenBucket{
		rate:     tokensPerSecond,
		capacity: float64(capacity),
		tokens:   float64(capacity),
		last:     t,
		now:      now,
	}
}

// Allow consumes one token if available.
func (b *TokenBucket) Allow() bool {
	return b.AllowN(1)
}

// AllowN consumes n tokens if available.
func (b *TokenBucket) AllowN(n int) bool {
	if n <= 0 {
		return true
	}

	b.mu.Lock()
	defer b.mu.Unlock()

	b.refillLocked()

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

	b.tokens -= need
	return true
}

// Tokens returns the currently available token count.
func (b *TokenBucket) Tokens() float64 {
	b.mu.Lock()
	defer b.mu.Unlock()

	b.refillLocked()
	return b.tokens
}

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

	b.tokens += elapsed * b.rate
	if b.tokens > b.capacity {
		b.tokens = b.capacity
	}
	b.last = now
}

func TestTokenBucketAllowsInitialBurstThenLimits(t *testing.T) {
	now := time.Unix(0, 0)
	bucket := newTokenBucket(10, 2, func() time.Time { return now })

	if !bucket.Allow() || !bucket.Allow() {
		t.Fatal("expected initial burst to be allowed")
	}
	if bucket.Allow() {
		t.Fatal("expected empty bucket to reject request")
	}
}

func TestTokenBucketRefillsOverTime(t *testing.T) {
	now := time.Unix(0, 0)
	bucket := newTokenBucket(2, 4, func() time.Time { return now })

	if !bucket.AllowN(4) {
		t.Fatal("expected full bucket to allow four tokens")
	}
	if bucket.Allow() {
		t.Fatal("expected depleted bucket to reject request")
	}

	now = now.Add(1500 * time.Millisecond)

	if !bucket.AllowN(3) {
		t.Fatal("expected three tokens after 1.5 seconds at two tokens per second")
	}
	if bucket.Allow() {
		t.Fatal("expected bucket to be empty again")
	}
}