package mcp

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestSendLimiter_AllowBurst(t *testing.T) {
	lim := newsendLimiter(5.0, 8)

	// Should allow a full burst of 8.
	for i := 0; i < 8; i++ {
		if !lim.Allow() {
			t.Fatalf("expected Allow()=true at burst %d", i)
		}
	}

	// 9th should be denied.
	if lim.Allow() {
		t.Fatal("expected Allow()=false after exhausting burst")
	}
}

func TestSendLimiter_Refill(t *testing.T) {
	lim := newsendLimiter(10.0, 2) // 10/sec, burst 2

	// Exhaust burst.
	lim.Allow()
	lim.Allow()
	if lim.Allow() {
		t.Fatal("expected Allow()=false after exhausting burst")
	}

	// Move the limiter clock back 200ms -> should refill 2 tokens.
	forceLimiterState(lim, 0, time.Now().Add(-200*time.Millisecond))
	if !lim.Allow() {
		t.Fatal("expected Allow()=true after refill")
	}
}

func TestSendLimiter_DoesNotExceedMax(t *testing.T) {
	lim := newsendLimiter(100.0, 3) // fast refill, small burst

	// Move the limiter clock back so tokens would exceed burst if uncapped.
	forceLimiterState(lim, 3, time.Now().Add(-50*time.Millisecond))

	// Should allow at most 3 (the burst cap).
	count := 0
	for lim.Allow() {
		count++
		if count > 10 {
			t.Fatal("infinite agentteam — tokens not capped")
		}
	}
	if count != 3 {
		t.Fatalf("expected burst cap of 3, got %d", count)
	}
}

func forceLimiterState(lim *sendLimiter, tokens float64, lastTime time.Time) {
	lim.mu.Lock()
	defer lim.mu.Unlock()
	lim.tokens = tokens
	lim.lastTime = lastTime
}

func TestSendLimiterMap_GetCreatesAndReuses(t *testing.T) {
	m := newsendLimiterMap(5.0, 8)

	lim1 := m.Get("agent-1")
	lim2 := m.Get("agent-1")
	if lim1 != lim2 {
		t.Fatal("expected same limiter instance for same key")
	}

	lim3 := m.Get("agent-2")
	if lim1 == lim3 {
		t.Fatal("expected different limiter instance for different key")
	}
}

func TestSendLimiterMap_Remove(t *testing.T) {
	m := newsendLimiterMap(5.0, 8)

	lim1 := m.Get("agent-1")
	// Exhaust tokens.
	for i := 0; i < 8; i++ {
		lim1.Allow()
	}
	if lim1.Allow() {
		t.Fatal("expected exhausted limiter")
	}

	// Remove and re-get: should be a fresh limiter.
	m.Remove("agent-1")
	lim2 := m.Get("agent-1")
	if !lim2.Allow() {
		t.Fatal("expected fresh limiter after Remove")
	}
}

func TestSendLimiterMap_GetRemovesStaleLimiters(t *testing.T) {
	m := newsendLimiterMap(5.0, 8)

	stale := m.Get("agent-1")
	active := m.Get("agent-2")

	m.mu.Lock()
	m.limiters["agent-1"] = sendLimiterEntry{
		limiter:  stale,
		lastUsed: time.Now().Add(-sendLimiterIdleTTL - time.Second),
	}
	m.limiters["agent-2"] = sendLimiterEntry{
		limiter:  active,
		lastUsed: time.Now(),
	}
	m.lastCleanup = time.Now().Add(-sendLimiterCleanupInterval - time.Second)
	m.mu.Unlock()

	if got := m.Get("agent-2"); got != active {
		t.Fatal("expected active limiter to be reused")
	}

	m.mu.Lock()
	_, staleExists := m.limiters["agent-1"]
	_, activeExists := m.limiters["agent-2"]
	m.mu.Unlock()
	if staleExists {
		t.Fatal("expected stale limiter to be removed")
	}
	if !activeExists {
		t.Fatal("expected active limiter to remain")
	}
}

func TestSendLimiterMap_GetBoundsLimiterCount(t *testing.T) {
	m := newsendLimiterMap(5.0, 8)
	m.maxBuckets = 2

	first := m.Get("agent-1")
	second := m.Get("agent-2")

	now := time.Now()
	m.mu.Lock()
	m.limiters["agent-1"] = sendLimiterEntry{
		limiter:  first,
		lastUsed: now.Add(-10 * time.Second),
	}
	m.limiters["agent-2"] = sendLimiterEntry{
		limiter:  second,
		lastUsed: now,
	}
	m.mu.Unlock()

	third := m.Get("agent-3")
	if third == nil {
		t.Fatal("expected limiter for new key")
	}

	m.mu.Lock()
	_, firstExists := m.limiters["agent-1"]
	gotSecond := m.limiters["agent-2"].limiter
	gotThird := m.limiters["agent-3"].limiter
	gotLen := len(m.limiters)
	m.mu.Unlock()

	if firstExists {
		t.Fatal("expected oldest limiter to be evicted")
	}
	if gotSecond != second {
		t.Fatal("expected recently used limiter to remain")
	}
	if gotThird != third {
		t.Fatal("expected new limiter to be stored")
	}
	if gotLen != 2 {
		t.Fatalf("len(limiters) = %d, want 2", gotLen)
	}
}

func TestSendLimiterKey_UsesConversationScope(t *testing.T) {
	got := sendLimiterKey("agentteam", "chat-1", "agent-1")
	want := "agentteam\x00chat-1"
	if got != want {
		t.Fatalf("sendLimiterKey() = %q, want %q", got, want)
	}

	if got := sendLimiterKey("", "", "agent-1"); got != "agent-1" {
		t.Fatalf("sendLimiterKey() fallback = %q, want agent-1", got)
	}
}

func TestSendLimiterKeySharesExactTargetAcrossAgents(t *testing.T) {
	first := sendLimiterKey("agentteam", "#ops:thread", "agent-a")
	second := sendLimiterKey("agentteam", "#ops:thread", "agent-b")
	require.Equal(t, first, second)
}
