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

package health

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

	"github.com/mark3labs/mcp-go/client/transport"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	"github.com/stacklok/toolhive/pkg/vmcp"
	authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
	"github.com/stacklok/toolhive/pkg/vmcp/mocks"
)

func TestNewHealthChecker(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	t.Cleanup(ctrl.Finish)

	mockClient := mocks.NewMockBackendClient(ctrl)

	tests := []struct {
		name    string
		timeout time.Duration
	}{
		{
			name:    "with timeout",
			timeout: 5 * time.Second,
		},
		{
			name:    "with zero timeout",
			timeout: 0,
		},
	}

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

			checker := NewHealthChecker(mockClient, tt.timeout, 0)
			require.NotNil(t, checker)

			// Type assert to access internals for verification
			hc, ok := checker.(*healthChecker)
			require.True(t, ok)
			assert.Equal(t, mockClient, hc.client)
			assert.Equal(t, tt.timeout, hc.timeout)
		})
	}
}

func TestHealthChecker_CheckHealth_Success(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockClient := mocks.NewMockBackendClient(ctrl)
	mockClient.EXPECT().
		ListCapabilities(gomock.Any(), gomock.Any()).
		Return(&vmcp.CapabilityList{}, nil).
		Times(1)

	checker := NewHealthChecker(mockClient, 5*time.Second, 0)
	target := &vmcp.BackendTarget{
		WorkloadID:   "backend-1",
		WorkloadName: "test-backend",
		BaseURL:      "http://localhost:8080",
	}

	status, err := checker.CheckHealth(context.Background(), target)
	assert.NoError(t, err)
	assert.Equal(t, vmcp.BackendHealthy, status)
}

func TestHealthChecker_CheckHealth_ContextCancellation(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockClient := mocks.NewMockBackendClient(ctrl)
	mockClient.EXPECT().
		ListCapabilities(gomock.Any(), gomock.Any()).
		DoAndReturn(func(ctx context.Context, _ *vmcp.BackendTarget) (*vmcp.CapabilityList, error) {
			<-ctx.Done()
			return nil, ctx.Err()
		}).
		Times(1)

	checker := NewHealthChecker(mockClient, 100*time.Millisecond, 0)
	target := &vmcp.BackendTarget{
		WorkloadID:   "backend-1",
		WorkloadName: "test-backend",
		BaseURL:      "http://localhost:8080",
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel() // Cancel immediately

	status, err := checker.CheckHealth(ctx, target)
	assert.Error(t, err)
	assert.Equal(t, vmcp.BackendUnhealthy, status)
}

func TestHealthChecker_CheckHealth_NoTimeout(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockClient := mocks.NewMockBackendClient(ctrl)
	mockClient.EXPECT().
		ListCapabilities(gomock.Any(), gomock.Any()).
		Return(&vmcp.CapabilityList{}, nil).
		Times(1)

	// Create checker with no timeout
	checker := NewHealthChecker(mockClient, 0, 0)
	target := &vmcp.BackendTarget{
		WorkloadID:   "backend-1",
		WorkloadName: "test-backend",
		BaseURL:      "http://localhost:8080",
	}

	status, err := checker.CheckHealth(context.Background(), target)
	assert.NoError(t, err)
	assert.Equal(t, vmcp.BackendHealthy, status)
}

func TestHealthChecker_CheckHealth_ErrorCategorization(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name           string
		err            error
		expectedStatus vmcp.BackendHealthStatus
		description    string
	}{
		{
			name:           "timeout error",
			err:            fmt.Errorf("context deadline exceeded"),
			expectedStatus: vmcp.BackendUnhealthy,
			description:    "should categorize timeout as unhealthy",
		},
		{
			name:           "connection refused",
			err:            fmt.Errorf("connection refused"),
			expectedStatus: vmcp.BackendUnhealthy,
			description:    "should categorize connection error as unhealthy",
		},
		{
			name:           "authentication failed",
			err:            fmt.Errorf("authentication failed: invalid token"),
			expectedStatus: vmcp.BackendUnauthenticated,
			description:    "should categorize auth failure as unauthenticated",
		},
		{
			name:           "401 unauthorized",
			err:            fmt.Errorf("HTTP 401 unauthorized"),
			expectedStatus: vmcp.BackendUnauthenticated,
			description:    "should categorize 401 as unauthenticated",
		},
		{
			name:           "403 forbidden",
			err:            fmt.Errorf("403 forbidden"),
			expectedStatus: vmcp.BackendUnauthenticated,
			description:    "should categorize 403 as unauthenticated",
		},
		{
			name:           "status code 401",
			err:            fmt.Errorf("status code 401"),
			expectedStatus: vmcp.BackendUnauthenticated,
			description:    "should recognize status code format",
		},
		{
			name:           "request unauthenticated",
			err:            fmt.Errorf("request unauthenticated"),
			expectedStatus: vmcp.BackendUnauthenticated,
			description:    "should recognize request unauthenticated",
		},
		{
			name:           "access denied",
			err:            fmt.Errorf("access denied"),
			expectedStatus: vmcp.BackendUnauthenticated,
			description:    "should recognize access denied",
		},
		{
			name:           "generic error",
			err:            fmt.Errorf("unknown error"),
			expectedStatus: vmcp.BackendUnhealthy,
			description:    "should default unknown errors to unhealthy",
		},
	}

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

			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			mockClient := mocks.NewMockBackendClient(ctrl)
			mockClient.EXPECT().
				ListCapabilities(gomock.Any(), gomock.Any()).
				Return(nil, tt.err).
				Times(1)

			checker := NewHealthChecker(mockClient, 5*time.Second, 0)
			target := &vmcp.BackendTarget{
				WorkloadID:   "backend-1",
				WorkloadName: "test-backend",
				BaseURL:      "http://localhost:8080",
			}

			status, err := checker.CheckHealth(context.Background(), target)
			assert.Error(t, err, tt.description)
			assert.Equal(t, tt.expectedStatus, status, tt.description)
		})
	}
}

