// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package auth

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"

	"github.com/cenkalti/backoff/v5"
	"go.uber.org/mock/gomock"
	"golang.org/x/oauth2"

	rt "github.com/stacklok/toolhive/pkg/container/runtime"
	statusMocks "github.com/stacklok/toolhive/pkg/workloads/statuses/mocks"
)

// mockStatusUpdater adapts a mock statuses.StatusManager to auth.StatusUpdater for testing
type mockStatusUpdater struct {
	sm *statusMocks.MockStatusManager
}

func newMockStatusUpdater(ctrl *gomock.Controller) (*mockStatusUpdater, *statusMocks.MockStatusManager) {
	mockSM := statusMocks.NewMockStatusManager(ctrl)
	return &mockStatusUpdater{sm: mockSM}, mockSM
}

func (m *mockStatusUpdater) SetWorkloadStatus(ctx context.Context, workloadName string, status rt.WorkloadStatus, reason string) error {
	return m.sm.SetWorkloadStatus(ctx, workloadName, status, reason)
}

// mockTokenSource is a simple mock implementation of oauth2.TokenSource for testing.
// It uses a callback function to allow flexible token/error configuration.
type mockTokenSource struct {
	mu        sync.Mutex
	tokenFn   func() (*oauth2.Token, error)
	callCount int
	notifyAt  int
	notify    chan struct{}
}

func newMockTokenSource() *mockTokenSource {
	return &mockTokenSource{
		tokenFn: func() (*oauth2.Token, error) {
			return nil, errors.New("no token configured")
		},
	}
}

func (m *mockTokenSource) setTokenFn(fn func() (*oauth2.Token, error)) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.tokenFn = fn
}

// notifyOnCall returns a channel that is closed when Token() is called for the nth time.
// Useful in tests to synchronise without time.Sleep.
func (m *mockTokenSource) notifyOnCall(n int) <-chan struct{} {
	m.mu.Lock()
	defer m.mu.Unlock()
	ch := make(chan struct{})
	m.notifyAt = n
	m.notify = ch
	return ch
}

func (m *mockTokenSource) Token() (*oauth2.Token, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.callCount++
	tok, err := m.tokenFn()
	if m.notify != nil && m.callCount >= m.notifyAt {
		close(m.notify)
		m.notify = nil
	}
	return tok, err
}

// createRetrieveError creates an error for testing token failures. ErrorCode
// is left unset, mirroring what golang.org/x/oauth2 produces when the response
// body is not a parseable RFC 6749 error response (e.g. an HTML page from a
// WAF or load balancer).
func createRetrieveError(statusCode int, body string) *oauth2.RetrieveError {
	response := &http.Response{
		StatusCode: statusCode,
		Body:       http.NoBody,
	}
	return &oauth2.RetrieveError{
		Response: response,
		Body:     []byte(body),
	}
}

// createRetrieveErrorWithCode is like createRetrieveError but also sets the
// ErrorCode field, mirroring what golang.org/x/oauth2 populates when the
// server responds with a parseable JSON error body containing an "error"
// field.
func createRetrieveErrorWithCode(statusCode int, errorCode, body string) *oauth2.RetrieveError {
	err := createRetrieveError(statusCode, body)
	err.ErrorCode = errorCode
	return err
}

func TestMonitoredTokenSource_SuccessfulTokenRetrieval(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, _ := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	validToken := &oauth2.Token{
		AccessToken:  "test-access-token",
		RefreshToken: "test-refresh-token",
		Expiry:       time.Now().Add(time.Hour),
	}
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		return validToken, nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)

	// Test successful token retrieval
	token, err := ats.Token()
	if err != nil {
		t.Fatalf("Expected no error, got %v", err)
	}

	if token.AccessToken != "test-access-token" {
		t.Errorf("Expected access token 'test-access-token', got %s", token.AccessToken)
	}

	// Should not have called SetWorkloadStatus for successful retrieval
	// (no expectations set means we expect it not to be called)
}

