package ratelimit

import (
	"sync"
	"time"
)

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

// NewTokenBucket creates a bucket with the given capacity and refill rate.
//
// Capacity is the maximum number of stored tokens. Rate is tokens per second.
func NewTokenBucket(capacity int, rate float64) *TokenBucket {
	if capacity <= 0 {
		panic("capacity must be positive")
	}
	if rate <= 0 {
		panic("rate must be positive")
	}

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

// Allow reports whether one token could be consumed immediately.
func (b *TokenBucket) Allow() bool {
	return b.AllowN(1)
}

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

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

	b.refill(time.Now())

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

	b.tokens -= need
	return true
}

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

	b.refill(time.Now())
	return b.tokens
}

func (b *TokenBucket) refill(now time.Time) {
	if now.Before(b.last) {
		b.last = now
		return
	}

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

/*
Example tests for a file named token_bucket_test.go:

func TestTokenBucketAllowsUntilEmpty(t *testing.T) {
	b := NewTokenBucket(2, 1)

	if !b.Allow() || !b.Allow() {
		t.Fatal("expected first two requests to be allowed")
	}
	if b.Allow() {
		t.Fatal("expected empty bucket to reject request")
	}
}

func TestTokenBucketRefills(t *testing.T) {
	b := NewTokenBucket(1, 20)

	if !b.Allow() {
		t.Fatal("expected initial token")
	}
	if b.Allow() {
		t.Fatal("expected empty bucket")
	}

	time.Sleep(60 * time.Millisecond)

	if !b.Allow() {
		t.Fatal("expected bucket to refill")
	}
}
*/