package balancer

import (
	"testing"
	"time"

	"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)

func TestKeyPool_Select_Basic(t *testing.T) {
	kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 5, 30*time.Second)
	ks := kp.Select()
	if ks == nil {
		t.Fatal("expected to select a key")
	}
	if ks.Key != "key1" && ks.Key != "key2" {
		t.Errorf("unexpected key: %s", ks.Key)
	}
}

func TestKeyPool_Select_SkipsDead(t *testing.T) {
	kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2"}, 1, 2, 1*time.Hour)
	ks1 := kp.Select()

	// Fail twice to trigger circuit breaker
	ks1.RecordFailure()
	ks1.RecordFailure()

	selected := kp.Select()
	if selected == nil {
		t.Fatal("expected fallback to second key")
	}
	if selected.Key == ks1.Key {
		t.Error("should not select the failed key")
	}
}

func TestKeyState_CombinedScore(t *testing.T) {
	ks := newKeyState(models.ProviderTavily, "test-key", 2, 5, 30*time.Second)
	score := ks.CombinedScore()
	if score <= 0 {
		t.Errorf("expected positive score, got %f", score)
	}

	// After recording a slow latency, score should drop
	ks.RecordSuccess(2 * time.Second)
	slowScore := ks.CombinedScore()
	if slowScore >= score {
		t.Error("expected slow latency to reduce score")
	}
}

func TestCircuitBreaker_OpenAndRecover(t *testing.T) {
	cb := NewCircuitBreaker(3, 50*time.Millisecond)

	// Should be closed initially
	if !cb.Allow() {
		t.Error("expected allow initially")
	}

	// Fail 3 times to open
	cb.Failure()
	cb.Failure()
	cb.Failure()
	if cb.Allow() {
		t.Error("expected circuit open after 3 failures")
	}

	// Wait for cooldown
	time.Sleep(60 * time.Millisecond)
	if !cb.Allow() {
		t.Error("expected allow after cooldown")
	}

	// Half-open: success should close
	cb.Success()
	if cb.State() != StateClosed {
		t.Error("expected circuit closed after half-open success")
	}
}

func TestCircuitBreaker_HalfOpenFailure(t *testing.T) {
	cb := NewCircuitBreaker(3, 50*time.Millisecond)

	cb.Failure()
	cb.Failure()
	cb.Failure()
	if cb.Allow() {
		t.Error("expected open")
	}

	time.Sleep(60 * time.Millisecond)
	if !cb.Allow() {
		t.Error("expected half-open allow")
	}

	// Half-open: failure should re-open immediately
	cb.Failure()
	if cb.State() != StateOpen {
		t.Error("expected open after half-open failure")
	}
}

func TestKeyPool_ActiveCount(t *testing.T) {
	kp := NewKeyPool(models.ProviderTavily, []string{"key1", "key2", "key3"}, 1, 3, 30*time.Second)
	if kp.TotalCount() != 3 {
		t.Errorf("expected 3 total, got %d", kp.TotalCount())
	}
	if kp.ActiveCount() != 3 {
		t.Errorf("expected 3 active, got %d", kp.ActiveCount())
	}
}

func TestKeyPool_AllKeysDead(t *testing.T) {
	kp := NewKeyPool(models.ProviderTavily, []string{"key1"}, 1, 1, 1*time.Hour)
	ks := kp.Select()
	ks.RecordFailure()

	selected := kp.Select()
	if selected != nil {
		t.Error("expected nil when all keys dead")
	}
}