package server

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"strconv"
	"strings"
	"testing"
	"time"

	"buf.build/go/protovalidate"
	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab"
	gapi "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_gitlab"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"go.uber.org/mock/gomock"
	"golang.org/x/oauth2"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// Custom HTTP client implementation
type customRoundTripper struct {
	responses       map[string]func() *http.Response
	defaultResponse func(req *http.Request) (*http.Response, error)
}

func (m *customRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	key := req.Method + " " + req.URL.String()
	if val, exists := m.responses[key]; exists {
		resp := val()
		// Clone the response to avoid issues with multiple reads
		return &http.Response{
			StatusCode: resp.StatusCode,
			Header:     resp.Header.Clone(),
			Body:       resp.Body,
			Request:    req,
		}, nil
	}

	return m.defaultResponse(req)
}

// Test setup helper
func createTestAuthHandler(t *testing.T, gitLabClient *gitlab.Client, gitLabURL string) (*gomock.Controller, *authHandler, *MockAuthStore) {
	t.Helper()

	ctrl := gomock.NewController(t)

	if gitLabClient == nil {
		gitLabClient = mock_gitlab.SetupClient(t, "", "/", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusNotImplemented)
		})
	}

	if gitLabURL == "" {
		gitLabURL = "https://gdk.test"
	}

	oidcProviderCfg := &oidc.ProviderConfig{
		IssuerURL:   gitLabURL,
		AuthURL:     gitLabURL + "/oauth/authorize",
		TokenURL:    gitLabURL + "/oauth/token",
		UserInfoURL: gitLabURL + "/oauth/userinfo",
		JWKSURL:     gitLabURL + "/oauth/discovery/keys",
		Algorithms:  []string{"RS256"},
	}
	oidcProvider := oidcProviderCfg.NewProvider(t.Context())

	oauthConfig := &oauth2.Config{
		ClientID:     "test-client-id",
		ClientSecret: "test-client-secret",
		RedirectURL:  gitLabURL + "/-/kubernetes-agent/workspaces/oauth/redirect",
		Scopes:       []string{oidc.ScopeOpenID},
		Endpoint:     oidcProvider.Endpoint(),
	}

	store := NewMockAuthStore(ctrl)
	validator, err := protovalidate.New()
	require.NoError(t, err)

	handler := &authHandler{
		log:                   testlogger.New(t),
		gitLabClient:          gitLabClient,
		oidcProvider:          oidcProvider,
		oauthConfig:           oauthConfig,
		store:                 store,
		handleProcessingError: func(string, error) {},
		validator:             validator,
	}

	return ctrl, handler, store
}

func createMockOIDCHTTPClientContext(t *testing.T, tokenExchangeHTTPStatus, userInfoHTTPStatus int) context.Context {
	responses := make(map[string]func() *http.Response, 2)

	if tokenExchangeHTTPStatus != 0 {
		// Mock OAuth2 token exchange
		tokenExchangeBody := createJSONResponseBody(t, map[string]any{
			"access_token": "mock-access-token",
			"token_type":   "Bearer",
			"expires_in":   3600,
		})
		if tokenExchangeHTTPStatus != http.StatusOK {
			tokenExchangeBody = createJSONResponseBody(t, map[string]any{})
		}
		responses["POST https://gdk.test/oauth/token"] = func() *http.Response {
			return &http.Response{
				StatusCode: tokenExchangeHTTPStatus,
				Header:     http.Header{"Content-Type": []string{"application/json"}},
				Body:       tokenExchangeBody,
			}
		}
	}

	if userInfoHTTPStatus != 0 {
		// Mock OIDC UserInfo endpoint
		userInfoBody := createJSONResponseBody(t, map[string]any{
			"sub": "123",
		})
		if userInfoHTTPStatus != http.StatusOK {
			userInfoBody = createJSONResponseBody(t, map[string]any{})
		}
		responses["GET https://gdk.test/oauth/userinfo"] = func() *http.Response {
			return &http.Response{
				StatusCode: userInfoHTTPStatus,
				Header:     http.Header{"Content-Type": []string{"application/json"}},
				Body:       userInfoBody,
			}
		}
	}

	mockClient := &http.Client{
		Transport: &customRoundTripper{
			responses: responses,
			defaultResponse: func(req *http.Request) (*http.Response, error) {
				return nil, fmt.Errorf("unexpected request: %s %s", req.Method, req.URL.String())
			},
		},
	}
	return oidc.ClientContext(t.Context(), mockClient)
}

func createTestRequest(host, method, path string) *http.Request {
	req := httptest.NewRequest(method, path, nil)
	req.Host = host
	req.Header.Set(httpz.XForwardedProtoHeader, "https")
	return req
}

func createJSONResponseBody(t *testing.T, data any) io.ReadCloser {
	jsonData, err := json.Marshal(data)
	require.NoError(t, err)
	return io.NopCloser(strings.NewReader(string(jsonData)))
}

func mustMarshalProto(t *testing.T, msg proto.Message) []byte {
	data, err := proto.Marshal(msg)
	require.NoError(t, err)
	return data
}

func mustEncodeProtoState(t *testing.T, stateParam *OAuthStateParam) string {
	data, err := proto.Marshal(stateParam)
	require.NoError(t, err)
	return base64.URLEncoding.EncodeToString(data)
}

func assertErrorCookie(t *testing.T, cookie *http.Cookie, protoValidator protovalidate.Validator, expectedOriginalURL string, expectedHTTPStatus int, expectedMsg string) {
	require.NotNil(t, cookie)
	require.Equal(t, errorCookieName, cookie.Name)
	require.NotEmpty(t, cookie.Value)
	teData, err := base64.URLEncoding.DecodeString(cookie.Value)
	require.NoError(t, err)
	te := &TransferError{}
	err = proto.Unmarshal(teData, te)
	require.NoError(t, err)
	require.NotNil(t, te)
	err = protoValidator.Validate(te)
	require.NoError(t, err)
	assert.Equal(t, expectedOriginalURL, te.OriginalUrl)
	assert.Equal(t, expectedHTTPStatus, int(te.HttpStatusCode))
	assert.Equal(t, expectedMsg, te.Message)
}

