package balancer

import (
	"context"
	"errors"
	"testing"

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

type mockSearcher struct {
	name   models.ProviderKind
	failOn int // 第几次调用失败（0 = 永不）
	calls  int
}

func (m *mockSearcher) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
	m.calls++
	if m.failOn > 0 && m.calls >= m.failOn {
		return nil, errors.New("mock failure")
	}
	return &models.SearchResult{Provider: m.name, Query: req.Query, Results: []models.ResultItem{{Title: "mock"}}}, nil
}
func (m *mockSearcher) HealthCheck(ctx context.Context) error { return nil }
func (m *mockSearcher) Name() models.ProviderKind             { return m.name }

func TestLoadBalancer_Search_Basic(t *testing.T) {
	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})

	result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Provider != models.ProviderTavily {
		t.Errorf("expected tavily, got %s", result.Provider)
	}
}

func TestLoadBalancer_Search_FallbackProvider(t *testing.T) {
	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 1})
	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})
	lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})

	result, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if string(result.Provider) != "brave" {
		t.Errorf("expected fallback to brave, got %s", result.Provider)
	}
}

func TestLoadBalancer_Search_AllExhausted(t *testing.T) {
	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily, failOn: 1})

	_, err := lb.Search(context.Background(), &models.SearchRequest{Query: "test"})
	if err == nil {
		t.Fatal("expected error when all keys exhausted")
	}
}

func TestLoadBalancer_Search_PreferredProviders(t *testing.T) {
	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
	lb.AddProvider(models.ProviderTavily, []string{"key1"}, 1, &mockSearcher{name: models.ProviderTavily})
	lb.AddProvider(models.ProviderBrave, []string{"key1"}, 1, &mockSearcher{name: models.ProviderBrave})

	result, err := lb.Search(context.Background(), &models.SearchRequest{
		Query:     "test",
		Providers: []string{"brave"},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if result.Provider != models.ProviderBrave {
		t.Errorf("expected bing (preferred), got %s", result.Provider)
	}
}

func TestLoadBalancer_ActiveKeyCount(t *testing.T) {
	lb := NewLoadBalancer(&config.BalancerConfig{RetryCount: 0})
	lb.AddProvider(models.ProviderTavily, []string{"key1", "key2"}, 1, &mockSearcher{name: models.ProviderTavily})

	counts := lb.ActiveKeyCount()
	if counts[models.ProviderTavily] != 2 {
		t.Errorf("expected 2 active keys, got %d", counts[models.ProviderTavily])
	}
}