package retry

import (
	"context"
	"errors"
	"math"
	"math/rand"
	"time"
)

// Config controls retry timing.
type Config struct {
	// Attempts is the total number of tries, including the first call.
	Attempts int

	// InitialDelay is the delay before the first retry.
	InitialDelay time.Duration

	// MaxDelay caps the exponential delay.
	MaxDelay time.Duration

	// Multiplier grows the delay after each failed attempt.
	Multiplier float64

	// Jitter randomly varies each delay by up to this fraction.
	// For example, 0.2 means +/-20%.
	Jitter float64
}

// DefaultConfig is a conservative retry policy for short-lived operations.
var DefaultConfig = Config{
	Attempts:     3,
	InitialDelay: 100 * time.Millisecond,
	MaxDelay:     2 * time.Second,
	Multiplier:   2,
	Jitter:       0.2,
}

// Retry calls fn until it succeeds, the attempts are exhausted, or ctx is done.
func Retry(ctx context.Context, cfg Config, fn func(context.Context) error) error {
	cfg = normalize(cfg)

	var lastErr error
	delay := cfg.InitialDelay

	for attempt := 1; attempt <= cfg.Attempts; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}

		if err := fn(ctx); err != nil {
			lastErr = err
		} else {
			return nil
		}

		if attempt == cfg.Attempts {
			break
		}

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

		delay = nextDelay(delay, cfg)
	}

	return lastErr
}

func normalize(cfg Config) Config {
	if cfg.Attempts <= 0 {
		cfg.Attempts = DefaultConfig.Attempts
	}
	if cfg.InitialDelay <= 0 {
		cfg.InitialDelay = DefaultConfig.InitialDelay
	}
	if cfg.MaxDelay <= 0 {
		cfg.MaxDelay = DefaultConfig.MaxDelay
	}
	if cfg.Multiplier < 1 {
		cfg.Multiplier = DefaultConfig.Multiplier
	}
	if cfg.Jitter < 0 {
		cfg.Jitter = 0
	}
	if cfg.Jitter > 1 {
		cfg.Jitter = 1
	}
	return cfg
}

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

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

func nextDelay(current time.Duration, cfg Config) time.Duration {
	next := float64(current) * cfg.Multiplier
	if next > float64(cfg.MaxDelay) || math.IsInf(next, 0) {
		return cfg.MaxDelay
	}
	return time.Duration(next)
}

func jitter(d time.Duration, fraction float64) time.Duration {
	if fraction <= 0 || d <= 0 {
		return d
	}

	delta := fraction * (rand.Float64()*2 - 1)
	return time.Duration(float64(d) * (1 + delta))
}

// TestRetrySucceedsAfterFailures verifies retrying a transient failure.
func TestRetrySucceedsAfterFailures(t interface {
	Fatalf(string, ...any)
}) {
	ctx := context.Background()
	attempts := 0

	err := Retry(ctx, Config{
		Attempts:     3,
		InitialDelay: time.Nanosecond,
		MaxDelay:     time.Nanosecond,
		Multiplier:   2,
	}, func(context.Context) error {
		attempts++
		if attempts < 3 {
			return errors.New("temporary failure")
		}
		return nil
	})

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

// TestRetryStopsOnContextCancel verifies cancellation wins over more retries.
func TestRetryStopsOnContextCancel(t interface {
	Fatalf(string, ...any)
}) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	calls := 0
	err := Retry(ctx, Config{Attempts: 3}, func(context.Context) error {
		calls++
		return errors.New("temporary failure")
	})

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