package main

import (
	"context"
	"errors"
	"net"
	"sync/atomic"
	"testing"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
)

func TestReadinessNotReadyUntilWarmed(t *testing.T) {
	r := &readiness{}
	if r.Ready() {
		t.Fatal("expected a fresh readiness to be not ready until warmed")
	}
}

func TestReadinessReadyAfterWarm(t *testing.T) {
	r := &readiness{}
	r.MarkWarm()
	if !r.Ready() {
		t.Fatal("expected readiness to be ready after MarkWarm")
	}
}

func TestReadinessNotReadyWhileDraining(t *testing.T) {
	r := &readiness{}
	r.MarkWarm()
	r.MarkDraining()
	if r.Ready() {
		t.Fatal("expected readiness to be not ready while draining, even after warm")
	}
}

func TestReadinessDrainIsTerminal(t *testing.T) {
	r := &readiness{}
	r.MarkDraining()
	r.MarkWarm() // a late warmup completion must not re-open readiness after drain
	if r.Ready() {
		t.Fatal("expected drain to be terminal: warming after draining must stay not ready")
	}
}

func TestReadinessMarkWarmReportsReady(t *testing.T) {
	r := &readiness{}
	if !r.MarkWarm() {
		t.Fatal("expected MarkWarm to report ready when not draining")
	}
	if !r.Ready() {
		t.Fatal("expected ready after MarkWarm")
	}
}

func TestReadinessMarkWarmIsNoOpAfterDraining(t *testing.T) {
	r := &readiness{}
	r.MarkDraining()
	if r.MarkWarm() {
		t.Fatal("expected MarkWarm to report not-ready (no-op) once draining")
	}
	if r.Ready() {
		t.Fatal("expected readiness to stay not ready after draining")
	}
}

func TestReadinessDrainingReflectsDrainState(t *testing.T) {
	r := &readiness{}
	if r.Draining() {
		t.Fatal("expected not draining before MarkDraining")
	}
	r.MarkDraining()
	if !r.Draining() {
		t.Fatal("expected draining after MarkDraining")
	}
}

func TestWarmUpDoesNotLogRetryWhenCanceledDuringProbe(t *testing.T) {
	core, logs := observer.New(zapcore.WarnLevel)

	ctx, cancel := context.WithCancel(context.Background())
	probe := func(context.Context) error {
		cancel() // simulate the context being cancelled while the probe runs
		return errors.New("boom")
	}

	err := warmUp(ctx, probe, time.Millisecond, zap.New(core))
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("expected context.Canceled, got %v", err)
	}
	if n := logs.FilterMessage("warmup probe failed, retrying").Len(); n != 0 {
		t.Fatalf("expected no retry log when cancelled during a probe, got %d", n)
	}
}

func TestWarmUpRetriesUntilProbeSucceeds(t *testing.T) {
	var calls int32
	probe := func(context.Context) error {
		if atomic.AddInt32(&calls, 1) < 3 {
			return errors.New("backend not ready")
		}
		return nil
	}

	if err := warmUp(context.Background(), probe, time.Millisecond, zap.NewNop()); err != nil {
		t.Fatalf("expected warmUp to succeed once the probe passes, got %v", err)
	}
	if got := atomic.LoadInt32(&calls); got != 3 {
		t.Fatalf("expected probe to be retried until success (3 calls), got %d", got)
	}
}

func TestWarmUpAbortsOnContextCancel(t *testing.T) {
	probe := func(context.Context) error { return errors.New("never ready") }

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	err := warmUp(ctx, probe, time.Millisecond, zap.NewNop())
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("expected warmUp to return context.Canceled when aborted, got %v", err)
	}
}

func TestTCPReachableProbeSucceedsWhenListenerUp(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	defer l.Close()

	probe := tcpReachableProbe(l.Addr().String(), time.Second)
	if err := probe(context.Background()); err != nil {
		t.Fatalf("expected probe to succeed against a live listener, got %v", err)
	}
}

func TestTCPReachableProbeFailsWhenNothingListening(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr := l.Addr().String()
	if err := l.Close(); err != nil { // free the port so nothing is listening
		t.Fatalf("close: %v", err)
	}

	probe := tcpReachableProbe(addr, time.Second)
	if err := probe(context.Background()); err == nil {
		t.Fatal("expected probe to fail when nothing is listening on the address")
	}
}