func TestMonitoredTokenSource_AuthenticationErrorMarksUnauthenticated(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, statusManager := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	// Create an error that simulates token retrieval failure
	retrieveErr := createRetrieveErrorWithCode(http.StatusBadRequest, "invalid_grant", `{"error":"invalid_grant","error_description":"refresh token expired"}`)
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		return nil, retrieveErr
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)

	// Expect SetWorkloadStatus to be called with unauthenticated status
	statusManager.EXPECT().
		SetWorkloadStatus(
			gomock.Any(),
			"test-workload",
			rt.WorkloadStatusUnauthenticated,
			gomock.Any(),
		).
		DoAndReturn(func(_ context.Context, _ string, _ rt.WorkloadStatus, reason string) error {
			if !strings.Contains(reason, "invalid_grant") {
				t.Errorf("Expected reason to contain 'invalid_grant', got %s", reason)
			}
			return nil
		}).
		Times(1)

	// Token retrieval should fail and mark as unauthenticated
	_, err := ats.Token()
	if err == nil {
		t.Fatal("Expected error, got nil")
	}

	// Give a moment for the async call to complete
	time.Sleep(50 * time.Millisecond)
}

func TestMonitoredTokenSource_ErrorMarksUnauthenticated(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, statusManager := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	// Any error should mark as unauthenticated
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		return nil, errors.New("some generic error")
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)

	// Expect SetWorkloadStatus to be called for any error
	statusManager.EXPECT().
		SetWorkloadStatus(
			gomock.Any(),
			"test-workload",
			rt.WorkloadStatusUnauthenticated,
			gomock.Any(),
		).
		Return(nil).
		Times(1)

	// Token retrieval should fail and mark as unauthenticated
	_, err := ats.Token()
	if err == nil {
		t.Fatal("Expected error, got nil")
	}

	// Give a moment for the async call to complete
	time.Sleep(50 * time.Millisecond)
}

func TestMonitoredTokenSource_BackgroundMonitoring(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, statusManager := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	callCount := 0
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		callCount++
		if callCount == 1 {
			// First call: return valid token with short expiry
			return &oauth2.Token{
				AccessToken:  "test-token",
				RefreshToken: "test-refresh",
				Expiry:       time.Now().Add(500 * time.Millisecond),
			}, nil
		}
		// Subsequent calls: return authentication error
		retrieveErr := createRetrieveErrorWithCode(http.StatusUnauthorized, "invalid_token", `{"error":"invalid_token"}`)
		return nil, retrieveErr
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)

	// Expect SetWorkloadStatus to be called when auth error occurs
	statusManager.EXPECT().
		SetWorkloadStatus(
			gomock.Any(),
			"test-workload",
			rt.WorkloadStatusUnauthenticated,
			gomock.Any(),
		).
		Return(nil).
		Times(1)

	ats.StartBackgroundMonitoring()

	// Wait for token to expire and background monitoring to detect failure
	// The timer is scheduled for when token expires (500ms), then it processes the error
	// Need enough time for: initial timer (1ms) + token expiry (500ms) + error processing
	time.Sleep(2 * time.Second)

	// Verify monitoring stopped by checking that SetWorkloadStatus was called
	// (the mock expectations already verify this)
}

func TestMonitoredTokenSource_BackgroundMonitoringStopsOnAnyError(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, statusManager := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	callCount := 0
	// Use a generic error - should mark as unauthenticated and stop monitoring
	genericErr := errors.New("network timeout")
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		callCount++
		if callCount == 1 {
			// First call: return valid token with short expiry
			return &oauth2.Token{
				AccessToken:  "test-token",
				RefreshToken: "test-refresh",
				Expiry:       time.Now().Add(500 * time.Millisecond),
			}, nil
		}
		// Subsequent calls: return generic error (should mark as unauthenticated)
		return nil, genericErr
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)

	// Expect SetWorkloadStatus to be called when any error occurs
	statusManager.EXPECT().
		SetWorkloadStatus(
			gomock.Any(),
			"test-workload",
			rt.WorkloadStatusUnauthenticated,
			gomock.Any(),
		).
		Return(nil).
		Times(1)

	ats.StartBackgroundMonitoring()

	// Wait for token to expire and background monitoring to detect failure
	// Flow: initial timer (1ms) → first check (gets token) → reschedule → wait → second check (gets error) → mark unauthenticated
	time.Sleep(2 * time.Second)

	// Verify monitoring stopped by checking that SetWorkloadStatus was called
	// (the mock expectations already verify this)
}

