package loadbalancer

import (
	"testing"

	"github.com/owen/vibeproxy/internal/models"
)

func TestRoundRobin(t *testing.T) {
	b1 := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckNone})
	b2 := NewBackend(models.Backend{Host: "b", Enabled: true, HealthCheckType: models.HealthCheckNone})
	pool := NewPool("round_robin", []*Backend{b1, b2})

	seen := map[string]int{}
	for i := 0; i < 10; i++ {
		b, err := pool.Pick()
		if err != nil {
			t.Fatal(err)
		}
		seen[b.Model.Host]++
	}
	if seen["a"] != 5 || seen["b"] != 5 {
		t.Fatalf("expected even split, got %v", seen)
	}
}

func TestLeastConn(t *testing.T) {
	b1 := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckNone})
	b2 := NewBackend(models.Backend{Host: "b", Enabled: true, HealthCheckType: models.HealthCheckNone})
	pool := NewPool("least_conn", []*Backend{b1, b2})

	b1.Acquire()
	b1.Acquire()
	picked, err := pool.Pick()
	if err != nil {
		t.Fatal(err)
	}
	if picked.Model.Host != "b" {
		t.Fatalf("expected b (fewer active conns), got %s", picked.Model.Host)
	}
}

func TestWeighted(t *testing.T) {
	b1 := NewBackend(models.Backend{Host: "a", Weight: 3, Enabled: true, HealthCheckType: models.HealthCheckNone})
	b2 := NewBackend(models.Backend{Host: "b", Weight: 1, Enabled: true, HealthCheckType: models.HealthCheckNone})
	pool := NewPool("weighted", []*Backend{b1, b2})

	counts := map[string]int{}
	for i := 0; i < 8; i++ {
		b, err := pool.Pick()
		if err != nil {
			t.Fatal(err)
		}
		counts[b.Model.Host]++
	}
	if counts["a"] != 6 || counts["b"] != 2 {
		t.Fatalf("expected 3:1 weighted split (6:2 of 8), got %v", counts)
	}
}

func TestNoHealthyBackend(t *testing.T) {
	b1 := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckTCP, FailureThreshold: 1})
	pool := NewPool("round_robin", []*Backend{b1})
	if _, err := pool.Pick(); err != ErrNoHealthyBackend {
		t.Fatalf("expected ErrNoHealthyBackend, got %v", err)
	}
}

func TestRecordCheckThresholds(t *testing.T) {
	b := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckTCP, FailureThreshold: 2, RecoveryThreshold: 2})
	if b.Healthy() {
		t.Fatal("expected backend to start unhealthy pending checks")
	}
	b.RecordCheck(true)
	if b.Healthy() {
		t.Fatal("should still be unhealthy after 1 of 2 required successes")
	}
	changed, healthy := b.RecordCheck(true)
	if !changed || !healthy {
		t.Fatal("expected backend to become healthy after 2 successes")
	}

	b.RecordCheck(false)
	if !b.Healthy() {
		t.Fatal("should still be healthy after 1 of 2 required failures")
	}
	changed, healthy = b.RecordCheck(false)
	if !changed || healthy {
		t.Fatal("expected backend to become unhealthy after 2 failures")
	}
}