package main

import (
	"sync/atomic"
	"testing"
	"time"

	"go.uber.org/zap"
)

func TestDrainControllerMarksNotReadyBeforeBeginDrain(t *testing.T) {
	rdy := &readiness{}
	rdy.MarkWarm()

	var calls atomic.Int32
	beginCh := make(chan struct{})
	delay := 40 * time.Millisecond

	controller := newDrainController(rdy, delay, func() int {
		calls.Add(1)
		close(beginCh)
		return 0
	}, zap.NewNop())

	start := time.Now()
	controller.Enter("test")

	if rdy.Ready() {
		t.Fatalf("expected readiness to flip to not-ready immediately")
	}

	select {
	case <-beginCh:
		t.Fatalf("expected beginDrain to wait for readiness delay")
	case <-time.After(15 * time.Millisecond):
	}

	controller.WaitForBeginDrain()

	if calls.Load() != 1 {
		t.Fatalf("expected beginDrain to run once, got %d", calls.Load())
	}
	if elapsed := time.Since(start); elapsed < delay {
		t.Fatalf("expected beginDrain after at least %s, got %s", delay, elapsed)
	}
}

func TestDrainControllerEnterOnlyBeginsDrainOnce(t *testing.T) {
	rdy := &readiness{}
	rdy.MarkWarm()

	var calls atomic.Int32
	controller := newDrainController(rdy, 0, func() int {
		return int(calls.Add(1))
	}, zap.NewNop())

	controller.Enter("first")
	controller.Enter("second")
	controller.WaitForBeginDrain()

	if calls.Load() != 1 {
		t.Fatalf("expected beginDrain to run once, got %d", calls.Load())
	}
}