func assertTransferErrorParam(t *testing.T, teDataEncoded string, protoValidator protovalidate.Validator, expectedOriginalURL string, expectedHTTPStatus int, expectedMsg string) {
	require.NotNil(t, teDataEncoded)
	teData, err := base64.URLEncoding.DecodeString(teDataEncoded)
	require.NoError(t, err)
	te := &TransferError{}
	err = proto.Unmarshal(teData, te)
	require.NoError(t, err)
	require.NotNil(t, te)
	err = protoValidator.Validate(te)
	require.NoError(t, err)
	assert.Equal(t, expectedOriginalURL, te.OriginalUrl)
	assert.Equal(t, expectedHTTPStatus, int(te.HttpStatusCode))
	assert.Equal(t, expectedMsg, te.Message)
}

// Tests for redirection to OAuth.

func TestAuthHandler_RedirectsToOAuth_Success(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest("5000-w1.workspaces.test", "GET", "/")
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_RedirectsToOAuth_SetInRedisError_ReturnsError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	// Mock failed OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(fmt.Errorf("redis error"))

	req := createTestRequest("5000-w1.workspaces.test", "GET", "/")
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusInternalServerError, w.Code)
	assert.Contains(t, w.Body.String(), "Something went wrong while redirecting to OAuth provider")
}

// Tests for error cookie.

func TestAuthHandler_Middleware_HandleErrorCookie_InvalidEncodedValue_ReturnsError(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  errorCookieName,
		Value: "Hello \xed\xa0\x80! world", // Invalid UTF-8 sequence
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Equal(t, "Something went wrong\n", w.Body.String())
	// Check cookie has been marked to be deleted
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assert.Equal(t, errorCookieName, cookies[0].Name)
	assert.Equal(t, -1, cookies[0].MaxAge)
}

func TestAuthHandler_Middleware_HandleErrorCookie_InvalidProtoValue_ReturnsError(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	teData := base64.StdEncoding.EncodeToString([]byte("invalid-proto-marshaled-data"))

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  errorCookieName,
		Value: teData,
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Equal(t, "Something went wrong\n", w.Body.String())
	// Check cookie has been marked to be deleted
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assert.Equal(t, errorCookieName, cookies[0].Name)
	assert.Equal(t, -1, cookies[0].MaxAge)
}

func TestAuthHandler_Middleware_HandleErrorCookie_InvalidHost_ReturnsError(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := "https://5000-w1.workspaces.test:9000/abc\n" // Invalid URL

	// Create a transfer error which is used to set the cookie value
	te := &TransferError{
		OriginalUrl:    originalURL,
		HttpStatusCode: http.StatusTeapot,
		Message:        "I'm a teapot",
	}
	rawData := mustMarshalProto(t, te)
	teData := base64.StdEncoding.EncodeToString(rawData)

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  errorCookieName,
		Value: teData,
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Equal(t, "Something went wrong\n", w.Body.String())
	// Check cookie has been marked to be deleted
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assert.Equal(t, errorCookieName, cookies[0].Name)
	assert.Equal(t, -1, cookies[0].MaxAge)
}

func TestAuthHandler_Middleware_HandleErrorCookie_HostMismatch_ReturnsError(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := "https://5000-w1.workspaces.test:9000/abc" // host mismatch

	// Create a transfer error which is used to set the cookie value
	te := &TransferError{
		OriginalUrl:    originalURL,
		HttpStatusCode: http.StatusTeapot,
		Message:        "I'm a teapot",
	}
	rawData := mustMarshalProto(t, te)
	teData := base64.StdEncoding.EncodeToString(rawData)

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  errorCookieName,
		Value: teData,
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Equal(t, "Something went wrong\n", w.Body.String())
	// Check cookie has been marked to be deleted
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assert.Equal(t, errorCookieName, cookies[0].Name)
	assert.Equal(t, -1, cookies[0].MaxAge)
}

func TestAuthHandler_Middleware_HandleErrorCookie_ValidValue_ReturnsErrorFromCookie(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/abc", host)

	// Create a transfer error which is used to set the cookie value
	te := &TransferError{
		OriginalUrl:    originalURL,
		HttpStatusCode: http.StatusTeapot,
		Message:        "I'm a teapot",
	}
	rawData := mustMarshalProto(t, te)
	teData := base64.StdEncoding.EncodeToString(rawData)

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  errorCookieName,
		Value: teData,
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTeapot, w.Code)
	assert.Contains(t, w.Body.String(), "I'm a teapot")
	// Check cookie has been marked to be deleted
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assert.Equal(t, errorCookieName, cookies[0].Name)
	assert.Equal(t, -1, cookies[0].MaxAge)
}

// Tests for user sessions.

func TestAuthHandler_Middleware_HandleUserSessionCookie_GetFromRedisError_ReturnsError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	sessionID := "session-id"

	// Mock session retrieval - returns error
	store.EXPECT().
		GetUserSession(gomock.Any(), sessionID).
		Return(nil, fmt.Errorf("redis error"))

	req := createTestRequest(host, "GET", "/workspace")
	req.AddCookie(&http.Cookie{
		Name:  userSessionCookieName,
		Value: sessionID,
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusInternalServerError, w.Code)
	assert.Contains(t, w.Body.String(), "Something went wrong\n")
}