func TestMonitoredTokenSource_ExpiredTokenHandling(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, _ := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	// Return an already-expired token (oauth2 library should try to refresh)
	expiredToken := &oauth2.Token{
		AccessToken:  "expired-token",
		RefreshToken: "refresh-token",
		Expiry:       time.Now().Add(-time.Hour),
	}
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		return expiredToken, nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)

	// Should not mark as unauthenticated just for expired token
	// (oauth2 library should handle refresh; we only mark on actual auth errors)
	// (no expectations set means we expect SetWorkloadStatus not to be called)

	ats.StartBackgroundMonitoring()

	// Wait a bit for monitoring to check
	time.Sleep(200 * time.Millisecond)

	cancel()
}

func TestMonitoredTokenSource_StopMonitoring(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, _ := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		return &oauth2.Token{
			AccessToken:  "test-token",
			RefreshToken: "refresh",
			Expiry:       time.Now().Add(time.Hour),
		}, nil
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)
	ats.StartBackgroundMonitoring()

	// Wait a bit to ensure monitoring started
	time.Sleep(100 * time.Millisecond)

	// Stop monitoring via context cancellation
	cancel()

	// Wait a bit for monitoring to stop
	time.Sleep(100 * time.Millisecond)

	// Verify monitoring stopped - context cancellation is handled internally
	// We can verify by ensuring no unexpected SetWorkloadStatus calls
	// (test passes if no errors occur)
}

func TestMonitoredTokenSource_MultipleCallsToToken(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, statusManager := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	retrieveErr := createRetrieveErrorWithCode(http.StatusUnauthorized, "invalid_token", `{"error":"invalid_token"}`)
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		return nil, retrieveErr
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)

	statusManager.EXPECT().
		SetWorkloadStatus(
			gomock.Any(),
			"test-workload",
			rt.WorkloadStatusUnauthenticated,
			gomock.Any(),
		).
		Return(nil).
		Times(3) // Each Token() call will mark it

	// Call Token() multiple times
	for i := 0; i < 3; i++ {
		_, err := ats.Token()
		if err == nil {
			t.Fatal("Expected error, got nil")
		}
	}

	time.Sleep(50 * time.Millisecond)
}

