package ratelimiter

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

// Limiter is a token-bucket rate limiter.
// It refills at a fixed rate up to its configured capacity.
type Limiter struct {
	mu       sync.Mutex
	rate     float64
	capacity float64
	tokens   float64
	last     time.Time
}

// New creates a limiter that refills rate tokens per second, up to capacity.
func New(rate float64, capacity int) *Limiter {
	if rate <= 0 {
		panic("rate must be positive")
	}
	if capacity <= 0 {
		panic("capacity must be positive")
	}

	now := time.Now()
	return &Limiter{
		rate:     rate,
		capacity: float64(capacity),
		tokens:   float64(capacity),
		last:     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 true
	}

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

	l.refill(time.Now())

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

	l.tokens -= need
	return true
}

// Available returns the current number of whole tokens available.
func (l *Limiter) Available() int {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.refill(time.Now())
	return int(l.tokens)
}

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

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

func TestLimiterAllowsInitialBurst(t *testing.T) {
	limiter := New(1, 2)

	if !limiter.Allow() {
		t.Fatal("expected first token to be allowed")
	}
	if !limiter.Allow() {
		t.Fatal("expected second token to be allowed")
	}
	if limiter.Allow() {
		t.Fatal("expected third token to be denied")
	}
}

func TestLimiterRefills(t *testing.T) {
	limiter := New(100, 1)

	if !limiter.Allow() {
		t.Fatal("expected initial token to be allowed")
	}
	if limiter.Allow() {
		t.Fatal("expected empty bucket to deny request")
	}

	time.Sleep(15 * time.Millisecond)

	if !limiter.Allow() {
		t.Fatal("expected token after refill")
	}
}