// Package backoff implements exponential backoff with jitter for retry logic.
package backoff

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

// Config holds the configuration for a Backoff instance.
type Config struct {
	// Initial is the starting upper bound for the backoff duration.
	// Must be > 0.
	Initial time.Duration `yaml:"initial"`
	// Max is the absolute maximum upper bound for the backoff duration.
	// Must be >= Initial.
	Max time.Duration `yaml:"max"`
	// RNG is an optional source of randomness.
	// If nil, the global math/rand/v2 source is used.
	RNG *rand.Rand `yaml:"-"`
}

// Validate checks that the configuration values are sensible.
func (cfg Config) Validate() error {
	if cfg.Initial <= 0 {
		return fmt.Errorf("invalid initial backoff: %v", cfg.Initial)
	}
	if cfg.Max < cfg.Initial {
		return fmt.Errorf("invalid max duration: %v", cfg.Max)
	}
	// MaxInt64 will overflow when calculating jitter.
	if cfg.Max == math.MaxInt64 {
		return fmt.Errorf("max duration cannot be MaxInt64")
	}
	return nil
}

// Backoff implements exponential backoff with full jitter.
// It is NOT safe for concurrent use.
type Backoff struct {
	initial time.Duration
	max     time.Duration
	current time.Duration
	rng     *rand.Rand
}

// New creates a new Backoff instance with the provided configuration.
func New(cfg Config) (*Backoff, error) {
	if err := cfg.Validate(); err != nil {
		return nil, fmt.Errorf("invalid config: %v", err)
	}

	return &Backoff{
		initial: cfg.Initial,
		max:     cfg.Max,
		current: cfg.Initial,
		rng:     cfg.RNG,
	}, nil
}

// Next returns the next backoff delay.
func (b *Backoff) Next() time.Duration {
	limit := b.current

	// Update state for the next call.
	if b.current >= b.max {
		b.current = b.max
	} else if b.current > b.max/2 {
		// If doubling would exceed max, just clamp to max.
		b.current = b.max
	} else {
		b.current *= 2
	}

	// Calculate jitter; return random in [0, limit].
	// Int64N(n) returns values in [0, n).
	// For [0, limit], we use limit + 1.
	var jitter int64
	if b.rng != nil {
		jitter = b.rng.Int64N(int64(limit) + 1)
	} else {
		jitter = rand.Int64N(int64(limit) + 1)
	}

	return time.Duration(jitter)
}

// Reset resets the current backoff cap to the initial value.
// This should be called after a successful operation to restart
// the backoff sequence for future retries.
func (b *Backoff) Reset() {
	b.current = b.initial
}