// TestTransientRefresher_SingleflightDeduplicatesConcurrentRetries verifies that
// concurrent Refresh() calls are funnelled through a single retry loop via
// singleflight, so the underlying token source is not hammered by independent
// retry loops ("thundering herd").
func TestTransientRefresher_SingleflightDeduplicatesConcurrentRetries(t *testing.T) {
	t.Parallel()

	const numCallers = 10

	tokenSource := newMockTokenSource()
	recoveredToken := &oauth2.Token{
		AccessToken: "recovered-token",
		Expiry:      time.Now().Add(time.Hour),
	}
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		return recoveredToken, nil // retry always succeeds immediately
	})

	transientErr := &net.OpError{
		Op: "dial", Net: "tcp",
		Err: &os.SyscallError{Syscall: "connect", Err: syscall.ECONNREFUSED},
	}

	// Two-phase synchronisation to guarantee deterministic singleflight deduplication:
	//
	// Phase 1 (beforeEntry): all numCallers goroutines arrive here before calling
	// Refresh. A WaitGroup barrier ensures they are all released simultaneously,
	// so they race to group.Do together.
	//
	// Phase 2 (afterEntry): the singleflight leader enters this hook from inside
	// group.Do and waits until all numCallers goroutines have signalled they are
	// about to call Refresh (i.e. finished Phase 1). At that point the leader is
	// still running inside Do, so any follower that subsequently calls Do will be
	// deduplicated rather than starting an independent retry loop.
	//
	// Without Phase 2 the leader could return before late goroutines reached Do,
	// causing each to start its own singleflight group and hammer the token source.
	allAtSingleflight := make(chan struct{})
	var atSingleflight sync.WaitGroup
	atSingleflight.Add(numCallers)
	var closeOnce sync.Once

	var beforeDo sync.WaitGroup
	beforeDo.Add(numCallers)

	ctx := context.Background()
	refresher := &transientRefresher{
		source:     tokenSource,
		workload:   "test-workload",
		newBackOff: fastBackOff,
		beforeEntry: func() {
			// Phase 1: barrier — release all goroutines simultaneously.
			atSingleflight.Done()
			closeOnce.Do(func() {
				atSingleflight.Wait()
				close(allAtSingleflight)
			})
			<-allAtSingleflight
			// Signal: I am about to call group.Do.
			beforeDo.Done()
		},
		afterEntry: func() {
			// Phase 2: leader waits until all goroutines have signalled they are
			// about to call group.Do, so the group is fully formed before retry returns.
			beforeDo.Wait()
		},
	}

	var wg sync.WaitGroup
	tokens := make([]*oauth2.Token, numCallers)
	errs := make([]error, numCallers)
	for i := range numCallers {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			tokens[idx], errs[idx] = refresher.Refresh(ctx, transientErr)
		}(i)
	}

	// Guard against a deadlock in the synchronisation barriers turning into a
	// silent hang. Use the test deadline if available; otherwise fall back to a
	// conservative fixed timeout.
	done := make(chan struct{})
	go func() { wg.Wait(); close(done) }()
	timeout := 10 * time.Second
	if deadline, ok := t.Deadline(); ok {
		timeout = time.Until(deadline) - 500*time.Millisecond
	}
	select {
	case <-done:
	case <-time.After(timeout):
		t.Fatal("test timed out — likely deadlock in synchronisation barriers")
	}

	// All callers must succeed with the recovered token.
	for i := range numCallers {
		if errs[i] != nil {
			t.Errorf("caller %d: unexpected error: %v", i, errs[i])
		}
		if tokens[i] == nil || tokens[i].AccessToken != "recovered-token" {
			t.Errorf("caller %d: expected recovered-token, got %v", i, tokens[i])
		}
	}

	// KEY ASSERTION: exactly 1 call via singleflight, not numCallers independent calls.
	// Independent retry loops would produce up to numCallers calls.
	tokenSource.mu.Lock()
	calls := tokenSource.callCount
	tokenSource.mu.Unlock()
	if calls != 1 {
		t.Errorf("expected 1 tokenSource.Token() call (singleflight deduplication), got %d", calls)
	}
}

// --- helpers for new tests ---

// timeoutNetError is a minimal net.Error with Timeout() == true.
type timeoutNetError struct{}

func (*timeoutNetError) Error() string   { return "i/o timeout" }
func (*timeoutNetError) Timeout() bool   { return true }
func (*timeoutNetError) Temporary() bool { return true }

var _ net.Error = (*timeoutNetError)(nil)

// fastBackOff returns a backoff with very short intervals so retry tests run quickly.
func fastBackOff() backoff.BackOff {
	b := backoff.NewExponentialBackOff()
	b.InitialInterval = 10 * time.Millisecond
	b.MaxInterval = 50 * time.Millisecond
	b.Reset()
	return b
}

// --- error classification via background monitor ---