func TestCategorizeError(t *testing.T) {
	t.Parallel()

	// Backends with an outgoing auth strategy configured: a 401/403 is the
	// expected response to a no-credential probe and must be treated as healthy.
	targetWithUpstreamInject := &vmcp.BackendTarget{
		AuthConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeUpstreamInject},
	}
	targetWithTokenExchange := &vmcp.BackendTarget{
		AuthConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeTokenExchange},
	}
	targetWithHeaderInjection := &vmcp.BackendTarget{
		AuthConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeHeaderInjection},
	}

	// Backends without an outgoing auth strategy: a 401/403 indicates operator
	// misconfiguration and must surface as BackendUnauthenticated.
	targetNoAuthConfig := &vmcp.BackendTarget{AuthConfig: nil}
	targetUnauthenticatedStrategy := &vmcp.BackendTarget{
		AuthConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeUnauthenticated},
	}

	tests := []struct {
		name           string
		target         *vmcp.BackendTarget
		err            error
		expectedStatus vmcp.BackendHealthStatus
	}{
		{
			name:           "nil error",
			target:         targetNoAuthConfig,
			err:            nil,
			expectedStatus: vmcp.BackendHealthy,
		},

		// Auth errors + outgoing auth configured -> healthy (probe challenge is expected).
		{
			name:           "auth error with upstream_inject strategy is healthy",
			target:         targetWithUpstreamInject,
			err:            vmcp.ErrAuthenticationFailed,
			expectedStatus: vmcp.BackendHealthy,
		},
		{
			name:           "auth error with token_exchange strategy is healthy",
			target:         targetWithTokenExchange,
			err:            vmcp.ErrAuthenticationFailed,
			expectedStatus: vmcp.BackendHealthy,
		},
		{
			name:           "auth error with header_injection strategy is healthy",
			target:         targetWithHeaderInjection,
			err:            vmcp.ErrAuthenticationFailed,
			expectedStatus: vmcp.BackendHealthy,
		},
		{
			name:           "authz error with upstream_inject strategy is healthy",
			target:         targetWithUpstreamInject,
			err:            vmcp.ErrAuthorizationFailed,
			expectedStatus: vmcp.BackendHealthy,
		},
		{
			name:           "string-based auth error with header_injection strategy is healthy",
			target:         targetWithHeaderInjection,
			err:            errors.New("HTTP 401"),
			expectedStatus: vmcp.BackendHealthy,
		},

		// Auth errors + no outgoing auth configured -> unauthenticated (misconfig signal).
		{
			name:           "auth error with nil AuthConfig is unauthenticated (misconfig)",
			target:         targetNoAuthConfig,
			err:            vmcp.ErrAuthenticationFailed,
			expectedStatus: vmcp.BackendUnauthenticated,
		},
		{
			name:           "auth error with StrategyTypeUnauthenticated is unauthenticated (misconfig)",
			target:         targetUnauthenticatedStrategy,
			err:            vmcp.ErrAuthenticationFailed,
			expectedStatus: vmcp.BackendUnauthenticated,
		},
		{
			name:           "authentication failed (string) with nil AuthConfig",
			target:         targetNoAuthConfig,
			err:            errors.New("authentication failed"),
			expectedStatus: vmcp.BackendUnauthenticated,
		},
		{
			name:           "authentication error (string) with nil AuthConfig",
			target:         targetNoAuthConfig,
			err:            errors.New("authentication error: invalid credentials"),
			expectedStatus: vmcp.BackendUnauthenticated,
		},
		{
			name:           "request unauthorized with nil AuthConfig",
			target:         targetNoAuthConfig,
			err:            errors.New("request unauthorized"),
			expectedStatus: vmcp.BackendUnauthenticated,
		},
		{
			name:           "HTTP 401 with nil AuthConfig",
			target:         targetNoAuthConfig,
			err:            errors.New("HTTP 401"),
			expectedStatus: vmcp.BackendUnauthenticated,
		},
		{
			name:           "HTTP 403 with nil AuthConfig",
			target:         targetNoAuthConfig,
			err:            errors.New("HTTP 403"),
			expectedStatus: vmcp.BackendUnauthenticated,
		},
		{
			name:           "nil target with auth error is unauthenticated",
			target:         nil,
			err:            vmcp.ErrAuthenticationFailed,
			expectedStatus: vmcp.BackendUnauthenticated,
		},

		// Non-auth errors: AuthConfig is irrelevant; classification is unchanged.
		{
			name:           "timeout with upstream_inject strategy is still unhealthy",
			target:         targetWithUpstreamInject,
			err:            errors.New("request timeout"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "timeout with nil AuthConfig is unhealthy",
			target:         targetNoAuthConfig,
			err:            errors.New("request timeout"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "deadline exceeded with nil AuthConfig is unhealthy",
			target:         targetNoAuthConfig,
			err:            errors.New("context deadline exceeded"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "connection refused with nil AuthConfig is unhealthy",
			target:         targetNoAuthConfig,
			err:            errors.New("connection refused"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "connection refused with header_injection strategy is still unhealthy",
			target:         targetWithHeaderInjection,
			err:            errors.New("connection refused"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "connection reset with nil AuthConfig is unhealthy",
			target:         targetNoAuthConfig,
			err:            errors.New("connection reset by peer"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "no route to host with nil AuthConfig is unhealthy",
			target:         targetNoAuthConfig,
			err:            errors.New("no route to host"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "network unreachable with nil AuthConfig is unhealthy",
			target:         targetNoAuthConfig,
			err:            errors.New("network is unreachable"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "generic error with nil AuthConfig is unhealthy",
			target:         targetNoAuthConfig,
			err:            errors.New("something went wrong"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
		{
			name:           "generic error with token_exchange strategy is still unhealthy",
			target:         targetWithTokenExchange,
			err:            errors.New("something went wrong"),
			expectedStatus: vmcp.BackendUnhealthy,
		},
	}

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

			status := categorizeError(tt.target, tt.err)
			assert.Equal(t, tt.expectedStatus, status)
		})
	}
}

func TestIsAuthenticationError(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		err       error
		expectErr bool
	}{
		// Positive cases
		{name: "authentication failed", err: errors.New("authentication failed"), expectErr: true},
		{name: "Authentication Failed (uppercase)", err: errors.New("Authentication Failed"), expectErr: true},
		{name: "authentication error", err: errors.New("authentication error: bad token"), expectErr: true},
		{name: "401 unauthorized", err: errors.New("401 unauthorized"), expectErr: true},
		{name: "403 forbidden", err: errors.New("403 forbidden"), expectErr: true},
		{name: "HTTP 401", err: errors.New("HTTP 401"), expectErr: true},
		{name: "HTTP 403", err: errors.New("HTTP 403"), expectErr: true},
		{name: "status code 401", err: errors.New("status code 401"), expectErr: true},
		{name: "status code 403", err: errors.New("status code 403"), expectErr: true},
		{name: "request unauthenticated", err: errors.New("request unauthenticated"), expectErr: true},
		{name: "request unauthorized", err: errors.New("request unauthorized"), expectErr: true},
		{name: "access denied", err: errors.New("access denied"), expectErr: true},

		// mcp-go ErrUnauthorized format: "unauthorized (401)" (reversed order vs "401 unauthorized")
		{name: "unauthorized (401) - mcp-go ErrUnauthorized format", err: errors.New("unauthorized (401)"), expectErr: true},

		// Negative cases - should NOT be detected as auth errors
		{name: "connection refused", err: errors.New("connection refused"), expectErr: false},
		{name: "timeout", err: errors.New("request timeout"), expectErr: false},
		{name: "generic error", err: errors.New("something went wrong"), expectErr: false},
		{name: "404 not found", err: errors.New("404 not found"), expectErr: false},
		{name: "500 internal server error", err: errors.New("500 internal server error"), expectErr: false},
		{name: "hostname with 401", err: errors.New("http://backend401.example.com"), expectErr: false},
		// Pin against accidental loosening of the "authorization required"
		// substring matcher: a validation message of the form "field
		// 'authorization' required" must not be misclassified as an auth
		// failure. The current matcher uses the contiguous substring
		// "authorization required" (one space, no punctuation), so this
		// string does not match — but a future loosening (e.g. allowing
		// any whitespace) would silently regress.
		{name: "validation message containing 'authorization' and 'required'", err: errors.New("field 'authorization' required"), expectErr: false},
		{name: "nil error", err: nil, expectErr: false},
	}

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

			result := vmcp.IsAuthenticationError(tt.err)
			assert.Equal(t, tt.expectErr, result)
		})
	}
}

func TestIsTimeoutError(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		err       error
		expectErr bool
	}{
		{name: "timeout", err: errors.New("request timeout"), expectErr: true},
		{name: "deadline exceeded", err: errors.New("deadline exceeded"), expectErr: true},
		{name: "context deadline exceeded", err: errors.New("context deadline exceeded"), expectErr: true},
		{name: "Timeout (uppercase)", err: errors.New("Request Timeout"), expectErr: true},
		{name: "connection refused", err: errors.New("connection refused"), expectErr: false},
		{name: "generic error", err: errors.New("something went wrong"), expectErr: false},
		{name: "nil error", err: nil, expectErr: false},
	}

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

			result := vmcp.IsTimeoutError(tt.err)
			assert.Equal(t, tt.expectErr, result)
		})
	}
}