func TestAuthHandler_Middleware_HandleUserSessionCookie_RedisNotFound_RedirectsToOAuth(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"

	// Create a valid user session
	sessionID := "test-session-id"

	// Mock session retrieval - not found
	store.EXPECT().
		GetUserSession(gomock.Any(), sessionID).
		Return(nil, nil)

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  userSessionCookieName,
		Value: sessionID,
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleUserSessionCookie_ExpiredSession_RedirectsToOAuth(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"

	// Create an expired user session
	userSession := &UserSession{
		Host:      host,
		User:      &User{Id: 123},
		CreatedAt: timestamppb.New(time.Now().Add(-25 * time.Hour)), // Expired
	}

	sessionID := "expired-session-id"

	// Mock session retrieval - found but expired
	store.EXPECT().
		GetUserSession(gomock.Any(), sessionID).
		Return(userSession, nil)

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  userSessionCookieName,
		Value: sessionID,
	})
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleUserSessionCookie_ValidSession_CallsNext(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"

	// Create a valid user session
	userSession := &UserSession{
		CreatedAt: timestamppb.Now(),
		Host:      host,
		User:      &User{Id: 123},
		Workspace: &Workspace{Id: 1, Port: 5000},
	}

	sessionID := "test-session-id"

	// Mock session retrieval - found and valid
	store.EXPECT().
		GetUserSession(gomock.Any(), sessionID).
		Return(userSession, nil)

	req := createTestRequest(host, "GET", "/")
	req.AddCookie(&http.Cookie{
		Name:  userSessionCookieName,
		Value: sessionID,
	})
	w := httptest.NewRecorder()

	nextCalled := false
	var contextUser *User
	var contextWorkspace *Workspace
	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		nextCalled = true
		user, err := getUserFromContext(r.Context())
		assert.NoError(t, err)
		contextUser = user
		workspace, err := getWorkspaceFromContext(r.Context())
		assert.NoError(t, err)
		contextWorkspace = workspace
	})

	handler.middleware(next).ServeHTTP(w, req)

	require.True(t, nextCalled, "Next handler should be called")
	require.NotNil(t, contextUser)
	assert.Equal(t, int64(123), contextUser.Id)
	require.NotNil(t, contextWorkspace)
	assert.Equal(t, int64(1), contextWorkspace.Id)
	assert.Equal(t, uint32(5000), contextWorkspace.Port)
}

// Tests for transfer error.

func TestAuthHandler_Middleware_HandleTransferError_InvalidEncodedValue_ReturnsWithBadRequestError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	teData := "VGhpcyBpcyBhIHRlc3QuCg!" // Contains an invalid character '!'

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/?"+transferErrorParamName+"="+teData)

	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleTransferError_InvalidProtoValue_ReturnsWithBadRequestError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	teData := base64.StdEncoding.EncodeToString([]byte("invalid-proto-marshaled-data"))

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/?"+transferErrorParamName+"="+teData)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleTransferError_HostMismatch_ReturnsWithBadRequestError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := "https://5000-w1.workspaces.test:9000/abc" // host mismatch

	te := &TransferError{
		OriginalUrl:    originalURL,
		HttpStatusCode: http.StatusTeapot,
		Message:        "I'm a teapot",
	}
	rawData := mustMarshalProto(t, te)
	teData := base64.StdEncoding.EncodeToString(rawData)

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/?"+transferErrorParamName+"="+teData)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleTransferError_RedirectsWithErrorCookie(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/abc", host)

	te := &TransferError{
		OriginalUrl:    originalURL,
		HttpStatusCode: http.StatusTeapot,
		Message:        "I'm a teapot",
	}
	rawData := mustMarshalProto(t, te)
	teData := base64.StdEncoding.EncodeToString(rawData)

	req := createTestRequest(host, "GET", "/?"+transferErrorParamName+"="+teData)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)

	// Check cookie was set
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assertErrorCookie(t, cookies[0], handler.validator, originalURL, http.StatusTeapot, "I'm a teapot")
}

// Tests for transfer tokens.