// TestMonitoredTokenSource_BackgroundMonitor_ErrorClassification verifies that the
// background monitor correctly distinguishes transient network errors (which trigger
// retries without marking the workload unauthenticated) from non-transient errors
// (which immediately mark the workload as unauthenticated and stop monitoring).
func TestMonitoredTokenSource_BackgroundMonitor_ErrorClassification(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name        string
		err         error
		isTransient bool // true → monitor retries; false → monitor marks unauthenticated
	}{
		// Non-transient: plain errors and OAuth protocol failures (4xx with a
		// populated RFC 6749 error code) must fail fast.
		{name: "plain error", err: errors.New("some error"), isTransient: false},
		{name: "context.Canceled", err: context.Canceled, isTransient: false},
		{name: "context.DeadlineExceeded", err: context.DeadlineExceeded, isTransient: false},
		{name: "oauth2.RetrieveError 400 invalid_grant", err: createRetrieveErrorWithCode(http.StatusBadRequest, "invalid_grant", `{"error":"invalid_grant"}`), isTransient: false},
		{name: "oauth2.RetrieveError 401 invalid_client", err: createRetrieveErrorWithCode(http.StatusUnauthorized, "invalid_client", `{"error":"invalid_client"}`), isTransient: false},
		{name: "oauth2.RetrieveError 403 unauthorized_client", err: createRetrieveErrorWithCode(http.StatusForbidden, "unauthorized_client", `{"error":"unauthorized_client"}`), isTransient: false},
		{name: "oauth2.RetrieveError nil response", err: &oauth2.RetrieveError{}, isTransient: false},
		// Transient: network-level errors must be retried.
		{name: "*net.DNSError timeout", err: &net.DNSError{Err: "i/o timeout", Name: "example.com", IsTimeout: true}, isTransient: true},
		{name: "*net.OpError connection refused", err: &net.OpError{Op: "dial", Net: "tcp", Err: &os.SyscallError{Syscall: "connect", Err: syscall.ECONNREFUSED}}, isTransient: true},
		{name: "*url.Error wrapping *net.OpError", err: &url.Error{Op: "Post", URL: "https://example.com/token", Err: &net.OpError{Op: "dial", Net: "tcp", Err: &os.SyscallError{Syscall: "connect", Err: syscall.ECONNREFUSED}}}, isTransient: true},
		{name: "net.Error timeout", err: &timeoutNetError{}, isTransient: true},
		// Transient: OAuth server 5xx errors (load balancer, server restart).
		{name: "oauth2.RetrieveError 500", err: createRetrieveError(http.StatusInternalServerError, "Internal Server Error"), isTransient: true},
		{name: "oauth2.RetrieveError 502", err: createRetrieveError(http.StatusBadGateway, "Bad Gateway"), isTransient: true},
		{name: "oauth2.RetrieveError 503", err: createRetrieveError(http.StatusServiceUnavailable, "Service Unavailable"), isTransient: true},
		{name: "oauth2.RetrieveError 504", err: createRetrieveError(http.StatusGatewayTimeout, "Gateway Timeout"), isTransient: true},
		// Transient: 4xx without an RFC 6749 error code in the body.
		// These are infrastructure-level errors (WAF, CDN, proxy) that
		// commonly resolve on their own, not OAuth protocol failures.
		{name: "oauth2.RetrieveError 401 with HTML body", err: createRetrieveError(http.StatusUnauthorized, "<html><body>Unauthorized</body></html>"), isTransient: true},
		{name: "oauth2.RetrieveError 403 WAF block", err: createRetrieveError(http.StatusForbidden, "<html><body>Cloudflare Firewall Block</body></html>"), isTransient: true},
		{name: "oauth2.RetrieveError 400 with empty body", err: createRetrieveError(http.StatusBadRequest, ""), isTransient: true},
		{name: "oauth2.RetrieveError 408 request timeout", err: createRetrieveError(http.StatusRequestTimeout, ""), isTransient: true},
		// Transient: 429 Too Many Requests is retryable per HTTP standard
		// regardless of body content.
		{name: "oauth2.RetrieveError 429 empty body", err: createRetrieveError(http.StatusTooManyRequests, ""), isTransient: true},
		{name: "oauth2.RetrieveError 429 with rate-limit error code", err: createRetrieveErrorWithCode(http.StatusTooManyRequests, "rate_limit_exceeded", `{"error":"rate_limit_exceeded"}`), isTransient: true},
		// Transient: unparsable OAuth responses (HTML from load balancer on 200).
		{name: "oauth2 cannot parse json", err: fmt.Errorf("oauth2: cannot parse json: invalid character '<'"), isTransient: true},
		{name: "wrapped oauth2 parse error", err: fmt.Errorf("refresh failed: %w", fmt.Errorf("oauth2: cannot parse json: invalid character '<'")), isTransient: true},
		{name: "oauth2 cannot parse response", err: fmt.Errorf("oauth2: cannot parse response: invalid URL escape"), isTransient: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			tokenSource := newMockTokenSource()
			tokenSource.setTokenFn(func() (*oauth2.Token, error) {
				if tokenSource.callCount == 1 {
					// Initial tick: short-lived token so the monitor retries quickly.
					return &oauth2.Token{
						AccessToken: "initial-token",
						Expiry:      time.Now().Add(10 * time.Millisecond),
					}, nil
				}
				return nil, tt.err
			})

			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			if tt.isTransient {
				// Transient: SetWorkloadStatus must NOT be called — no EXPECT set.
				statusUpdater, _ := newMockStatusUpdater(ctrl)
				retrying := tokenSource.notifyOnCall(2)

				ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)
				ats.refresher.newBackOff = fastBackOff
				ats.StartBackgroundMonitoring()

				<-retrying // Ensure the retry loop has been entered before cancelling.
				cancel()
				<-ats.Stopped()
			} else {
				// Non-transient: SetWorkloadStatus must be called exactly once.
				statusUpdater, statusManager := newMockStatusUpdater(ctrl)
				statusManager.EXPECT().
					SetWorkloadStatus(
						gomock.Any(),
						"test-workload",
						rt.WorkloadStatusUnauthenticated,
						gomock.Any(),
					).
					Return(nil).
					Times(1)

				ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)
				ats.refresher.newBackOff = fastBackOff
				ats.StartBackgroundMonitoring()

				<-ats.Stopped() // Monitor stops itself after marking unauthenticated.
			}
		})
	}
}

