package agentrouter

import (
	"context"
	"math/rand"
	"time"

	"connectrpc.com/connect"
)

// RetryPolicy configures the SDK's retry interceptor.
type RetryPolicy struct {
	// MaxAttempts is the total attempts including the initial call. Zero
	// or negative disables retries.
	MaxAttempts int

	// InitialBackoff is the wait before the first retry.
	InitialBackoff time.Duration

	// MaxBackoff is the cap on individual sleep intervals.
	MaxBackoff time.Duration

	// TotalBudget is the absolute deadline applied across all attempts.
	// Zero means no budget.
	TotalBudget time.Duration

	// JitterFraction is the random fraction added to each backoff (0.2
	// means +-20%). Zero disables jitter.
	JitterFraction float64
}

// DefaultRetryPolicy is the SDK default: 3 attempts, 100ms initial backoff,
// 5s cap, 90s total budget, 20% jitter.
var DefaultRetryPolicy = RetryPolicy{
	MaxAttempts:    3,
	InitialBackoff: 100 * time.Millisecond,
	MaxBackoff:     5 * time.Second,
	TotalBudget:    90 * time.Second,
	JitterFraction: 0.2,
}

// retryInterceptor returns an interceptor that retries idempotent failures
// (Unavailable, DeadlineExceeded, Aborted) per the supplied policy.
func retryInterceptor(p RetryPolicy) connect.UnaryInterceptorFunc {
	return func(next connect.UnaryFunc) connect.UnaryFunc {
		return func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) {
			if p.MaxAttempts < 1 {
				return next(ctx, req)
			}
			deadline := time.Time{}
			if p.TotalBudget > 0 {
				deadline = time.Now().Add(p.TotalBudget)
			}
			backoff := p.InitialBackoff
			var resp connect.AnyResponse
			var err error
			for attempt := 0; attempt < p.MaxAttempts; attempt++ {
				resp, err = next(ctx, req)
				if err == nil {
					return resp, nil
				}
				if !shouldRetry(err) {
					return nil, wrap(err)
				}
				if attempt == p.MaxAttempts-1 {
					break
				}
				wait := withJitter(backoff, p.JitterFraction)
				if !deadline.IsZero() && time.Now().Add(wait).After(deadline) {
					break
				}
				select {
				case <-time.After(wait):
				case <-ctx.Done():
					return nil, wrap(ctx.Err())
				}
				backoff *= 2
				if backoff > p.MaxBackoff {
					backoff = p.MaxBackoff
				}
			}
			return nil, wrap(err)
		}
	}
}

func shouldRetry(err error) bool {
	switch connect.CodeOf(err) {
	case connect.CodeUnavailable, connect.CodeDeadlineExceeded, connect.CodeAborted:
		return true
	}
	return false
}

func withJitter(d time.Duration, frac float64) time.Duration {
	if frac <= 0 {
		return d
	}
	delta := float64(d) * frac
	noise := (rand.Float64()*2 - 1) * delta
	out := time.Duration(float64(d) + noise)
	if out < 0 {
		out = 0
	}
	return out
}
