package retry

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

// Retry calls fn up to attempts times, doubling the delay after each failure.
// It stops early when fn succeeds or the context is canceled.
func Retry(ctx context.Context, attempts int, initialDelay time.Duration, fn func() error) error {
	return retry(ctx, attempts, initialDelay, fn, wait)
}

func retry(
	ctx context.Context,
	attempts int,
	initialDelay time.Duration,
	fn func() error,
	sleep func(context.Context, time.Duration) error,
) error {
	if attempts < 1 {
		return errors.New("retry: attempts must be positive")
	}
	if initialDelay < 0 {
		return errors.New("retry: initial delay must not be negative")
	}
	if fn == nil {
		return errors.New("retry: function must not be nil")
	}

	delay := initialDelay
	var lastErr error

	for attempt := 0; attempt < attempts; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}

		if lastErr = fn(); lastErr == nil {
			return nil
		}
		if attempt == attempts-1 {
			break
		}

		if err := sleep(ctx, delay); err != nil {
			return err
		}

		// Saturate instead of overflowing time.Duration.
		if delay > time.Duration(1<<63-1)/2 {
			delay = time.Duration(1<<63 - 1)
		} else {
			delay *= 2
		}
	}

	return lastErr
}

func wait(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()

	select {
	case <-timer.C:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func TestRetryEventuallySucceeds(t *testing.T) {
	expected := errors.New("temporary failure")
	calls := 0
	var delays []time.Duration

	err := retry(
		context.Background(),
		4,
		10*time.Millisecond,
		func() error {
			calls++
			if calls < 3 {
				return expected
			}
			return nil
		},
		func(_ context.Context, delay time.Duration) error {
			delays = append(delays, delay)
			return nil
		},
	)

	if err != nil {
		t.Fatalf("Retry returned %v", err)
	}
	if calls != 3 {
		t.Fatalf("got %d calls, want 3", calls)
	}
	if len(delays) != 2 ||
		delays[0] != 10*time.Millisecond ||
		delays[1] != 20*time.Millisecond {
		t.Fatalf("got delays %v, want [10ms 20ms]", delays)
	}
}

func TestRetryStopsWhenContextIsCanceled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	calls := 0

	err := retry(
		ctx,
		5,
		time.Second,
		func() error {
			calls++
			return errors.New("temporary failure")
		},
		func(ctx context.Context, _ time.Duration) error {
			cancel()
			return ctx.Err()
		},
	)

	if !errors.Is(err, context.Canceled) {
		t.Fatalf("got %v, want context.Canceled", err)
	}
	if calls != 1 {
		t.Fatalf("got %d calls, want 1", calls)
	}
}