// TestIsPermanentTokenEndpointError verifies that isPermanentTokenEndpointError
// is the strict inverse of classifyOAuthRetrieveError on the
// *oauth2.RetrieveError branch (with a non-nil Response). The DCR/CIMD
// remediation Warn fires only when the OAuth server returned a structured
// RFC 6749 error code; non-spec-compliant responses (HTML pages from a WAF,
// CDN, or reverse proxy) should not trigger that Warn because they carry no
// OAuth-protocol verdict.
//
// Existing Token() / markAsUnauthenticated tests reach this function through
// indirect call paths and yield 100% line coverage, but none of them assert
// on the boolean it returns. This test pins the behavioral contract directly.
func TestIsPermanentTokenEndpointError(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name        string
		err         error
		isPermanent bool
	}{
		// Not an *oauth2.RetrieveError at all — no OAuth verdict to act on.
		{name: "plain error", err: errors.New("some error"), isPermanent: false},
		{name: "*oauth2.RetrieveError with nil Response", err: &oauth2.RetrieveError{}, isPermanent: false},
		// Transient HTTP-level conditions — never permanent.
		{name: "5xx server error", err: createRetrieveError(http.StatusInternalServerError, "Internal Server Error"), isPermanent: false},
		{name: "429 Too Many Requests", err: createRetrieveError(http.StatusTooManyRequests, ""), isPermanent: false},
		{name: "408 Request Timeout", err: createRetrieveError(http.StatusRequestTimeout, ""), isPermanent: false},
		// 4xx without an RFC 6749 error code — infrastructure response, no OAuth verdict.
		{name: "401 with HTML body (WAF)", err: createRetrieveError(http.StatusUnauthorized, "<html><body>Unauthorized</body></html>"), isPermanent: false},
		{name: "403 with HTML body (Cloudflare)", err: createRetrieveError(http.StatusForbidden, "<html><body>Firewall Block</body></html>"), isPermanent: false},
		// 4xx with an RFC 6749 error code — OAuth server rendered a verdict.
		{name: "400 invalid_grant", err: createRetrieveErrorWithCode(http.StatusBadRequest, "invalid_grant", `{"error":"invalid_grant"}`), isPermanent: true},
		{name: "401 invalid_client", err: createRetrieveErrorWithCode(http.StatusUnauthorized, "invalid_client", `{"error":"invalid_client"}`), isPermanent: true},
		{name: "403 unauthorized_client", err: createRetrieveErrorWithCode(http.StatusForbidden, "unauthorized_client", `{"error":"unauthorized_client"}`), isPermanent: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			got := isPermanentTokenEndpointError(tt.err)
			if got != tt.isPermanent {
				t.Errorf("isPermanentTokenEndpointError(%v) = %v, want %v",
					tt.err, got, tt.isPermanent)
			}
		})
	}
}