func TestIsConnectionError(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		err       error
		expectErr bool
	}{
		{name: "connection refused", err: errors.New("connection refused"), expectErr: true},
		{name: "connection reset", err: errors.New("connection reset by peer"), expectErr: true},
		{name: "no route to host", err: errors.New("no route to host"), expectErr: true},
		{name: "network unreachable", err: errors.New("network is unreachable"), expectErr: true},
		{name: "Connection Refused (uppercase)", err: errors.New("Connection Refused"), expectErr: true},
		{name: "timeout", err: errors.New("request timeout"), expectErr: false},
		{name: "authentication failed", err: errors.New("authentication failed"), expectErr: false},
		{name: "generic error", err: errors.New("something went wrong"), expectErr: false},
		{name: "nil error", err: nil, expectErr: false},
	}

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

			result := vmcp.IsConnectionError(tt.err)
			assert.Equal(t, tt.expectErr, result)
		})
	}
}

func TestHealthChecker_CheckHealth_Timeout(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockClient := mocks.NewMockBackendClient(ctrl)
	mockClient.EXPECT().
		ListCapabilities(gomock.Any(), gomock.Any()).
		DoAndReturn(func(ctx context.Context, _ *vmcp.BackendTarget) (*vmcp.CapabilityList, error) {
			// Simulate slow backend
			select {
			case <-time.After(2 * time.Second):
				return &vmcp.CapabilityList{}, nil
			case <-ctx.Done():
				return nil, ctx.Err()
			}
		}).
		Times(1)

	checker := NewHealthChecker(mockClient, 100*time.Millisecond, 0)
	target := &vmcp.BackendTarget{
		WorkloadID:   "backend-1",
		WorkloadName: "test-backend",
		BaseURL:      "http://localhost:8080",
	}

	status, err := checker.CheckHealth(context.Background(), target)
	assert.Error(t, err)
	assert.Equal(t, vmcp.BackendUnhealthy, status)
}