func TestAuthHandler_Middleware_HandleTransferToken_GetFromRedisError_RedirectsWithErrorCookie(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/", host)
	token := "token-id"

	// Mock token retrieval - returns error
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(nil, fmt.Errorf("redis error"))

	req := createTestRequest(host, "GET", "/?"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	// Check transfer token query param is removed
	assert.NotContains(t, w.Header().Get("Location"), transferTokenParamName)
	// Check cookie was set
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assertErrorCookie(t, cookies[0], handler.validator, originalURL, http.StatusInternalServerError, "Something went wrong")
}

func TestAuthHandler_Middleware_HandleTransferToken_NotFound_RedirectsToOAuth(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	path := "/test"
	token := "non-existent-token"

	// Mock transfer token retrieval - not found
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(nil, nil)

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", path+"?"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleTransferToken_InvalidOriginalURL_RedirectsWithErrorCookie(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	path := "/abc?xyz=123"
	originalURL := fmt.Sprintf("https://%s%s", host, path)

	// Create a transfer token with different host
	transferToken := &TransferToken{
		CreatedAt:   timestamppb.New(time.Now().Add(10 * time.Minute)),
		OriginalUrl: fmt.Sprintf("https://%s/%s\n", host, path), // Invalid URL
		User:        &User{Id: 123},
	}
	token := "host-mismatch-token"

	// Mock transfer token retrieval
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(transferToken, nil)

	req := createTestRequest(host, "GET", path+"&"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	// Check transfer token query param is removed
	assert.NotContains(t, w.Header().Get("Location"), transferTokenParamName)
	// Check cookie was set
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	// Since the original URL in the transfer token is not valid,
	// the expected original URL will be constructed from the request.
	assertErrorCookie(t, cookies[0], handler.validator, originalURL, http.StatusInternalServerError, "Something went wrong")
}

func TestAuthHandler_Middleware_HandleTransferToken_ExpiredToken_RedirectsToOAuth(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/abc?xyz=123", host)

	// Create an expired transfer token
	transferToken := &TransferToken{
		CreatedAt:   timestamppb.New(time.Now().Add(-2 * time.Minute)), // Expired
		OriginalUrl: originalURL,
		User:        &User{Id: 123},
	}

	token := "expired-transfer-token"

	// Mock transfer token retrieval
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(transferToken, nil)

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/?"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleTransferToken_HostMismatch_RedirectsToOAuth(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := "https://5000-w1.workspaces.test:9000/abc?xyz=123" // host mismatch

	// Create a transfer token with different host
	transferToken := &TransferToken{
		CreatedAt:   timestamppb.New(time.Now().Add(10 * time.Minute)),
		OriginalUrl: originalURL,
		User:        &User{Id: 123},
	}
	token := "host-mismatch-token"

	// Mock transfer token retrieval
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(transferToken, nil)

	// Mock successful OAuth state storage
	store.EXPECT().
		StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/?"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), "gdk.test/oauth/authorize")
}

func TestAuthHandler_Middleware_HandleTransferToken_DeleteFromRedisError_RedirectsWithErrorCookie(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	token := "token-id"
	originalURL := fmt.Sprintf("https://%s/abc?xyz=123", host)

	// Create a transfer token
	transferToken := &TransferToken{
		CreatedAt:   timestamppb.New(time.Now().Add(10 * time.Minute)),
		OriginalUrl: originalURL,
		User:        &User{Id: 123},
		Workspace:   &Workspace{Id: 1, Port: 5000},
	}

	// Mock transfer token retrieval
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(transferToken, nil)

	// Mock token deletion - returns error
	store.EXPECT().
		DeleteTransferToken(gomock.Any(), token).
		Return(fmt.Errorf("redis error"))

	req := createTestRequest(host, "GET", "/workspace?"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	// Check transfer token query param is removed
	assert.NotContains(t, w.Header().Get("Location"), transferTokenParamName)
	// Check error cookie was set
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assertErrorCookie(t, cookies[0], handler.validator, originalURL, http.StatusInternalServerError, "Something went wrong")
}

func TestAuthHandler_Middleware_HandleTransferToken_SetInRedis_RedirectsWithErrorCookie(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/abc?xyz=123", host)

	// Create a transfer token with different host
	transferToken := &TransferToken{
		CreatedAt:   timestamppb.New(time.Now().Add(10 * time.Minute)),
		OriginalUrl: originalURL,
		User:        &User{Id: 123},
		Workspace:   &Workspace{Id: 1, Port: 5000},
	}

	token := "token-id"

	// Mock transfer token retrieval
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(transferToken, nil)

	// Mock token deletion
	store.EXPECT().
		DeleteTransferToken(gomock.Any(), token).
		Return(nil)

	// Mock user sessions set - returns error
	store.EXPECT().
		StoreUserSession(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(fmt.Errorf("redis error"))

	req := createTestRequest(host, "GET", "/?"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	// Check transfer token query param is removed
	assert.NotContains(t, w.Header().Get("Location"), transferTokenParamName)
	// Check error cookie was set
	cookies := w.Result().Cookies()
	assertErrorCookie(t, cookies[0], handler.validator, originalURL, http.StatusInternalServerError, "Something went wrong")
}

func TestAuthHandler_Middleware_HandleTransferToken_ValidToken_CreatesUserSession(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/abc?xyz=123", host)
	user := &User{Id: 123}
	workspace := &Workspace{Id: 1, Port: 5000}

	// Create a valid transfer token
	transferToken := &TransferToken{
		CreatedAt:   timestamppb.Now(),
		OriginalUrl: originalURL,
		User:        user,
		Workspace:   workspace,
	}

	token := "valid-transfer-token"

	// Mock transfer token retrieval
	store.EXPECT().
		GetTransferToken(gomock.Any(), token).
		Return(transferToken, nil)

	// Mock token deletion
	store.EXPECT().
		DeleteTransferToken(gomock.Any(), token).
		Return(nil)

	// Mock session storage
	store.EXPECT().
		StoreUserSession(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	req := createTestRequest(host, "GET", "/?"+transferTokenParamName+"="+token)
	w := httptest.NewRecorder()

	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("Next handler should not be called")
	})

	handler.middleware(next).ServeHTTP(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)

	// Check cookie was set
	cookies := w.Result().Cookies()
	require.Len(t, cookies, 1)
	assert.Equal(t, userSessionCookieName, cookies[0].Name)
	assert.NotEmpty(t, cookies[0].Value)
}

// Tests for oauth redirect callback.

func TestAuthHandler_OAuthRedirectCallback_NonGetMethod_ReturnsMethodNotAllowed(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	req := createTestRequest("gdk.test", "POST", "/-/kubernetes-agent/workspaces/oauth/redirect")
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusMethodNotAllowed, w.Code)
}

func TestAuthHandler_OAuthRedirectCallback_MissingCode_ReturnsBadRequest(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?state=some-state")
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Contains(t, w.Body.String(), errTextMissingOrCodeState)
}

func TestAuthHandler_OAuthRedirectCallback_MissingState_ReturnsBadRequest(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code=test-code")
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Contains(t, w.Body.String(), errTextMissingOrCodeState)
}

func TestAuthHandler_OAuthRedirectCallback_InvalidState_ReturnsInternalServerError(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code=test-code&state=invalid-state")
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Contains(t, w.Body.String(), errTextInvalidOrMissingState)
}

func TestAuthHandler_OAuthRedirectCallback_InvalidState_InvalidOriginalURL_ReturnsBadRequest(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	// Create encoded state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: "abc.cyx\n", // Invalid URL
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code=test-code&state="+encodedState)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusInternalServerError, w.Code)
	assert.Contains(t, w.Body.String(), "Something went wrong")
}

func TestAuthHandler_OAuthRedirectCallback_GetFromRedisError_RedirectsWithInternalServerError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	originalURL := "https://5000-w1.workspaces.test/path"

	// Create valid encoded state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Mock OAuth state retrieval - returns error
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(nil, fmt.Errorf("redis error"))

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code=test-code&state="+encodedState)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusInternalServerError, errTextSomethingWentWrongOAuthCallback)
}

func TestAuthHandler_OAuthRedirectCallback_StateNotFoundInRedis_RedirectsWithBadRequestError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	originalURL := "https://5000-w1.workspaces.test/path"

	// Create valid encoded state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Mock OAuth state retrieval - not found
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(nil, nil)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code=test-code&state="+encodedState)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusBadRequest, errTextInvalidSession)
}

func TestAuthHandler_OAuthRedirectCallback_StateValueMismatch_RedirectsWithBadRequestError(t *testing.T) {
	ctrl, handler, store := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState + "random",
	}

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code=test-code&state="+encodedState)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusBadRequest, errTextInvalidSession)
}

