package main

import (
	"sync"
	"testing"
	"time"
)

func newTestMgr(n int) *proxyManager {
	// Zero settle/warmup so tests don't sleep; tiny maxHold disabled by being long.
	return newProxyManager(n, 0, 0, time.Hour)
}

func TestAcquireAnySkipsDraining(t *testing.T) {
	m := newTestMgr(3)
	m.drain(0, time.Second) // slot 0 out
	m.drain(1, time.Second) // slot 1 out

	// Only slot 2 should ever come back from acquireAny.
	for i := 0; i < 50; i++ {
		id, err := m.acquireAny()
		if err != nil {
			t.Fatalf("acquireAny: %v", err)
		}
		if id != 2 {
			t.Fatalf("expected only slot 2 selectable, got %d", id)
		}
		m.release(id)
	}
}

func TestAcquireAnyNoneAvailable(t *testing.T) {
	m := newTestMgr(2)
	m.drain(0, time.Second)
	m.drain(1, time.Second)
	if _, err := m.acquireAny(); err != errNoProxies {
		t.Fatalf("expected errNoProxies, got %v", err)
	}
}

func TestAcquirePinnedDraining(t *testing.T) {
	m := newTestMgr(2)
	if err := m.acquirePinned(1); err != nil {
		t.Fatalf("healthy slot should acquire: %v", err)
	}
	m.release(1)

	m.drain(1, time.Second)
	if err := m.acquirePinned(1); err != errDraining {
		t.Fatalf("expected errDraining for pinned drained slot, got %v", err)
	}
}

func TestDrainWaitsForInflight(t *testing.T) {
	m := newTestMgr(1)
	if _, err := m.acquireAny(); err != nil { // hold slot 0 in flight
		t.Fatal(err)
	}

	done := make(chan int, 1)
	go func() { done <- m.drain(0, 2*time.Second) }()

	// drain must not return while the request is still in flight.
	select {
	case <-done:
		t.Fatal("drain returned before in-flight request released")
	case <-time.After(150 * time.Millisecond):
	}

	m.release(0)
	select {
	case residual := <-done:
		if residual != 0 {
			t.Fatalf("expected clean drain (0 residual), got %d", residual)
		}
	case <-time.After(time.Second):
		t.Fatal("drain did not return after release")
	}
}

func TestDrainTimeoutReportsResidual(t *testing.T) {
	m := newTestMgr(1)
	if _, err := m.acquireAny(); err != nil { // never released
		t.Fatal(err)
	}
	if residual := m.drain(0, 100*time.Millisecond); residual != 1 {
		t.Fatalf("expected residual 1 on timeout, got %d", residual)
	}
}

func TestUndrainRestoresSelection(t *testing.T) {
	m := newTestMgr(1)
	m.drain(0, time.Second)
	if _, err := m.acquireAny(); err != errNoProxies {
		t.Fatalf("expected slot unavailable while draining, got %v", err)
	}
	m.undrain(0) // warmup is zero in test mgr → immediately selectable
	if _, err := m.acquireAny(); err != nil {
		t.Fatalf("expected slot available after undrain, got %v", err)
	}
}

func TestWarmupHoldsSlotOut(t *testing.T) {
	m := newProxyManager(1, 0, 200*time.Millisecond, time.Hour)
	m.undrain(0) // sets warmupAt = now + 200ms
	if _, err := m.acquireAny(); err != errNoProxies {
		t.Fatalf("expected slot held out during warmup, got %v", err)
	}
	time.Sleep(250 * time.Millisecond)
	if _, err := m.acquireAny(); err != nil {
		t.Fatalf("expected slot available after warmup, got %v", err)
	}
}

func TestSelfHealOnMissedUndrain(t *testing.T) {
	m := newProxyManager(1, 0, 0, 100*time.Millisecond) // maxHold 100ms
	m.drain(0, time.Second)
	if _, err := m.acquireAny(); err != errNoProxies {
		t.Fatalf("expected unavailable right after drain, got %v", err)
	}
	time.Sleep(200 * time.Millisecond) // self-heal timer should fire
	if _, err := m.acquireAny(); err != nil {
		t.Fatalf("expected self-heal to restore slot, got %v", err)
	}
}

func TestSnapshotEgressAndRotation(t *testing.T) {
	m := newProxyManager(3, 0, 0, time.Hour)
	m.setEgress(2, "9.9.9.9", 12*time.Millisecond, true)

	m.drain(0, time.Second) // first rotation: establishes lastRotated only
	m.undrain(0)
	time.Sleep(5 * time.Millisecond)
	m.drain(1, time.Second) // second rotation: gap becomes the interval; slot 1 left draining

	snap := m.snapshot()

	if got := snap.Proxies[2]; got.EgressIP != "9.9.9.9" || !got.Reachable || got.CheckedAgeSec == nil {
		t.Fatalf("egress not reflected on slot 2: %+v", got)
	}
	if snap.Proxies[1].State != "draining" {
		t.Fatalf("slot 1 should be draining, got %q", snap.Proxies[1].State)
	}
	if snap.Rotation.LastSlot == nil || *snap.Rotation.LastSlot != 1 {
		t.Fatalf("last_slot want 1, got %v", snap.Rotation.LastSlot)
	}
	if snap.Rotation.NextSlot == nil || *snap.Rotation.NextSlot != 2 {
		t.Fatalf("next_slot want 2 (round-robin), got %v", snap.Rotation.NextSlot)
	}
	if snap.Rotation.NextEgressIP != "9.9.9.9" {
		t.Fatalf("next_egress_ip want 9.9.9.9 (slot 2's IP), got %q", snap.Rotation.NextEgressIP)
	}
	if snap.Rotation.ETASec == nil {
		t.Fatal("expected an ETA once cadence is known")
	}
}

func TestConcurrentAcquireRelease(t *testing.T) {
	m := newTestMgr(4)
	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if id, err := m.acquireAny(); err == nil {
				m.release(id)
			}
		}()
	}
	wg.Wait()
	for i := 0; i < 4; i++ {
		if m.inflight[i] != 0 {
			t.Fatalf("slot %d leaked in-flight count: %d", i, m.inflight[i])
		}
	}
}