// TestHealthChecker_CheckHealth_ContextCarriesHealthCheckMarker verifies that CheckHealth
// passes a context with the health check marker to ListCapabilities.
// This is critical because the auth strategies (header_injection, token_exchange) read
// this marker to decide how to authenticate probe requests.
func TestHealthChecker_CheckHealth_ContextCarriesHealthCheckMarker(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	var capturedCtx context.Context
	mockClient := mocks.NewMockBackendClient(ctrl)
	mockClient.EXPECT().
		ListCapabilities(gomock.Any(), gomock.Any()).
		DoAndReturn(func(ctx context.Context, _ *vmcp.BackendTarget) (*vmcp.CapabilityList, error) {
			capturedCtx = ctx
			return &vmcp.CapabilityList{}, nil
		}).
		Times(1)

	checker := NewHealthChecker(mockClient, 5*time.Second, 0)
	target := &vmcp.BackendTarget{
		WorkloadID:   "backend-1",
		WorkloadName: "test-backend",
		BaseURL:      "http://localhost:8080",
	}

	status, err := checker.CheckHealth(context.Background(), target)
	require.NoError(t, err)
	assert.Equal(t, vmcp.BackendHealthy, status)

	// The context passed to ListCapabilities must carry the health check marker so
	// that auth strategies (header_injection, token_exchange) apply the correct
	// authentication path for probe requests.
	require.NotNil(t, capturedCtx, "context must have been captured")
	assert.True(t, IsHealthCheck(capturedCtx),
		"ListCapabilities must receive a context with the health check marker; "+
			"without it, header_injection and token_exchange strategies cannot "+
			"apply outgoing auth to health check probes")
}

