package retry

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

type Config struct {
	InitialDelay time.Duration
	MaxDelay     time.Duration
	Multiplier   float64
	Jitter        float64
	MaxAttempts  int
}

func DefaultConfig() Config {
	return Config{
		InitialDelay: 100 * time.Millisecond,
		MaxDelay:     5 * time.Second,
		Multiplier:   2,
		Jitter:        0.2,
		MaxAttempts:  5,
	}
}

// Do retries fn with exponential backoff until it succeeds, the context is done,
// or MaxAttempts is reached. MaxAttempts <= 0 means retry indefinitely.
func Do(ctx context.Context, cfg Config, fn func() error) error {
	if cfg.InitialDelay <= 0 {
		cfg.InitialDelay = 100 * time.Millisecond
	}
	if cfg.MaxDelay <= 0 {
		cfg.MaxDelay = 5 * time.Second
	}
	if cfg.Multiplier < 1 {
		cfg.Multiplier = 2
	}
	if cfg.Jitter < 0 {
		cfg.Jitter = 0
	}

	var lastErr error
	delay := cfg.InitialDelay

	for attempt := 1; cfg.MaxAttempts <= 0 || attempt <= cfg.MaxAttempts; attempt++ {
		if err := ctx.Err(); err != nil {
			return errors.Join(lastErr, err)
		}

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

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

		wait := withJitter(delay, cfg.Jitter)
		timer := time.NewTimer(wait)

		select {
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return errors.Join(lastErr, ctx.Err())
		case <-timer.C:
		}

		delay = nextDelay(delay, cfg.MaxDelay, cfg.Multiplier)
	}

	return lastErr
}

func nextDelay(current, max time.Duration, multiplier float64) time.Duration {
	next := time.Duration(float64(current) * multiplier)
	if next <= 0 || next > max {
		return max
	}
	return next
}

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

	spread := float64(delay) * jitter
	min := float64(delay) - spread
	max := float64(delay) + spread

	return time.Duration(min + rand.Float64()*(max-min))
}

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

	err := Do(context.Background(), Config{
		InitialDelay: time.Millisecond,
		MaxDelay:     time.Millisecond,
		Multiplier:   2,
		MaxAttempts:  3,
	}, func() error {
		attempts++
		if attempts < 2 {
			return errors.New("temporary failure")
		}
		return nil
	})

	if err != nil {
		t.Fatalf("expected success, got %v", err)
	}
	if attempts != 2 {
		t.Fatalf("expected 2 attempts, got %d", attempts)
	}
}

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

	err := Do(context.Background(), Config{
		InitialDelay: time.Millisecond,
		MaxDelay:     time.Millisecond,
		Multiplier:   2,
		MaxAttempts:  3,
	}, func() error {
		attempts++
		return want
	})

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