package main

import (
	"context"
	"errors"
	"fmt"
	"math/rand/v2"
	"sync/atomic"
	"time"
)

var ErrBusy = errors.New("bulkhead full")
var ErrBudget = errors.New("retry budget exhausted")
var ErrTransient = errors.New("transient")

type Budget struct{ remaining atomic.Int64 }

func NewBudget(n int64) *Budget { b := &Budget{}; b.remaining.Store(n); return b }
func (b *Budget) Take() bool {
	for {
		n := b.remaining.Load()
		if n <= 0 {
			return false
		}
		if b.remaining.CompareAndSwap(n, n-1) {
			return true
		}
	}
}

// Backoff uses saturating multiplication: the duration never wraps around.
func Backoff(base, cap time.Duration, retry int) time.Duration {
	if base <= 0 || cap <= 0 {
		return 0
	}
	d := min(base, cap)
	for i := 0; i < retry && d < cap; i++ {
		if d > cap/2 {
			return cap
		}
		d *= 2
	}
	return d
}
func FullJitter(d time.Duration) time.Duration {
	if d <= 0 {
		return 0
	}
	return time.Duration(rand.Int64N(int64(d)))
}
func Wait(ctx context.Context, d time.Duration) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-t.C:
		return ctx.Err()
	}
}

// attempts includes the initial attempt. The caller supplies a safe-to-retry operation.
func Retry(ctx context.Context, attempts int, b *Budget, delay func(int) time.Duration, fn func(context.Context) error) error {
	if attempts < 1 {
		return errors.New("attempts must be positive")
	}
	var err error
	for i := 0; i < attempts; i++ {
		if e := ctx.Err(); e != nil {
			return e
		}
		if i > 0 {
			if !b.Take() {
				return errors.Join(ErrBudget, err)
			}
			if e := Wait(ctx, delay(i-1)); e != nil {
				return e
			}
		}
		if e := ctx.Err(); e != nil {
			return e
		}
		err = fn(ctx)
		if err == nil || !errors.Is(err, ErrTransient) {
			return err
		}
	}
	return err
}

type Bulkhead chan struct{}

func NewBulkhead(n int) Bulkhead {
	if n < 1 {
		panic("positive capacity required")
	}
	return make(Bulkhead, n)
}
func (b Bulkhead) Try(ctx context.Context, fn func(context.Context) error) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	select {
	case b <- struct{}{}:
		defer func() { <-b }()
	default:
		return ErrBusy
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	return fn(ctx)
}
func main() {
	p, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Hour))
	defer cancel()
	c, done := context.WithTimeout(p, 2*time.Hour)
	defer done()
	pd, _ := p.Deadline()
	cd, _ := c.Deadline()
	fmt.Println("child inherits earlier deadline:", pd.Equal(cd))
	calls := 0
	err := Retry(context.Background(), 5, NewBudget(2), func(int) time.Duration { return 0 }, func(context.Context) error { calls++; return ErrTransient })
	fmt.Println("attempts:", calls, "budget exhausted:", errors.Is(err, ErrBudget))
	b := NewBulkhead(1)
	_ = b.Try(context.Background(), func(ctx context.Context) error {
		fmt.Println("second request rejected:", errors.Is(b.Try(ctx, func(context.Context) error { return nil }), ErrBusy))
		return nil
	})
	fmt.Println("slot released:", b.Try(context.Background(), func(context.Context) error { return nil }) == nil)
	fmt.Println("capped backoff:", Backoff(10*time.Millisecond, 80*time.Millisecond, 0), Backoff(10*time.Millisecond, 80*time.Millisecond, 1), Backoff(10*time.Millisecond, 80*time.Millisecond, 20))
}