package routing

import (
	"testing"

	"aigw/internal/catalog"
	"aigw/internal/config"
	"aigw/internal/domain"
)

func TestPlanHonorsPriorityAndProtocol(t *testing.T) {
	cfg := config.Config{
		Providers: []config.ProviderConfig{
			{ID: "openai-primary", Protocol: domain.ProtocolOpenAI, BaseURL: "https://one.test", APIKey: "one"},
			{ID: "openai-fallback", Protocol: domain.ProtocolOpenAI, BaseURL: "https://two.test", APIKey: "two"},
			{ID: "anthropic", Protocol: domain.ProtocolAnthropic, BaseURL: "https://three.test", APIKey: "three"},
		},
		Models: []config.ModelConfig{{
			ID: "public/model",
			Routes: []config.RouteConfig{
				{Provider: "openai-fallback", UpstreamModel: "fallback", Priority: 10, Weight: 1},
				{Provider: "anthropic", UpstreamModel: "claude", Priority: 0, Weight: 1},
				{Provider: "openai-primary", UpstreamModel: "primary", Priority: 0, Weight: 1},
			},
		}},
	}
	router := New(catalog.New(cfg))
	plan, err := router.Plan("public/model", domain.ProtocolOpenAI)
	if err != nil {
		t.Fatal(err)
	}
	if len(plan) != 2 || plan[0].Provider.ID != "openai-primary" || plan[1].Provider.ID != "openai-fallback" {
		t.Fatalf("unexpected plan: %+v", plan)
	}
}

func TestPlanUsesWeightsForPrimarySelection(t *testing.T) {
	cfg := config.Config{
		Providers: []config.ProviderConfig{
			{ID: "one", Protocol: domain.ProtocolOpenAI, BaseURL: "https://one.test", APIKey: "one"},
			{ID: "two", Protocol: domain.ProtocolOpenAI, BaseURL: "https://two.test", APIKey: "two"},
		},
		Models: []config.ModelConfig{{ID: "public/model", Routes: []config.RouteConfig{
			{Provider: "one", UpstreamModel: "one", Weight: 3},
			{Provider: "two", UpstreamModel: "two", Weight: 1},
		}}},
	}
	router := New(catalog.New(cfg))
	counts := map[string]int{}
	for range 8 {
		plan, err := router.Plan("public/model", domain.ProtocolOpenAI)
		if err != nil {
			t.Fatal(err)
		}
		counts[plan[0].Provider.ID]++
	}
	if counts["one"] != 6 || counts["two"] != 2 {
		t.Fatalf("unexpected weighted distribution: %+v", counts)
	}
}