package grpcclient

import (
	"context"
	"math"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	backoffBase       = 100 * time.Millisecond
	backoffMaxRetries = 8
	backoffMax        = 30 * time.Second
)

// retryable returns true when the status code is a transient error worth
// retrying within the retry budget.
func retryable(code codes.Code) bool {
	switch code {
	case codes.Unavailable, codes.ResourceExhausted, codes.Internal:
		return true
	default:
		return false
	}
}

// doBackoff sleeps for an interval derived from attempt and the
// server-supplied retryAfterMs.  It respects context cancellation.
func doBackoff(ctx context.Context, attempt int, retryAfterMs int) error {
	var delay time.Duration
	if retryAfterMs > 0 {
		delay = time.Duration(retryAfterMs) * time.Millisecond
	} else {
		delay = time.Duration(float64(backoffBase) * math.Pow(2, float64(attempt)))
		if delay > backoffMax {
			delay = backoffMax
		}
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(delay):
		return nil
	}
}

// Retry calls fn repeatedly (up to c.maxRetries) until it succeeds,
// ctx is cancelled, or a non-retryable error is returned.
// fn should perform a single gRPC operation that may fail transiently.
func (c *Client) Retry(ctx context.Context, fn func() error) error {
	var lastErr error
	for attempt := 0; attempt <= c.maxRetries; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}
		lastErr = fn()
		if lastErr == nil {
			return nil
		}
		st, _ := status.FromError(lastErr)
		if !retryable(st.Code()) {
			return lastErr
		}
		if attempt < c.maxRetries {
			if err := doBackoff(ctx, attempt, 0); err != nil {
				return err
			}
		}
	}
	return lastErr
}