package retry

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

var (
	ErrInvalidAttempts = errors.New("retry: max attempts must be positive")
	ErrInvalidBackoff  = errors.New("retry: initial backoff cannot be negative")
)

// Retry calls operation until it succeeds, the context is canceled, or all
// attempts are exhausted. The delay doubles after each failed attempt.
func Retry(
	ctx context.Context,
	maxAttempts int,
	initialBackoff time.Duration,
	operation func(context.Context) error,
) error {
	if maxAttempts <= 0 {
		return ErrInvalidAttempts
	}
	if initialBackoff < 0 {
		return ErrInvalidBackoff
	}

	delay := initialBackoff
	var lastErr error

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

		lastErr = operation(ctx)
		if lastErr == nil {
			return nil
		}
		if attempt == maxAttempts-1 {
			break
		}

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

	return lastErr
}

func wait(ctx context.Context, delay time.Duration) error {
	if delay == 0 {
		return ctx.Err()
	}

	timer := time.NewTimer(delay)
	defer timer.Stop()

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

// doubled doubles a duration while avoiding overflow.
func doubled(delay time.Duration) time.Duration {
	const maxDuration = time.Duration(1<<63 - 1)
	if delay > maxDuration/2 {
		return maxDuration
	}
	return delay * 2
}

func TestRetryEventuallySucceeds(t *testing.T) {
	attempts := 0

	err := Retry(context.Background(), 3, 0, func(context.Context) error {
		attempts++
		if attempts < 3 {
			return errors.New("temporary failure")
		}
		return nil
	})

	if err != nil {
		t.Fatalf("Retry returned an unexpected error: %v", err)
	}
	if attempts != 3 {
		t.Fatalf("got %d attempts, want 3", attempts)
	}
}

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

	err := Retry(ctx, 5, time.Hour, func(context.Context) error {
		attempts++
		cancel()
		return errors.New("temporary failure")
	})

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