func TestAuthHandler_OAuthRedirectCallback_UserTokenExchangeFails_RedirectsWithInternalServerError(t *testing.T) {
	gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
		resp := &gapi.AuthorizeUserWorkspaceAccessResponse{
			Status: gapi.AuthorizeUserWorkspaceAccessStatus_AUTHORIZED,
			Info: &gapi.AuthorizeUserWorkspaceAccessInfo{
				Id:   1,
				Port: 5000,
			},
		}
		testhelpers.RespondWithJSON(t, w, resp)
	})
	ctrl, handler, store := createTestAuthHandler(t, gitLabClient, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)
	code := "test-auth-code"

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState,
	}
	oidcCtx := createMockOIDCHTTPClientContext(t, http.StatusForbidden, 0)

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code="+code+"&state="+encodedState)
	req = req.WithContext(oidcCtx)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusInternalServerError, errTextSomethingWentWrongOAuthCallback)
}

func TestAuthHandler_OAuthRedirectCallback_UserInfoFails_RedirectsWithInternalServerError(t *testing.T) {
	gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
		resp := &gapi.AuthorizeUserWorkspaceAccessResponse{
			Status: gapi.AuthorizeUserWorkspaceAccessStatus_AUTHORIZED,
			Info: &gapi.AuthorizeUserWorkspaceAccessInfo{
				Id:   1,
				Port: 5000,
			},
		}
		testhelpers.RespondWithJSON(t, w, resp)
	})
	ctrl, handler, store := createTestAuthHandler(t, gitLabClient, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)
	code := "test-auth-code"

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState,
	}
	oidcCtx := createMockOIDCHTTPClientContext(t, http.StatusOK, http.StatusForbidden)

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code="+code+"&state="+encodedState)
	req = req.WithContext(oidcCtx)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusInternalServerError, errTextSomethingWentWrongOAuthCallback)
}

func TestAuthHandler_OAuthRedirectCallback_AuthorizationRequestFails_RedirectsWithInternalServerError(t *testing.T) {
	gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	})
	ctrl, handler, store := createTestAuthHandler(t, gitLabClient, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)
	code := "test-auth-code"

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState,
	}
	oidcCtx := createMockOIDCHTTPClientContext(t, http.StatusOK, http.StatusOK)

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code="+code+"&state="+encodedState)
	req = req.WithContext(oidcCtx)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusInternalServerError, errTextSomethingWentWrongOAuthCallback)
}

func TestAuthHandler_OAuthRedirectCallback_UserNotAuthorized_RedirectsWithNotFoundError(t *testing.T) {
	gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
		resp := &gapi.AuthorizeUserWorkspaceAccessResponse{
			Status: gapi.AuthorizeUserWorkspaceAccessStatus_NOT_AUTHORIZED,
			Info:   &gapi.AuthorizeUserWorkspaceAccessInfo{},
		}
		testhelpers.RespondWithJSON(t, w, resp)
	})
	ctrl, handler, store := createTestAuthHandler(t, gitLabClient, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)
	code := "test-auth-code"

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState,
	}
	oidcCtx := createMockOIDCHTTPClientContext(t, http.StatusOK, http.StatusOK)

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code="+code+"&state="+encodedState)
	req = req.WithContext(oidcCtx)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusNotFound, http.StatusText(http.StatusNotFound))
}

func TestAuthHandler_OAuthRedirectCallback_SetInRedisError_RedirectsWithInternalServerError(t *testing.T) {
	gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
		resp := &gapi.AuthorizeUserWorkspaceAccessResponse{
			Status: gapi.AuthorizeUserWorkspaceAccessStatus_AUTHORIZED,
			Info: &gapi.AuthorizeUserWorkspaceAccessInfo{
				Id:   1,
				Port: 5000,
			},
		}
		testhelpers.RespondWithJSON(t, w, resp)
	})
	ctrl, handler, store := createTestAuthHandler(t, gitLabClient, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)
	code := "test-auth-code"

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState,
	}
	oidcCtx := createMockOIDCHTTPClientContext(t, http.StatusOK, http.StatusOK)

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	// Mock transfer token storage - returns error
	store.EXPECT().
		StoreTransferToken(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(fmt.Errorf("redis error"))

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code="+code+"&state="+encodedState)
	req = req.WithContext(oidcCtx)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusInternalServerError, errTextSomethingWentWrongOAuthCallback)
}

func TestAuthHandler_OAuthRedirectCallback_DeleteFromRedisError_RedirectsWithInternalServerError(t *testing.T) {
	gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
		resp := &gapi.AuthorizeUserWorkspaceAccessResponse{
			Status: gapi.AuthorizeUserWorkspaceAccessStatus_AUTHORIZED,
			Info: &gapi.AuthorizeUserWorkspaceAccessInfo{
				Id:   1,
				Port: 5000,
			},
		}
		testhelpers.RespondWithJSON(t, w, resp)
	})
	ctrl, handler, store := createTestAuthHandler(t, gitLabClient, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)
	code := "test-auth-code"

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState,
	}
	oidcCtx := createMockOIDCHTTPClientContext(t, http.StatusOK, http.StatusOK)

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	// Mock transfer token storage
	store.EXPECT().
		StoreTransferToken(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	// Mock OAuth state cleanup
	store.EXPECT().
		DeleteOAuthState(gomock.Any(), encodedState).
		Return(fmt.Errorf("redis error"))

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code="+code+"&state="+encodedState)
	req = req.WithContext(oidcCtx)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	redirectURL, err := url.Parse(w.Header().Get("Location"))
	require.NoError(t, err)
	require.NotNil(t, redirectURL)
	teDataEncoded := redirectURL.Query().Get(transferErrorParamName)
	assertTransferErrorParam(t, teDataEncoded, handler.validator, originalURL, http.StatusInternalServerError, errTextSomethingWentWrongOAuthCallback)
}

