package metrics

import (
	"testing"

	"github.com/prometheus/client_golang/prometheus"

	controllerPool "proxy-pool/internal/controller/pool"
)

func TestCapacityCollectorAggregatesActiveUpstreamsWithoutLabels(t *testing.T) {
	registry := prometheus.NewRegistry()
	collector, err := NewCapacityCollector(registry)
	if err != nil {
		t.Fatalf("NewCapacityCollector() error = %v", err)
	}
	collector.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "provider-a", SourceID: "term-a", Result: controllerPool.CapacityReadSuccess,
		Managed: 2, AvailableSlots: 5, EffectiveSlots: 7, PendingExpected: 2,
	})
	collector.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "provider-b", SourceID: "term-b", Result: controllerPool.CapacityReadSuccess,
		Managed: 3, AvailableSlots: 4, EffectiveSlots: 6, PendingExpected: 1,
	})
	collector.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "provider-a", SourceID: "term-a", Result: controllerPool.CapacityReadError,
	})
	collector.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "client-a", SourceID: "term-client", Result: "invalid", Managed: 100,
	})

	assertMetricValue(t, registry, "proxy_pool_controller_capacity_inventory_reads_total", map[string]string{"result": "success"}, 2)
	assertMetricValue(t, registry, "proxy_pool_controller_capacity_inventory_reads_total", map[string]string{"result": "error"}, 1)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_managed_proxies", 5)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_available_slots", 9)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_effective_slots", 13)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_pending_expected_proxies", 3)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_active_upstreams", 2)

	collector.RemoveCapacityUpstream("provider-a", "term-a")
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_managed_proxies", 3)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_available_slots", 4)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_effective_slots", 6)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_pending_expected_proxies", 1)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_active_upstreams", 1)
}

func TestNewCapacityCollectorReusesRegisteredCollectors(t *testing.T) {
	registry := prometheus.NewRegistry()
	first, err := NewCapacityCollector(registry)
	if err != nil {
		t.Fatalf("first NewCapacityCollector() error = %v", err)
	}
	second, err := NewCapacityCollector(registry)
	if err != nil {
		t.Fatalf("second NewCapacityCollector() error = %v", err)
	}
	first.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "provider-a", SourceID: "term-a", Result: controllerPool.CapacityReadSuccess,
		Managed: 1, AvailableSlots: 2, EffectiveSlots: 3,
	})
	second.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "provider-b", SourceID: "term-b", Result: controllerPool.CapacityReadSuccess,
		Managed: 2, AvailableSlots: 3, EffectiveSlots: 5, PendingExpected: 1,
	})

	assertMetricValue(t, registry, "proxy_pool_controller_capacity_inventory_reads_total", map[string]string{"result": "success"}, 2)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_managed_proxies", 3)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_available_slots", 5)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_effective_slots", 8)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_pending_expected_proxies", 1)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_active_upstreams", 2)
}

func TestCapacityCollectorKeepsNewTermWhenOldTermStops(t *testing.T) {
	registry := prometheus.NewRegistry()
	collector, err := NewCapacityCollector(registry)
	if err != nil {
		t.Fatalf("NewCapacityCollector() error = %v", err)
	}
	collector.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "provider-a", SourceID: "old-term", Result: controllerPool.CapacityReadSuccess,
		Managed: 1, AvailableSlots: 2, EffectiveSlots: 3,
	})
	collector.ObserveCapacity(controllerPool.CapacityObservation{
		UpstreamID: "provider-a", SourceID: "new-term", Result: controllerPool.CapacityReadSuccess,
		Managed: 4, AvailableSlots: 5, EffectiveSlots: 6, PendingExpected: 1,
	})
	collector.RemoveCapacityUpstream("provider-a", "old-term")

	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_managed_proxies", 4)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_available_slots", 5)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_effective_slots", 6)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_pending_expected_proxies", 1)
	assertGaugeValue(t, registry, "proxy_pool_controller_capacity_active_upstreams", 1)
}

func assertGaugeValue(t *testing.T, registry *prometheus.Registry, name string, want float64) {
	t.Helper()
	metrics, err := registry.Gather()
	if err != nil {
		t.Fatalf("Gather() = %v", err)
	}
	for _, family := range metrics {
		if family.GetName() != name {
			continue
		}
		for _, metric := range family.GetMetric() {
			if len(metric.GetLabel()) == 0 && metric.GetGauge().GetValue() == want {
				return
			}
		}
	}
	t.Fatalf("gauge %s value=%v was not found", name, want)
}