// TestHealthChecker_CheckHealth_AuthErrorsCategorizesAsUnauthenticated verifies that
// auth errors from health checks are categorised as BackendUnauthenticated when the
// backend target has no outgoing auth strategy configured (AuthConfig nil in these
// cases). This represents a misconfiguration: the backend requires authentication
// but no strategy was configured on the target. A 401/403 from a backend that *does*
// have an outgoing auth strategy is treated as BackendHealthy by the checker and
// is covered in TestCategorizeError.
func TestHealthChecker_CheckHealth_AuthErrorsCategorizesAsUnauthenticated(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name string
		err  error
	}{
		{
			name: "header injection auth failure - http 401",
			err:  fmt.Errorf("transport error: http 401"),
		},
		{
			name: "token exchange auth failure - status code 401",
			err:  fmt.Errorf("backend unavailable: failed to initialize client for backend my-backend: status code 401"),
		},
		{
			name: "sentinel auth error",
			err:  vmcp.ErrAuthenticationFailed,
		},
		{
			name: "sentinel authz error",
			err:  vmcp.ErrAuthorizationFailed,
		},
		{
			name: "wrapped sentinel auth error",
			err:  fmt.Errorf("client credentials grant failed: %w", vmcp.ErrAuthenticationFailed),
		},
		{
			// transport.ErrUnauthorized is wrapped with ErrAuthenticationFailed in wrapBackendError,
			// so a 401 from the mcp-go transport layer reaches health monitoring as
			// BackendUnauthenticated instead of BackendUnhealthy.
			name: "mcp-go ErrUnauthorized wrapped as ErrAuthenticationFailed by wrapBackendError",
			err:  fmt.Errorf("%w: failed to initialize for backend my-backend: unauthorized (401)", vmcp.ErrAuthenticationFailed),
		},
		{
			// Issue #5223: mcp-go v0.49.0+ returns *AuthorizationRequiredError for
			// 401 with WWW-Authenticate, wrapped in *transport.Error. Both layers
			// Unwrap() to transport.ErrAuthorizationRequired. The string fallback
			// must recognize "authorization required" so the probe error reaches
			// health monitoring as BackendUnauthenticated rather than falling
			// through to BackendUnhealthy and producing WARN spam.
			name: "transport.Error wrapping AuthorizationRequiredError",
			err:  transport.NewError(&transport.AuthorizationRequiredError{}),
		},
	}

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

			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			mockClient := mocks.NewMockBackendClient(ctrl)
			mockClient.EXPECT().
				ListCapabilities(gomock.Any(), gomock.Any()).
				Return(nil, tt.err).
				Times(1)

			checker := NewHealthChecker(mockClient, 5*time.Second, 0)
			target := &vmcp.BackendTarget{
				WorkloadID:   "backend-1",
				WorkloadName: "test-backend",
				BaseURL:      "http://localhost:8080",
			}

			status, err := checker.CheckHealth(context.Background(), target)
			assert.Error(t, err)
			assert.Equal(t, vmcp.BackendUnauthenticated, status,
				"auth failure from a health probe should be BackendUnauthenticated, not BackendUnhealthy")
		})
	}
}