func TestAuthHandler_OAuthRedirectCallback_UserAuthorized_CreatesTransferToken(t *testing.T) {
	gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
		resp := &gapi.AuthorizeUserWorkspaceAccessResponse{
			Status: gapi.AuthorizeUserWorkspaceAccessStatus_AUTHORIZED,
			Info: &gapi.AuthorizeUserWorkspaceAccessInfo{
				Id:   1,
				Port: 5000,
			},
		}
		testhelpers.RespondWithJSON(t, w, resp)
	})
	ctrl, handler, store := createTestAuthHandler(t, gitLabClient, "")
	defer ctrl.Finish()

	// Test data
	host := "5000-w1.workspaces.test"
	originalURL := fmt.Sprintf("https://%s/path", host)
	code := "test-auth-code"

	// Create state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL,
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.Now(),
	}
	encodedState := mustEncodeProtoState(t, stateParam)

	// Create OAuth state
	oauthState := &OAuthState{
		CreatedAt:    timestamppb.Now(),
		OriginalUrl:  originalURL,
		CodeVerifier: "test-verifier",
		StateValue:   encodedState,
	}
	oidcCtx := createMockOIDCHTTPClientContext(t, http.StatusOK, http.StatusOK)

	// Mock OAuth state retrieval
	store.EXPECT().
		GetOAuthState(gomock.Any(), encodedState).
		Return(oauthState, nil)

	// Mock transfer token storage
	store.EXPECT().
		StoreTransferToken(gomock.Any(), gomock.Any(), gomock.Any()).
		Return(nil)

	// Mock OAuth state cleanup
	store.EXPECT().
		DeleteOAuthState(gomock.Any(), encodedState).
		Return(nil)

	req := createTestRequest("gdk.test", "GET", "/-/kubernetes-agent/workspaces/oauth/redirect?code="+code+"&state="+encodedState)
	req = req.WithContext(oidcCtx)
	w := httptest.NewRecorder()

	handler.oauthRedirectCallback(w, req)

	assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
	assert.Contains(t, w.Header().Get("Location"), fmt.Sprintf("%s?%s=", originalURL, transferTokenParamName))
}

// Tests for auth handler's complete flow

