package tokenbucket

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

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

// New creates a full bucket with the given capacity and refill rate.
// It panics when capacity or tokensPerSecond is not positive.
func New(capacity int, tokensPerSecond float64) *Bucket {
	if capacity <= 0 || tokensPerSecond <= 0 {
		panic("capacity and tokensPerSecond must be positive")
	}
	return newBucket(capacity, tokensPerSecond, time.Now)
}

func newBucket(capacity int, tokensPerSecond float64, now func() time.Time) *Bucket {
	t := now()
	return &Bucket{
		rate:     tokensPerSecond,
		capacity: float64(capacity),
		tokens:   float64(capacity),
		last:     t,
		now:      now,
	}
}

// Allow reports whether one token was available and consumes it if so.
func (b *Bucket) Allow() bool {
	return b.AllowN(1)
}

// AllowN reports whether n tokens were available and consumes them atomically.
func (b *Bucket) AllowN(n int) bool {
	if n <= 0 {
		return n == 0
	}

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

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

	if float64(n) > b.tokens {
		return false
	}
	b.tokens -= float64(n)
	return true
}

func TestBucketExhaustion(t *testing.T) {
	now := time.Unix(0, 0)
	b := newBucket(2, 1, func() time.Time { return now })

	if !b.Allow() || !b.Allow() {
		t.Fatal("initial tokens should be available")
	}
	if b.Allow() {
		t.Fatal("exhausted bucket unexpectedly allowed a request")
	}
}

func TestBucketRefillAndCapacity(t *testing.T) {
	now := time.Unix(0, 0)
	b := newBucket(2, 2, func() time.Time { return now })

	if !b.AllowN(2) {
		t.Fatal("initial capacity should be available")
	}

	now = now.Add(500 * time.Millisecond)
	if !b.Allow() {
		t.Fatal("expected one refilled token")
	}
	if b.Allow() {
		t.Fatal("bucket refilled too many tokens")
	}

	now = now.Add(10 * time.Second)
	if !b.AllowN(2) {
		t.Fatal("bucket should refill to capacity")
	}
	if b.Allow() {
		t.Fatal("bucket must not refill beyond capacity")
	}
}