// TestHealthChecker_CheckHealth_AuthErrorWithOutgoingAuthIsHealthy verifies that a
// 401/403 from a backend that has an outgoing auth strategy configured (e.g.,
// upstream_inject, token_exchange, header_injection) is treated as BackendHealthy.
// Health probes deliberately do not carry user credentials, so the backend's auth
// challenge is the expected response and proves reachability. This is the behavior
// change introduced by the fix for issue #4920.
func TestHealthChecker_CheckHealth_AuthErrorWithOutgoingAuthIsHealthy(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name       string
		authConfig *authtypes.BackendAuthStrategy
		err        error
	}{
		{
			name:       "upstream_inject + sentinel auth error",
			authConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeUpstreamInject},
			err:        vmcp.ErrAuthenticationFailed,
		},
		{
			name:       "token_exchange + status code 401",
			authConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeTokenExchange},
			err:        fmt.Errorf("backend unavailable: failed to initialize client for backend my-backend: status code 401"),
		},
		{
			name:       "header_injection + HTTP 403",
			authConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeHeaderInjection},
			err:        errors.New("HTTP 403 forbidden"),
		},
		{
			name:       "upstream_inject + wrapped sentinel",
			authConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeUpstreamInject},
			err:        fmt.Errorf("%w: unauthorized (401)", vmcp.ErrAuthenticationFailed),
		},
		{
			// Issue #5223 exact reproducer: North's github-copilot-entry backend
			// behind an upstreamInject auth strategy. Probe carries no user creds;
			// mcp-go returns the typed authorization-required chain. Once correctly
			// classified as auth, #4935's logic must take over and report
			// BackendHealthy with nil err so the monitor records a successful
			// check, the circuit breaker stays closed, and the WARN spam stops.
			name:       "upstream_inject + transport.Error wrapping AuthorizationRequiredError",
			authConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeUpstreamInject},
			err:        transport.NewError(&transport.AuthorizationRequiredError{}),
		},
		{
			// Same mcp-go chain as the upstream_inject row above; this row pins
			// that token_exchange flows through the same authErrorStatus branch.
			name:       "token_exchange + transport.Error wrapping AuthorizationRequiredError",
			authConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeTokenExchange},
			err:        transport.NewError(&transport.AuthorizationRequiredError{}),
		},
		{
			// Same mcp-go chain as the upstream_inject row above; this row pins
			// that header_injection flows through the same authErrorStatus branch.
			name:       "header_injection + transport.Error wrapping AuthorizationRequiredError",
			authConfig: &authtypes.BackendAuthStrategy{Type: authtypes.StrategyTypeHeaderInjection},
			err:        transport.NewError(&transport.AuthorizationRequiredError{}),
		},
	}

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

			ctrl := gomock.NewController(t)
			t.Cleanup(ctrl.Finish)

			mockClient := mocks.NewMockBackendClient(ctrl)
			mockClient.EXPECT().
				ListCapabilities(gomock.Any(), gomock.Any()).
				Return(nil, tt.err).
				Times(1)

			checker := NewHealthChecker(mockClient, 5*time.Second, 0)
			target := &vmcp.BackendTarget{
				WorkloadID:   "backend-1",
				WorkloadName: "test-backend",
				BaseURL:      "http://localhost:8080",
				AuthConfig:   tt.authConfig,
			}

			status, err := checker.CheckHealth(t.Context(), target)
			// When the status is BackendHealthy (expected auth challenge) the
			// checker returns a nil error so the monitor records it as a
			// successful check and does not increment failure counters or open
			// the circuit breaker.
			assert.NoError(t, err,
				"auth challenge from an auth-configured backend must be reported "+
					"as a successful check (nil error) so the monitor records "+
					"success and the circuit breaker stays closed")
			assert.Equal(t, vmcp.BackendHealthy, status,
				"auth failure from a probe against a backend with an outgoing "+
					"auth strategy configured must be BackendHealthy — the challenge "+
					"is the expected response to a no-credential probe")
		})
	}
}