func TestAuthHandler_CompleteFlow(t *testing.T) {
	user := &User{Id: 123}
	workspace := &Workspace{Id: 1, Port: 5000}
	tests := []struct {
		name           string
		originalURLStr string
	}{
		{
			name:           "URL with root path",
			originalURLStr: "https://5000-w1.workspaces.test/",
		},
		{
			name:           "URL with subpath",
			originalURLStr: "https://5000-w1.workspaces.test/workspace",
		},
		{
			name:           "URL with query params",
			originalURLStr: "https://5000-w1.workspaces.test/workspace?project=123&path=src/main.go&line=42",
		},
		{
			name:           "URL with encoded query params",
			originalURLStr: "https://5000-w1.workspaces.test/workspace/project=123&path=src%2Fmain.go&line=42",
		},
		{
			name:           "URL with query params contains transfer token",
			originalURLStr: "https://5000-w1.workspaces.test/workspace/project=123&path=src%2Fmain.go&line=42&" + transferTokenParamName + "=007",
		},
		{
			name:           "URL with query params contains transfer error",
			originalURLStr: "https://5000-w1.workspaces.test/workspace/project=123&path=src%2Fmain.go&line=42&" + transferErrorParamName + "=000",
		},
		{
			name:           "URL with query params contains transfer token and transfer error",
			originalURLStr: "https://5000-w1.workspaces.test/workspace/project=123&path=src%2Fmain.go&line=42&" + transferErrorParamName + "=000&" + transferTokenParamName + "=007",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			originalURL, err := url.Parse(tt.originalURLStr)
			require.NoError(t, err)

			// Create test OAuth server
			oauthServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				switch r.URL.Path {
				case "/oauth/authorize":
					// Simulate OAuth provider authorize endpoint
					// Extract state and redirect back with auth code
					state := r.URL.Query().Get("state")
					// We could technically use "gdk.test/-/kubernetes-agents/workspaces/...".
					// But the KAS endpoint can be configured in multiple ways.
					redirectURL := "https://kas.test/workspaces/oauth/redirect?code=test-auth-code&state=" + state
					http.Redirect(w, r, redirectURL, http.StatusTemporaryRedirect)
				case "/oauth/token":
					// Mock token exchange
					w.Header().Set("Content-Type", "application/json")
					err = json.NewEncoder(w).Encode(map[string]any{
						"access_token": "mock-access-token",
						"token_type":   "Bearer",
						"expires_in":   3600,
					})
					assert.NoError(t, err)
				case "/oauth/userinfo":
					// Mock user info
					w.Header().Set("Content-Type", "application/json")
					err = json.NewEncoder(w).Encode(map[string]any{
						"sub": strconv.FormatInt(user.Id, 10),
					})
					assert.NoError(t, err)
				}
			}))
			defer oauthServer.Close()
			oauthSrvURL, err := url.Parse(oauthServer.URL)
			require.NoError(t, err)

			gitLabClient := mock_gitlab.SetupClient(t, "", gapi.AuthorizeUserWorkspaceAccessAPIPath, func(w http.ResponseWriter, r *http.Request) {
				resp := &gapi.AuthorizeUserWorkspaceAccessResponse{
					Status: gapi.AuthorizeUserWorkspaceAccessStatus_AUTHORIZED,
					Info: &gapi.AuthorizeUserWorkspaceAccessInfo{
						Id:   workspace.Id,
						Port: workspace.Port,
					},
				}
				testhelpers.RespondWithJSON(t, w, resp)
			})
			ctrl, handler, store := createTestAuthHandler(t, gitLabClient, oauthSrvURL.String())
			defer ctrl.Finish()

			var nextHandlerCalled bool
			nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				nextHandlerCalled = true

				assert.Equal(t, tt.originalURLStr, buildOriginalURL(r).String(), "Original URL should be preserved")

				cookie, err := r.Cookie(userSessionCookieName) //nolint: govet
				assert.NoError(t, err)
				assert.NotNil(t, cookie)
				assert.NotEmpty(t, cookie.Value)

				ctxUser, err := getUserFromContext(r.Context())
				assert.NoError(t, err)
				assert.Equal(t, user, ctxUser)

				ctxWorkspace, err := getWorkspaceFromContext(r.Context())
				assert.NoError(t, err)
				assert.Equal(t, workspace, ctxWorkspace)

				w.WriteHeader(http.StatusOK)
				_, err = w.Write([]byte("Success"))
				assert.NoError(t, err)
			})

			// Create test Workspaces server and setup routes
			mux := http.NewServeMux()
			mux.HandleFunc("/", handler.middleware(nextHandler))
			mux.HandleFunc("kas.test/workspaces/oauth/redirect", func(w http.ResponseWriter, r *http.Request) {
				handler.oauthRedirectCallback(w, r)
			})
			workspacesSrv := httptest.NewServer(mux)
			defer workspacesSrv.Close()
			workspacesSrvURL, err := url.Parse(workspacesSrv.URL)
			require.NoError(t, err)

			// Set up mock store expectations
			store.EXPECT().
				StoreOAuthState(gomock.Any(), gomock.Any(), gomock.Any()).
				Return(nil)

			store.EXPECT().
				GetOAuthState(gomock.Any(), gomock.Any()).
				DoAndReturn(
					func(ctx context.Context, stateValue string) (*OAuthState, error) {
						return &OAuthState{
							CreatedAt:    timestamppb.Now(),
							OriginalUrl:  tt.originalURLStr,
							CodeVerifier: "test-verifier",
							StateValue:   stateValue,
						}, nil
					},
				)

			store.EXPECT().
				StoreTransferToken(gomock.Any(), gomock.Any(), gomock.Any()).
				Return(nil)

			store.EXPECT().
				DeleteOAuthState(gomock.Any(), gomock.Any()).
				Return(nil)

			store.EXPECT().
				GetTransferToken(gomock.Any(), gomock.Any()).
				Return(&TransferToken{
					CreatedAt:   timestamppb.Now(),
					OriginalUrl: tt.originalURLStr,
					User:        user,
					Workspace:   workspace,
				}, nil)

			store.EXPECT().
				DeleteTransferToken(gomock.Any(), gomock.Any()).
				Return(nil)

			store.EXPECT().
				StoreUserSession(gomock.Any(), gomock.Any(), gomock.Any()).
				Return(nil)

			store.EXPECT().
				GetUserSession(gomock.Any(), gomock.Any()).
				Return(&UserSession{
					CreatedAt: timestamppb.Now(),
					Host:      originalURL.Hostname(),
					User:      user,
					Workspace: workspace,
				}, nil).
				Times(2)

			// Create HTTP client that follows redirects and manages cookies
			jar, _ := cookiejar.New(nil)
			client := &http.Client{
				Jar: jar,
				CheckRedirect: func(req *http.Request, via []*http.Request) error {
					if len(via) >= 10 {
						return fmt.Errorf("too many redirects")
					}
					return nil
				},
				Transport: &customRoundTripper{
					responses: map[string]func() *http.Response{},
					defaultResponse: func(req *http.Request) (*http.Response, error) {
						// Create a copy of the request to avoid modifying the original request
						newReq := req.Clone(req.Context())

						// Set correct host header based on which server we are hitting
						switch newReq.URL.Host {
						case oauthSrvURL.Host:
							// OAuth server requests
							newReq.Host = "gdk.test"
						case "kas.test":
							// Workspaces server API requests - e.g. OAuth Redirect Callback requests
							newReq.Host = "kas.test"
							newReq.URL.Host = workspacesSrvURL.Host
							newReq.URL.Scheme = workspacesSrvURL.Scheme
						case originalURL.Hostname():
							// Workspaces server requests
							newReq.Host = originalURL.Hostname()
							newReq.URL.Host = workspacesSrvURL.Host
							newReq.URL.Scheme = workspacesSrvURL.Scheme
						default:
							return nil, fmt.Errorf("unknown host: %s", newReq.URL.Host)
						}

						newReq.Header.Set(httpz.XForwardedProtoHeader, originalURL.Scheme)

						return http.DefaultTransport.RoundTrip(newReq)
					},
				},
			}

			// Single request that follows all redirects automatically
			resp1, err := client.Get(originalURL.String()) //nolint:noctx
			require.NoError(t, err)
			defer resp1.Body.Close()

			// Verify next handler was reached with correct data
			assert.True(t, nextHandlerCalled, "Final handler should be reached")
			assert.Equal(t, http.StatusOK, resp1.StatusCode)

			body1, err := io.ReadAll(resp1.Body)
			require.NoError(t, err)
			assert.Equal(t, "Success", string(body1))

			// Reset next handler fixture
			nextHandlerCalled = false

			// Request with the existing session set already in the above steps
			resp2, err := client.Get(originalURL.String()) //nolint:noctx
			require.NoError(t, err)
			defer resp2.Body.Close()

			// Verify next handler was reached with correct data
			assert.True(t, nextHandlerCalled, "Next handler should be reached")
			assert.Equal(t, http.StatusOK, resp2.StatusCode)

			body2, err := io.ReadAll(resp2.Body)
			require.NoError(t, err)
			assert.Equal(t, "Success", string(body2))
		})
	}
}

// Tests for encoding and decoding oauth state.

