package ratelimit

import (
	"context"
	"testing"
	"time"
)

func TestWaitUnlimited(t *testing.T) {
	l := New(0, 0)
	start := time.Now()
	for i := 0; i < 100; i++ {
		if err := l.Wait(context.Background(), "direct"); err != nil {
			t.Fatal(err)
		}
	}
	if elapsed := time.Since(start); elapsed > time.Second {
		t.Fatalf("unlimited limiter blocked for %v", elapsed)
	}
}

func TestWaitHonorsCanceledContext(t *testing.T) {
	l := New(1, 0)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if err := l.Wait(ctx, "direct"); err == nil {
		t.Fatal("canceled context must surface an error")
	}
}

func TestWaitPacesGlobally(t *testing.T) {
	// rate 4/s -> burst 5; requests 6 and 7 must wait for token refill.
	l := New(4, 0)
	start := time.Now()
	for i := 0; i < 7; i++ {
		if err := l.Wait(context.Background(), "direct"); err != nil {
			t.Fatal(err)
		}
	}
	if elapsed := time.Since(start); elapsed < 300*time.Millisecond {
		t.Fatalf("global bucket did not pace requests: %v", elapsed)
	}
}

func TestWaitKeysAreIndependent(t *testing.T) {
	// Key buckets must not drain each other: exhausting "a" leaves "b" free.
	l := New(0, 2)
	for i := 0; i < 3; i++ {
		if err := l.Wait(context.Background(), "a"); err != nil {
			t.Fatal(err)
		}
	}
	start := time.Now()
	for i := 0; i < 3; i++ {
		if err := l.Wait(context.Background(), "b"); err != nil {
			t.Fatal(err)
		}
	}
	if elapsed := time.Since(start); elapsed > 200*time.Millisecond {
		t.Fatalf("key bucket 'b' was throttled by 'a': %v", elapsed)
	}
}