package retry

import (
	"context"
	"time"
)

// Config controls retry attempts and exponential backoff.
type Config struct {
	MaxAttempts int
	InitialDelay time.Duration
	MaxDelay     time.Duration
}

// Retry calls fn until it succeeds, the context is canceled, or attempts run out.
// The attempt number passed to fn starts at zero.
func Retry(ctx context.Context, cfg Config, fn func(context.Context, int) error) error {
	if cfg.MaxAttempts <= 0 {
		cfg.MaxAttempts = 3
	}
	if cfg.InitialDelay <= 0 {
		cfg.InitialDelay = 100 * time.Millisecond
	}
	if cfg.MaxDelay <= 0 {
		cfg.MaxDelay = 5 * time.Second
	}
	if cfg.InitialDelay > cfg.MaxDelay {
		cfg.InitialDelay = cfg.MaxDelay
	}

	delay := cfg.InitialDelay
	var lastErr error

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

		if lastErr = fn(ctx, attempt); lastErr == nil {
			return nil
		}
		if attempt == cfg.MaxAttempts-1 {
			break
		}

		timer := time.NewTimer(delay)
		select {
		case <-timer.C:
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return ctx.Err()
		}

		// Double the delay without overflowing or exceeding the cap.
		if delay >= cfg.MaxDelay-delay {
			delay = cfg.MaxDelay
		} else {
			delay *= 2
		}
	}

	return lastErr
}