func TestAuthHandler_StateEncodeDecodeRoundTrip(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	originalURL, _ := url.Parse("https://5000-w1.workspaces.test/path?param=value")

	// Test encoding and decoding
	encoded, err := handler.encodeOAuthState(originalURL)
	require.NoError(t, err)
	assert.NotEmpty(t, encoded)

	decoded, err := handler.decodeOAuthState(encoded)
	require.NoError(t, err)
	assert.Equal(t, originalURL.String(), decoded.OriginalUrl)
	assert.NotEmpty(t, decoded.Nonce)
}

func TestAuthHandler_InvalidJSONDecoding_ReturnsError(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	_, err := handler.decodeOAuthState("invalid-json-encoded")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "invalid state format")
}

func TestAuthHandler_ExpiredStateDecoding_ReturnsError(t *testing.T) {
	ctrl, handler, _ := createTestAuthHandler(t, nil, "")
	defer ctrl.Finish()

	// Create an old state parameter
	stateParam := &OAuthStateParam{
		OriginalUrl: "https:/5000-w1.workspaces.test/path?param=value",
		Nonce:       "test-nonce",
		CreatedAt:   timestamppb.New(time.Now().Add(-10 * time.Minute)), // Expired
	}

	encoded := mustEncodeProtoState(t, stateParam)

	_, err := handler.decodeOAuthState(encoded)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "state expired")
}

// Tests for url appending and removing query params.

func TestAuthHandler_URLManipulation(t *testing.T) {
	transferToken := "abc"
	tests := []struct {
		name                         string
		originalURLString            string
		expectedURLStringAfterAppend string
		expectedURLStringAfterRemove string
	}{
		{
			name:                         "URL contains no query params",
			originalURLString:            "https://5000-w1.workspaces.test/path",
			expectedURLStringAfterAppend: fmt.Sprintf("https://5000-w1.workspaces.test/path?%s=abc", transferTokenParamName),
			expectedURLStringAfterRemove: "https://5000-w1.workspaces.test/path",
		},
		{
			name:                         "URL contains query param",
			originalURLString:            "https://5000-w1.workspaces.test/path?existing=param",
			expectedURLStringAfterAppend: fmt.Sprintf("https://5000-w1.workspaces.test/path?%s=abc&existing=param", transferTokenParamName), // The order of the query params are changed
			expectedURLStringAfterRemove: "https://5000-w1.workspaces.test/path?existing=param",
		},
		{
			name:                         "URL contains query param which are empty",
			originalURLString:            "https://5000-w1.workspaces.test/path?existing=param&empty=",
			expectedURLStringAfterAppend: fmt.Sprintf("https://5000-w1.workspaces.test/path?%s=abc&empty=&existing=param", transferTokenParamName), // The order of the query params are changed
			expectedURLStringAfterRemove: "https://5000-w1.workspaces.test/path?empty=&existing=param",                                             // The order of the query params are changed
		},
		{
			name:                         "URL contains query param along with transfer token query param",
			originalURLString:            fmt.Sprintf("https://5000-w1.workspaces.test/path?existing=param&%s=xyz", transferTokenParamName),
			expectedURLStringAfterAppend: fmt.Sprintf("https://5000-w1.workspaces.test/path?%s=xyz&%s=abc&existing=param", transferTokenParamName, transferTokenParamName), // The order of the query params are changed
			expectedURLStringAfterRemove: fmt.Sprintf("https://5000-w1.workspaces.test/path?%s=xyz&existing=param", transferTokenParamName),                                // The order of the query params are changed
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			originalURL, err := url.Parse(tt.originalURLString)
			require.NoError(t, err)

			expectedURLAfterAppend, err := url.Parse(tt.expectedURLStringAfterAppend)
			require.NoError(t, err)

			expectedURLAfterRemove, err := url.Parse(tt.expectedURLStringAfterRemove)
			require.NoError(t, err)

			// Test appending value to query param
			appendValueToParam(originalURL, transferTokenParamName, transferToken)
			assert.Equal(t, expectedURLAfterAppend, originalURL)

			// Test removing value from query param
			// NOTE: originalURL has been modified in the code above
			removeLastValueForParam(originalURL, transferTokenParamName)
			assert.Equal(t, expectedURLAfterRemove, originalURL)
		})
	}
}

func TestAuthHandler_URLManipulation_RemoveLastValueForParam_EmptyRawQuery(t *testing.T) {
	originalURL, err := url.Parse("https://5000-w1.workspaces.test/path")
	require.NoError(t, err)
	removeLastValueForParam(originalURL, transferTokenParamName)
	assert.Equal(t, "https://5000-w1.workspaces.test/path", originalURL.String())
}

func Test_buildOriginalURL(t *testing.T) {
	req := createTestRequest("5000-w1.workspaces.test", "GET", "/")
	actualURL := buildOriginalURL(req)
	assert.Equal(t, "5000-w1.workspaces.test", actualURL.Host)
	assert.Equal(t, "/", actualURL.Path)
	// Verify the original request's URL is not modified
	assert.Empty(t, req.URL.Scheme)
}

func Test_getRequestScheme(t *testing.T) {
	tests := []struct {
		name      string
		headerSet bool
		tlsSet    bool
		expected  string
	}{
		{
			name:      "When request arrives with TLS and is proxied with X-Forwarded-Proto header set",
			headerSet: true,
			tlsSet:    false,
			expected:  "https",
		},
		{
			name:      "When request arrives directly with TLS",
			headerSet: false,
			tlsSet:    true,
			expected:  "https",
		},
		{
			name:      "When request arrives directly without TLS",
			headerSet: false,
			tlsSet:    false,
			expected:  "http",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, "/", nil)
			if tt.headerSet {
				req.Header.Add(httpz.XForwardedProtoHeader, "https")
			}
			if tt.tlsSet {
				req.TLS = &tls.ConnectionState{}
			}
			actual := getRequestScheme(req)
			assert.Equal(t, tt.expected, actual)
		})
	}
}
