package main

import (
	"context"
	"errors"
	"github.com/sony/gobreaker/v2"
	"golang.org/x/time/rate"
	"testing"
	"time"
)

func TestRateAndCanceledWait(t *testing.T) {
	a, b, c, d := RateDemo()
	if !a || !b || c || !d {
		t.Fatal(a, b, c, d)
	}
	l := rate.NewLimiter(1, 1)
	l.Allow()
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if err := l.Wait(ctx); !errors.Is(err, context.Canceled) {
		t.Fatal(err)
	}
	if err := l.WaitN(context.Background(), 2); err == nil {
		t.Fatal("request larger than burst accepted")
	}
}
func TestBreakerAndRecovery(t *testing.T) {
	calls, blocked, state := BreakerDemo()
	if calls != 2 || !blocked || state != "open" {
		t.Fatal(calls, blocked, state)
	}
	cb := NewBreaker(time.Millisecond)
	for range 2 {
		_, _ = cb.Execute(func() (int, error) { return 0, transient })
	}
	deadline := time.Now().Add(time.Second)
	for cb.State() != gobreaker.StateHalfOpen {
		if time.Now().After(deadline) {
			t.Fatal("no half-open")
		}
		time.Sleep(time.Millisecond)
	}
	if _, err := cb.Execute(func() (int, error) { return 1, nil }); err != nil || cb.State() != gobreaker.StateClosed {
		t.Fatal(err, cb.State())
	}
}
func TestRetryAndCancellation(t *testing.T) {
	calls, value, err := RetryDemo(context.Background())
	if calls != 3 || value != 42 || err != nil {
		t.Fatal(calls, value, err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	calls, _, err = RetryDemo(ctx)
	if calls != 0 || !errors.Is(err, context.Canceled) {
		t.Fatal(calls, err)
	}
}
func TestSentinelReleasesEntry(t *testing.T) {
	blocked, released, err := SentinelDemo()
	if !blocked || !released || err != nil {
		t.Fatal(blocked, released, err)
	}
}
func TestVegetaRecordsFailures(t *testing.T) {
	n, s, a, b := VegetaDemo()
	if n != 4 || s != .5 || a != 2 || b != 2 {
		t.Fatal(n, s, a, b)
	}
}

func TestReservationCancellation(t *testing.T) {
	l := rate.NewLimiter(1, 1)
	now := time.Now()
	if !l.AllowN(now, 1) {
		t.Fatal("initial token")
	}
	reserved := l.ReserveN(now, 1)
	if !reserved.OK() || reserved.DelayFrom(now) != time.Second {
		t.Fatal("unexpected reservation")
	}
	reserved.CancelAt(now)
	next := l.ReserveN(now, 1)
	defer next.CancelAt(now)
	if next.DelayFrom(now) != time.Second {
		t.Fatal("canceled reservation still consumed capacity", next.DelayFrom(now))
	}
}

func waitHalfOpen(t *testing.T, cb *gobreaker.CircuitBreaker[int]) {
	t.Helper()
	deadline := time.Now().Add(time.Second)
	for cb.State() != gobreaker.StateHalfOpen {
		if time.Now().After(deadline) {
			t.Fatal("no half-open")
		}
		time.Sleep(time.Millisecond)
	}
}
func TestHalfOpenFailureReopens(t *testing.T) {
	var transitions []gobreaker.State
	cb := gobreaker.NewCircuitBreaker[int](gobreaker.Settings{Timeout: time.Millisecond, MaxRequests: 1, ReadyToTrip: func(c gobreaker.Counts) bool { return c.ConsecutiveFailures >= 2 }, OnStateChange: func(_ string, _, to gobreaker.State) { transitions = append(transitions, to) }})
	for range 2 {
		_, _ = cb.Execute(func() (int, error) { return 0, transient })
	}
	waitHalfOpen(t, cb)
	_, err := cb.Execute(func() (int, error) { return 0, transient })
	if !errors.Is(err, transient) || transitions[len(transitions)-1] != gobreaker.StateOpen {
		t.Fatal(err, transitions)
	}
}
func TestHalfOpenAllowsOnlyOneConcurrentProbe(t *testing.T) {
	cb := NewBreaker(time.Millisecond)
	for range 2 {
		_, _ = cb.Execute(func() (int, error) { return 0, transient })
	}
	waitHalfOpen(t, cb)
	started := make(chan struct{})
	release := make(chan struct{})
	done := make(chan error, 1)
	go func() {
		_, err := cb.Execute(func() (int, error) { close(started); <-release; return 1, nil })
		done <- err
	}()
	defer func() {
		close(release)
		if err := <-done; err != nil {
			t.Error(err)
		}
	}()
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("probe did not start")
	}
	called := false
	_, err := cb.Execute(func() (int, error) { called = true; return 1, nil })
	if !errors.Is(err, gobreaker.ErrTooManyRequests) || called {
		t.Fatal(err, called)
	}
}