package retry

import (
	"context"
	"errors"
	"testing"
	"time"
)

func TestRunSucceedsOnFirstTry(t *testing.T) {
	calls := 0
	res := Run(context.Background(), Default(), func(ctx context.Context, attempt int) error {
		calls++
		return nil
	})
	if calls != 1 {
		t.Fatalf("want 1 call, got %d", calls)
	}
	if res.Attempts != 1 {
		t.Fatalf("want Attempts=1, got %d", res.Attempts)
	}
	if res.LastError != nil {
		t.Fatalf("want nil err, got %v", res.LastError)
	}
}

func TestRunRetriesThenSucceeds(t *testing.T) {
	calls := 0
	cfg := Config{
		MaxAttempts: 5,
		BaseDelay:   1 * time.Millisecond,
		MaxDelay:    5 * time.Millisecond,
		Budget:      5 * time.Second,
	}
	res := Run(context.Background(), cfg, func(ctx context.Context, attempt int) error {
		calls++
		if attempt < 3 {
			return errors.New("transient")
		}
		return nil
	})
	if calls != 3 {
		t.Fatalf("want 3 calls, got %d", calls)
	}
	if res.Attempts != 3 {
		t.Fatalf("want Attempts=3, got %d", res.Attempts)
	}
	if res.LastError != nil {
		t.Fatalf("want nil err, got %v", res.LastError)
	}
}

func TestRunExhaustsAttempts(t *testing.T) {
	calls := 0
	cfg := Config{
		MaxAttempts: 3,
		BaseDelay:   1 * time.Millisecond,
		MaxDelay:    5 * time.Millisecond,
		Budget:      5 * time.Second,
	}
	res := Run(context.Background(), cfg, func(ctx context.Context, attempt int) error {
		calls++
		return errors.New("nope")
	})
	if calls != 3 {
		t.Fatalf("want 3 calls, got %d", calls)
	}
	if res.Attempts != 3 {
		t.Fatalf("want Attempts=3, got %d", res.Attempts)
	}
	if res.LastError == nil {
		t.Fatalf("want non-nil err")
	}
}

func TestRunShortCircuitsOnPermanent(t *testing.T) {
	calls := 0
	cfg := Config{
		MaxAttempts: 5,
		BaseDelay:   1 * time.Millisecond,
		MaxDelay:    5 * time.Millisecond,
		Budget:      5 * time.Second,
	}
	res := Run(context.Background(), cfg, func(ctx context.Context, attempt int) error {
		calls++
		return &PermanentError{Err: errors.New("bad token")}
	})
	if calls != 1 {
		t.Fatalf("want 1 call, got %d", calls)
	}
	if !IsPermanent(res.LastError) {
		t.Fatalf("want PermanentError, got %T: %v", res.LastError, res.LastError)
	}
}

func TestRunRespectsBudget(t *testing.T) {
	calls := 0
	cfg := Config{
		MaxAttempts: 100,
		BaseDelay:   50 * time.Millisecond,
		MaxDelay:    200 * time.Millisecond,
		Budget:      100 * time.Millisecond, // tight budget
	}
	start := time.Now()
	res := Run(context.Background(), cfg, func(ctx context.Context, attempt int) error {
		calls++
		return errors.New("nope")
	})
	elapsed := time.Since(start)
	if elapsed > 500*time.Millisecond {
		t.Fatalf("budget exceeded: took %s", elapsed)
	}
	// With Budget=100ms and base 50ms, we expect at most
	// 3-4 calls (50 + 100 = 150ms would already bust the
	// budget for the next wait).
	if calls > 5 {
		t.Fatalf("too many calls under tight budget: %d", calls)
	}
	if res.LastError == nil {
		t.Fatalf("want non-nil err")
	}
}

func TestRunRespectsContextCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	calls := 0
	cfg := Config{
		MaxAttempts: 10,
		BaseDelay:   100 * time.Millisecond,
		MaxDelay:    500 * time.Millisecond,
		Budget:      5 * time.Second,
	}
	// Cancel the ctx after the first call.
	res := Run(ctx, cfg, func(ctx context.Context, attempt int) error {
		calls++
		if calls == 1 {
			// Cancel while we're sleeping.
			go func() {
				time.Sleep(20 * time.Millisecond)
				cancel()
			}()
		}
		return errors.New("nope")
	})
	if calls >= 10 {
		t.Fatalf("ctx cancel did not stop the loop; calls=%d", calls)
	}
	if res.LastError == nil {
		t.Fatalf("want non-nil err")
	}
}

func TestBackoffMonotonic(t *testing.T) {
	// Per-attempt wait is the time we sleep BEFORE that
	// attempt. So attempt 1 = no wait, attempt 2 = base,
	// attempt 3 = base*2, etc., capped at max.
	base := 100 * time.Millisecond
	max := 2 * time.Second
	prev := time.Duration(0)
	for n := 2; n <= 10; n++ {
		got := backoff(base, max, n)
		if n == 2 && got != base {
			t.Fatalf("attempt 2: want %s, got %s", base, got)
		}
		if got > max {
			t.Fatalf("attempt %d: exceeded max: %s", n, got)
		}
		// Wait must be >= previous (or equal under cap).
		if got < prev {
			t.Fatalf("attempt %d: wait decreased: %s < %s", n, got, prev)
		}
		prev = got
	}
}
