package workflows

import (
	"context"
	"errors"
	"testing"
	"time"

	"go.temporal.io/sdk/client"
)

// TestDialWithRetry_SucceedsAfterTransientFailures models the cold-start JWKS
// race: the first attempts fail with an "unauthorized" error, then the server's
// key cache refreshes and the dial succeeds.
func TestDialWithRetry_SucceedsAfterTransientFailures(t *testing.T) {
	calls := 0
	dial := func() (client.Client, error) {
		calls++
		if calls < 3 {
			return nil, errors.New("failed reaching server: Request unauthorized.")
		}
		return nil, nil
	}
	c, err := dialWithRetry(context.Background(), dial, time.Second, time.Millisecond, 5*time.Millisecond, nil)
	if err != nil {
		t.Fatalf("expected success after retries, got: %v", err)
	}
	if c != nil {
		t.Fatalf("expected the stub's nil client, got non-nil")
	}
	if calls != 3 {
		t.Fatalf("expected 3 dial attempts, got %d", calls)
	}
}

// TestDialWithRetry_ExhaustsBudget asserts a persistently failing dial returns the
// last error (wrapped) once the budget elapses.
func TestDialWithRetry_ExhaustsBudget(t *testing.T) {
	calls := 0
	wantErr := errors.New("failed reaching server: Request unauthorized.")
	dial := func() (client.Client, error) {
		calls++
		return nil, wantErr
	}
	_, err := dialWithRetry(context.Background(), dial, 20*time.Millisecond, time.Millisecond, 5*time.Millisecond, nil)
	if err == nil {
		t.Fatal("expected an error after exhausting the budget")
	}
	if !errors.Is(err, wantErr) {
		t.Fatalf("expected the wrapped last error, got: %v", err)
	}
	if calls < 2 {
		t.Fatalf("expected multiple attempts within the budget, got %d", calls)
	}
}

// TestDialWithRetry_HealthySucceedsFirstTry asserts a healthy server connects on
// the first attempt with no backoff delay.
func TestDialWithRetry_HealthySucceedsFirstTry(t *testing.T) {
	calls := 0
	dial := func() (client.Client, error) {
		calls++
		return nil, nil
	}
	start := time.Now()
	if _, err := dialWithRetry(context.Background(), dial, time.Second, 50*time.Millisecond, time.Second, nil); err != nil {
		t.Fatalf("expected immediate success, got: %v", err)
	}
	if calls != 1 {
		t.Fatalf("expected exactly 1 attempt, got %d", calls)
	}
	if elapsed := time.Since(start); elapsed > 20*time.Millisecond {
		t.Fatalf("expected no backoff delay on first-try success, took %s", elapsed)
	}
}

// TestDialWithRetry_RespectsContextCancellation asserts a canceled context aborts
// the retry loop promptly.
func TestDialWithRetry_RespectsContextCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	calls := 0
	dial := func() (client.Client, error) {
		calls++
		if calls == 1 {
			cancel() // cancel before the first backoff sleep
		}
		return nil, errors.New("still starting")
	}
	_, err := dialWithRetry(ctx, dial, 5*time.Second, 50*time.Millisecond, time.Second, nil)
	if err == nil {
		t.Fatal("expected a cancellation error")
	}
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("expected context.Canceled, got: %v", err)
	}
}