// TestIsTransientNetworkError_AgainstRealOAuth2Library is a contract test,
// not a unit test of the package's own logic. It pins the assumption that
// isTransientRetrieveError relies on: that golang.org/x/oauth2 populates
// RetrieveError.ErrorCode iff the response carries a parseable RFC 6749
// 'error' field (whether JSON or form-encoded), and leaves ErrorCode empty
// for non-spec-compliant response shapes (HTML pages from a WAF, CDN, or
// reverse proxy).
//
// Cases here are deliberately limited to response shapes where the
// synthetic test helpers (createRetrieveError, createRetrieveErrorWithCode)
// could plausibly diverge from reality. Cases unambiguously covered by the
// synthetic table (clearly populated ErrorCode JSON, status-code-only
// branches like 5xx and 429) are intentionally not duplicated here.
func TestIsTransientNetworkError_AgainstRealOAuth2Library(t *testing.T) {
	t.Parallel()

	const refreshToken = "test-refresh-token"

	tests := []struct {
		name        string
		handler     http.HandlerFunc
		isTransient bool
	}{
		{
			// HTML 4xx is the canonical WAF/CDN block shape. Pinning that
			// the library leaves ErrorCode empty here is what underpins the
			// "infrastructure error" branch of isTransientRetrieveError.
			name: "403 with HTML body (Cloudflare WAF)",
			handler: func(w http.ResponseWriter, _ *http.Request) {
				w.Header().Set("Content-Type", "text/html")
				w.WriteHeader(http.StatusForbidden)
				_, _ = w.Write([]byte("<html><body>Cloudflare Firewall Block</body></html>"))
			},
			isTransient: true,
		},
		{
			name: "401 with HTML body",
			handler: func(w http.ResponseWriter, _ *http.Request) {
				w.Header().Set("Content-Type", "text/html")
				w.WriteHeader(http.StatusUnauthorized)
				_, _ = w.Write([]byte("<html><body>Unauthorized</body></html>"))
			},
			isTransient: true,
		},
		{
			// Form-encoded error responses are non-spec but supported by
			// the library. Pinning that the library DOES populate ErrorCode
			// from form-encoded bodies — a synthetic helper used naively
			// (createRetrieveError without WithCode) would lie about this.
			name: "400 with form-encoded invalid_grant",
			handler: func(w http.ResponseWriter, _ *http.Request) {
				w.Header().Set("Content-Type", "application/x-www-form-urlencoded")
				w.WriteHeader(http.StatusBadRequest)
				_, _ = w.Write([]byte("error=invalid_grant&error_description=refresh+token+expired"))
			},
			isTransient: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			server := httptest.NewServer(tt.handler)
			t.Cleanup(server.Close)

			cfg := &oauth2.Config{
				ClientID:     "test-client",
				ClientSecret: "test-secret",
				Endpoint:     oauth2.Endpoint{TokenURL: server.URL},
			}

			expired := &oauth2.Token{
				AccessToken:  "expired-access-token",
				RefreshToken: refreshToken,
				Expiry:       time.Now().Add(-time.Hour),
			}

			_, err := cfg.TokenSource(context.Background(), expired).Token()
			if err == nil {
				t.Fatalf("expected refresh to fail, got nil error")
			}

			got := isTransientNetworkError(err)
			if got != tt.isTransient {
				t.Errorf("isTransientNetworkError(%v) = %v, want %v",
					err, got, tt.isTransient)
			}
		})
	}
}

// --- background monitor transient-error behaviour ---