func TestHealthChecker_CheckHealth_MultipleBackends(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockClient := mocks.NewMockBackendClient(ctrl)

	// Setup different responses for different backends
	mockClient.EXPECT().
		ListCapabilities(gomock.Any(), gomock.Any()).
		DoAndReturn(func(_ context.Context, target *vmcp.BackendTarget) (*vmcp.CapabilityList, error) {
			switch target.WorkloadID {
			case "backend-healthy":
				return &vmcp.CapabilityList{}, nil
			case "backend-auth-error":
				return nil, errors.New("authentication failed")
			case "backend-timeout":
				return nil, errors.New("context deadline exceeded")
			default:
				return nil, errors.New("unknown error")
			}
		}).
		Times(4)

	checker := NewHealthChecker(mockClient, 5*time.Second, 0)

	// Test healthy backend
	status, err := checker.CheckHealth(context.Background(), &vmcp.BackendTarget{
		WorkloadID:   "backend-healthy",
		WorkloadName: "Healthy Backend",
		BaseURL:      "http://localhost:8080",
	})
	assert.NoError(t, err)
	assert.Equal(t, vmcp.BackendHealthy, status)

	// Test auth error backend
	status, err = checker.CheckHealth(context.Background(), &vmcp.BackendTarget{
		WorkloadID:   "backend-auth-error",
		WorkloadName: "Auth Error Backend",
		BaseURL:      "http://localhost:8081",
	})
	assert.Error(t, err)
	assert.Equal(t, vmcp.BackendUnauthenticated, status)

	// Test timeout backend
	status, err = checker.CheckHealth(context.Background(), &vmcp.BackendTarget{
		WorkloadID:   "backend-timeout",
		WorkloadName: "Timeout Backend",
		BaseURL:      "http://localhost:8082",
	})
	assert.Error(t, err)
	assert.Equal(t, vmcp.BackendUnhealthy, status)

	// Test unknown error backend
	status, err = checker.CheckHealth(context.Background(), &vmcp.BackendTarget{
		WorkloadID:   "backend-unknown",
		WorkloadName: "Unknown Backend",
		BaseURL:      "http://localhost:8083",
	})
	assert.Error(t, err)
	assert.Equal(t, vmcp.BackendUnhealthy, status)
}
