package retry

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

type Config struct {
	Attempts int
	Initial  time.Duration
	MaxDelay time.Duration
	Factor   float64
	Jitter   float64
}

func (c Config) withDefaults() Config {
	if c.Attempts <= 0 {
		c.Attempts = 3
	}
	if c.Initial <= 0 {
		c.Initial = 100 * time.Millisecond
	}
	if c.MaxDelay <= 0 {
		c.MaxDelay = 5 * time.Second
	}
	if c.Factor <= 1 {
		c.Factor = 2
	}
	if c.Jitter < 0 {
		c.Jitter = 0
	}
	if c.Jitter > 1 {
		c.Jitter = 1
	}
	return c
}

// Do runs fn until it succeeds, the context is canceled, or attempts are exhausted.
func Do(ctx context.Context, cfg Config, fn func(context.Context) error) error {
	cfg = cfg.withDefaults()
	delay := cfg.Initial

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

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

		if attempt == cfg.Attempts {
			break
		}

		wait := withJitter(delay, cfg.Jitter)
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}

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

	return lastErr
}

// withJitter adjusts d by up to pct in either direction.
func withJitter(d time.Duration, pct float64) time.Duration {
	if pct <= 0 {
		return d
	}
	spread := float64(d) * pct
	offset := (rand.Float64()*2 - 1) * spread
	if adjusted := time.Duration(float64(d) + offset); adjusted > 0 {
		return adjusted
	}
	return d
}

func TestDoRetriesUntilSuccess(t *testing.T) {
	ctx := context.Background()
	attempts := 0

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

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

func TestDoReturnsLastError(t *testing.T) {
	want := errors.New("final failure")
	attempts := 0

	err := Do(context.Background(), Config{Attempts: 2, Initial: time.Nanosecond}, func(context.Context) error {
		attempts++
		return want
	})

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