package ratelimit

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

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

// NewBucket creates a full bucket. rate is the number of tokens added per second.
func NewBucket(capacity int, rate float64) (*Bucket, error) {
	return newBucket(capacity, rate, time.Now)
}

func newBucket(capacity int, rate float64, now func() time.Time) (*Bucket, error) {
	if capacity <= 0 {
		return nil, errors.New("capacity must be positive")
	}
	if rate <= 0 {
		return nil, errors.New("rate must be positive")
	}

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

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

// AllowN reports whether n tokens can be consumed atomically.
func (b *Bucket) AllowN(n int) bool {
	if n <= 0 || float64(n) > b.capacity {
		return false
	}

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

	b.refill(b.now())
	if b.tokens < float64(n) {
		return false
	}

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

func (b *Bucket) refill(now time.Time) {
	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
}

type fakeClock struct {
	now time.Time
}

func (c *fakeClock) Now() time.Time {
	return c.now
}

func (c *fakeClock) Advance(d time.Duration) {
	c.now = c.now.Add(d)
}

func TestBucketRefills(t *testing.T) {
	clock := &fakeClock{now: time.Unix(0, 0)}
	bucket, err := newBucket(2, 2, clock.Now)
	if err != nil {
		t.Fatal(err)
	}

	if !bucket.AllowN(2) {
		t.Fatal("expected initial tokens to be available")
	}
	if bucket.Allow() {
		t.Fatal("expected empty bucket to reject request")
	}

	clock.Advance(500 * time.Millisecond)
	if !bucket.Allow() {
		t.Fatal("expected one token after refill")
	}
	if bucket.Allow() {
		t.Fatal("expected refill token to be consumed")
	}
}

func TestBucketDoesNotExceedCapacity(t *testing.T) {
	clock := &fakeClock{now: time.Unix(0, 0)}
	bucket, err := newBucket(3, 1, clock.Now)
	if err != nil {
		t.Fatal(err)
	}

	if !bucket.AllowN(3) {
		t.Fatal("expected initial tokens to be available")
	}

	clock.Advance(10 * time.Second)
	if !bucket.AllowN(3) {
		t.Fatal("expected bucket to refill to capacity")
	}
	if bucket.Allow() {
		t.Fatal("expected refill to be capped at capacity")
	}
}