package retry

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

// Backoff controls retry timing. Zero values are replaced with sensible defaults.
type Backoff struct {
	InitialDelay time.Duration
	MaxDelay     time.Duration
	Multiplier   float64
	MaxAttempts  int
	Jitter       float64
	Sleep        func(context.Context, time.Duration) error
}

// Retry calls fn until it succeeds, the context is canceled, or MaxAttempts is reached.
func Retry(ctx context.Context, cfg Backoff, fn func(context.Context) error) error {
	if fn == nil {
		return errors.New("retry: nil function")
	}

	cfg = cfg.withDefaults()
	delay := cfg.InitialDelay

	var lastErr error
	for attempt := 1; ; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}

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

		if cfg.MaxAttempts > 0 && attempt >= cfg.MaxAttempts {
			return lastErr
		}

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

		delay = time.Duration(float64(delay) * cfg.Multiplier)
		if delay > cfg.MaxDelay {
			delay = cfg.MaxDelay
		}
	}
}

func (b Backoff) withDefaults() Backoff {
	if b.InitialDelay <= 0 {
		b.InitialDelay = 100 * time.Millisecond
	}
	if b.MaxDelay <= 0 {
		b.MaxDelay = 5 * time.Second
	}
	if b.Multiplier < 1 {
		b.Multiplier = 2
	}
	if b.Jitter < 0 {
		b.Jitter = 0
	}
	if b.Jitter > 1 {
		b.Jitter = 1
	}
	if b.Sleep == nil {
		b.Sleep = sleepContext
	}
	return b
}

func sleepContext(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 jitter(d time.Duration, pct float64) time.Duration {
	if pct <= 0 || d <= 0 {
		return d
	}

	spread := int64(float64(d) * pct)
	if spread <= 0 {
		return d
	}

	return time.Duration(int64(d) - spread + rand.Int63n(spread*2+1))
}

func TestRetrySucceedsAfterFailures(t *testing.T) {
	attempts := 0
	err := Retry(context.Background(), Backoff{
		InitialDelay: time.Millisecond,
		MaxAttempts:  3,
		Sleep: func(context.Context, time.Duration) error {
			return nil
		},
	}, func(context.Context) error {
		attempts++
		if attempts < 3 {
			return errors.New("temporary")
		}
		return nil
	})

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

func TestRetryStopsAtMaxAttempts(t *testing.T) {
	want := errors.New("still failing")
	attempts := 0

	err := Retry(context.Background(), Backoff{
		InitialDelay: time.Millisecond,
		MaxAttempts:  2,
		Sleep: func(context.Context, time.Duration) error {
			return nil
		},
	}, func(context.Context) error {
		attempts++
		return want
	})

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