// TestMonitoredTokenSource_TransientErrorRetriesAndSucceeds verifies that when the
// background monitor encounters a transient network error it retries with backoff and,
// once the network recovers, does NOT mark the workload as unauthenticated.
func TestMonitoredTokenSource_TransientErrorRetriesAndSucceeds(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	// No SetWorkloadStatus calls expected — the workload must stay authenticated.
	statusUpdater, _ := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	transientErr := &net.OpError{Op: "dial", Net: "tcp", Err: &os.SyscallError{Syscall: "connect", Err: syscall.ECONNREFUSED}}
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		switch tokenSource.callCount {
		case 1:
			// Initial monitor kick: valid token that expires soon.
			return &oauth2.Token{
				AccessToken: "initial-token",
				Expiry:      time.Now().Add(10 * time.Millisecond),
			}, nil
		case 2, 3, 4:
			// Transient failures during the retry window.
			return nil, transientErr
		default:
			// Network recovered — return a long-lived token.
			return &oauth2.Token{
				AccessToken: "renewed-token",
				Expiry:      time.Now().Add(time.Hour),
			}, nil
		}
	})

	// Wait for call 5: the recovery token return.
	recovered := tokenSource.notifyOnCall(5)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)
	ats.refresher.newBackOff = fastBackOff
	ats.StartBackgroundMonitoring()

	// Block until the monitor has successfully recovered, then stop it.
	<-recovered
	cancel()
	<-ats.Stopped()
	// gomock verifies SetWorkloadStatus was NOT called (no EXPECT set).
}

// TestMonitoredTokenSource_TransientErrorContextCancellation verifies that cancelling
// the monitoring context while the retry loop is running does NOT mark the workload
// as unauthenticated (the workload was simply removed, not broken).
func TestMonitoredTokenSource_TransientErrorContextCancellation(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	// No SetWorkloadStatus calls expected.
	statusUpdater, _ := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	transientErr := &net.OpError{Op: "dial", Net: "tcp", Err: &os.SyscallError{Syscall: "connect", Err: syscall.ECONNREFUSED}}
	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		if tokenSource.callCount == 1 {
			return &oauth2.Token{
				AccessToken: "initial-token",
				Expiry:      time.Now().Add(10 * time.Millisecond),
			}, nil
		}
		// All subsequent calls: perpetual transient error.
		return nil, transientErr
	})

	// Wait for the first retry attempt before cancelling.
	retrying := tokenSource.notifyOnCall(2)

	ctx, cancel := context.WithCancel(context.Background())

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)
	ats.refresher.newBackOff = fastBackOff
	ats.StartBackgroundMonitoring()

	// Cancel once we know the retry loop is running, then wait for clean exit.
	<-retrying
	cancel()
	<-ats.Stopped()
	// gomock verifies SetWorkloadStatus was NOT called (no EXPECT set).
}

// TestMonitoredTokenSource_TransientThenNonTransientMarksUnauthenticated verifies that
// after a few retryable failures, a non-transient error (e.g. 401) stops the retry loop
// and marks the workload as unauthenticated exactly once.
func TestMonitoredTokenSource_TransientThenNonTransientMarksUnauthenticated(t *testing.T) {
	t.Parallel()
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	statusUpdater, statusManager := newMockStatusUpdater(ctrl)
	tokenSource := newMockTokenSource()

	statusManager.EXPECT().
		SetWorkloadStatus(
			gomock.Any(),
			"test-workload",
			rt.WorkloadStatusUnauthenticated,
			gomock.Any(),
		).
		Return(nil).
		Times(1)

	transientErr := &net.OpError{Op: "dial", Net: "tcp", Err: &os.SyscallError{Syscall: "connect", Err: syscall.ECONNREFUSED}}
	nonTransientErr := createRetrieveErrorWithCode(http.StatusUnauthorized, "invalid_token", `{"error":"invalid_token"}`)

	tokenSource.setTokenFn(func() (*oauth2.Token, error) {
		switch tokenSource.callCount {
		case 1:
			// Initial tick: short-lived valid token.
			return &oauth2.Token{
				AccessToken: "initial-token",
				Expiry:      time.Now().Add(10 * time.Millisecond),
			}, nil
		case 2, 3:
			// Transient errors — retried.
			return nil, transientErr
		default:
			// Non-transient auth failure — must stop retrying and mark unauthenticated.
			return nil, nonTransientErr
		}
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	ats := NewMonitoredTokenSource(ctx, tokenSource, "test-workload", "", "", statusUpdater)
	ats.refresher.newBackOff = fastBackOff
	ats.StartBackgroundMonitoring()

	// Monitor stops itself after the non-transient error; wait for that.
	<-ats.Stopped()
	// gomock verifies SetWorkloadStatus was called exactly once.
}
