package retry

import (
	"math/rand/v2"
	"time"
)

// Backoff computes the delay before the next retry attempt.
type Backoff interface {
	// Delay returns the wait duration for the given attempt number (zero-based).
	Delay(attempt int) time.Duration
}

// ExponentialBackoff returns a Backoff that doubles the delay on each attempt.
// The delay is calculated as base * 2^attempt, capped at max. When withJitter
// is true, a random duration in [0, delay*0.5) is added.
func ExponentialBackoff(base, max time.Duration, withJitter bool) Backoff {
	return &exponentialBackoff{
		base:       base,
		max:        max,
		withJitter: withJitter,
	}
}

// ConstantBackoff returns a Backoff that always returns the same delay.
func ConstantBackoff(d time.Duration) Backoff {
	return constantBackoff{delay: d}
}

type exponentialBackoff struct {
	base       time.Duration
	max        time.Duration
	withJitter bool
}

func (b *exponentialBackoff) Delay(attempt int) time.Duration {
	delay := b.base
	for range attempt {
		delay *= 2
		if delay >= b.max {
			delay = b.max
			break
		}
	}

	if b.withJitter {
		jitter := time.Duration(rand.Int64N(int64(delay / 2)))
		delay += jitter
	}

	if delay > b.max {
		delay = b.max
	}

	return delay
}

type constantBackoff struct {
	delay time.Duration
}

func (b constantBackoff) Delay(_ int) time.Duration {
	return b.delay
}