package server

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"buf.build/go/protovalidate"
	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/redis/rueidis"
	"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/logz"
	"golang.org/x/oauth2"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type (
	userCtxKey      struct{}
	workspaceCtxKey struct{}
)

const (
	// errorCookieName is the cookie set if there are any errors during authentication and authorization.
	// This is to render the errors on the user's requested workspace host rather than the server's api host where
	// the error might occur. This cookie is deleted as soon as the error is rendered.
	errorCookieName = "gitlab-workspace-error"
	// userSessionCookieName is the cookie set after successful authentication and authorization.
	userSessionCookieName = "gitlab-workspace-session"
	// transferErrorParamName is the url query param used to transfer errors from server's api host, where an error
	// occurred, to the user's requested url. This value is used to set the error cookie on the user's requested url host.
	transferErrorParamName = "__gitlab_workspace_transfer_error"
	// transferTokenParamName is the url query param used to transfer token from server's api host to the user's
	// requested url. This is required to set the user session cookie on the user's requested url host, after successful
	// authentication and authorization.
	transferTokenParamName = "__gitlab_workspace_transfer_token" //nolint: gosec

	// errorCookieTTL is the time duration for which is the error cookie is valid.
	errorCookieTTL = 1 * time.Minute
	// oauthStateParamTTL is the time duration for which the oauth state param is valid.
	oauthStateParamTTL = 5 * time.Minute
	// oauthStateTTL is the time duration for which the oauth state is valid. It is also used for redis cache ttl.
	oauthStateTTL = 5 * time.Minute
	// transferTokenTTL is the time duration for which the transfer token is valid. It is also used for redis cache ttl.
	transferTokenTTL = 1 * time.Minute
	// userSessionTTL is the time duration for which the user session is valid. It is also used for redis cache ttl.
	userSessionTTL = 24 * time.Hour

	// OAuth callback query params
	oauthCodeParamName  = "code"
	oauthStateParamName = "state"

	// Error texts for responses
	errTextSomethingWentWrong              = "Something went wrong"
	errTextSomethingWentWrongOAuthCallback = "Something went wrong while handling OAuth redirect callback"
	errTextMissingOrCodeState              = "Missing code or state"
	errTextInvalidOrMissingState           = "Invalid or expired state"
	errTextInvalidSession                  = "Invalid or expired session"
)

type authHandlerInterface interface {
	middleware(next http.HandlerFunc) http.HandlerFunc
	oauthRedirectCallback(w http.ResponseWriter, r *http.Request)
}

type authHandler struct {
	log                   *slog.Logger
	gitLabClient          gitlab.ClientInterface
	oidcProvider          *oidc.Provider
	oauthConfig           *oauth2.Config
	store                 AuthStore
	handleProcessingError func(string, error)
	validator             protovalidate.Validator
}

func newAuthHandler(log *slog.Logger, gitLabClient gitlab.ClientInterface, oidcProvider *oidc.Provider, oauthConfig *oauth2.Config, redisClient rueidis.Client, redisPrefix string, handleProcessingError func(string, error), validator protovalidate.Validator) *authHandler {
	store := NewRedisAuthStore(redisClient, redisPrefix, oauthStateTTL, userSessionTTL, transferTokenTTL)

	return &authHandler{
		log:                   log,
		gitLabClient:          gitLabClient,
		oidcProvider:          oidcProvider,
		oauthConfig:           oauthConfig,
		store:                 store,
		handleProcessingError: handleProcessingError,
		validator:             validator,
	}
}

// middleware handles OAuth authentication for traffic on user's workspace host.
//
// User's workspace host - `port-ws1.workspaces.example.com`
// OAuth callback URL - `kas.example.com/workspaces/oauth/redirect`
// Since the two are different, we need to perform the following flow to set the cookie on the user's workspace host.
//
//  1. User accesses `port-ws1.workspaces.example.com`.
//  2. Auth middleware detects there is no valid user session cookie or error cookie or a transfer token or transfer error
//     in the query param. It stores state data in oauthStates and redirects to the OAuth provider - `gitlab.example.com/oauth/authorize`.
//  3. OAuth provider performs authentication and redirects back to `kas.example.com/workspaces/oauth/redirect`.
//  4. OAuth redirect callback handler validates the state received in the request with the data stored in oauthStates,
//     exchanges the OAuth code for OAuth token, gets the authenticated user from OAuth token,
//     checks if user is authorized to access the workspace, stores the transfer token in transferTokens cache and
//     then redirects back to the original URL with an additional query param transferTokenParamName
//     which contains the transfer token id. If there are any errors, it redirects back to the original URL with an
//     additional query param transferErrorParamName which contains the error to be rendered.
//  5. Auth middleware detects there is no valid error cookie or user session cookie. If there is a transfer error in the
//     query param, it sets the error cookie for the domain `ws1.workspaces.example.com`. It then redirects back
//     to the original URL without the additional query param transferErrorParamName. Else, if there is a transfer token
//     in the query param transferTokenParamName, uses that value to look up in the transferTokens
//     cache, delete it from the cache since it is meant to be single use, checks if the transfer token is valid,
//     checks if the transfer token's host and the request's host match, generate a new user session and store it in
//     the userSessions cache and set the user session cookie for the domain `ws1.workspaces.example.com`. If there are
//     any errors while handling of transfer token, it set the error cookie for the domain `ws1.workspaces.example.com`.
//     It then redirects back to the original URL without the additional query param transferTokenParamName.
//  6. Auth middleware checks if the error cookie is set, returns the error from cookie value and deletes the cookie.
//     Else, if user session cookie is set, extracts the user session id from cookie value,
//     validates the session in userSessions cache, inject the user and workspace details in the request's context and
//     serves the next HTTP request handler.
//
// In the above flow, when the user session cookie is being set, the origin of the redirect is the
// OAuth callback URL - `kas.example.com/workspaces/oauth/redirect`. This is the reason we have to use `Lax` as the
// SameSite attribute for the cookie. If we want to use `Strict` attribute, we will need to use a two-cookie strategy.
// OAuth callback URL -> User's workspace host sets cookie1 with Lax -> User's workspace host sets cookie1 with Strict.
// This addition to the flow would further complicate it without much value. Discussions on this matter can be viewed
// at - https://gitlab.com/gitlab-org/cluster-integration/gitlab-agent/-/merge_requests/2654#note_2632971823 .
func (h *authHandler) middleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		cookies := r.Cookies()

		// 1. Check error cookie.
		if errCookie := getCookie(cookies, errorCookieName); errCookie != nil && errCookie.Value != "" {
			h.handleErrorCookie(w, r, errCookie.Value)
			return
		}

		// 2. Check user session cookie.
		if userSessionCookie := getCookie(cookies, userSessionCookieName); userSessionCookie != nil && userSessionCookie.Value != "" {
			h.handleUserSessionCookie(w, r, next, userSessionCookie.Value)
			return
		}

		query := r.URL.Query()

		// 3. Check transfer token.
		// Transfer token is checked before transfer error to avoid OAuth loop for situations where the user's
		// original URL contains the transfer error in the param.
		if transferTokens, ok := query[transferTokenParamName]; ok && len(transferTokens) > 0 {
			h.handleTransferToken(w, r, transferTokens)
			return
		}

		// 4. Check transfer error.
		if transferErrs, ok := query[transferErrorParamName]; ok && len(transferErrs) > 0 {
			h.handleTransferError(w, r, transferErrs)
			return
		}

		// 5. Redirect to OAuth
		h.redirectToOAuth(w, r)
	}
}

// handleErrorCookie processes error cookie, removes the cookie and serves the error to the user.
// We do not check for the validity of the cookie since we are only using it for rendering an error and
// deleting it immediately.
func (h *authHandler) handleErrorCookie(w http.ResponseWriter, r *http.Request, value string) {
	// Invalidate the cookie immediately to prevent loop and weird issues.
	setCookie(w, r, errorCookieName, "", -1*time.Second)
	teData, err := base64.URLEncoding.DecodeString(value)
	if err != nil {
		h.log.Debug("Failed to decode transfer error protobuf", logz.Error(err))
		http.Error(w, errTextSomethingWentWrong, http.StatusBadRequest)
		return
	}
	var te TransferError
	if err = proto.Unmarshal(teData, &te); err != nil {
		h.log.Debug("Failed to unmarshal transfer error protobuf", logz.Error(err))
		http.Error(w, errTextSomethingWentWrong, http.StatusBadRequest)
		return
	}
	if err = h.validator.Validate(&te); err != nil {
		h.log.Debug("Failed to validate transfer error protobuf", logz.Error(err))
		http.Error(w, errTextSomethingWentWrong, http.StatusBadRequest)
		return
	}
	valid, err := te.IsValid(r.Host)
	if err != nil {
		h.log.Debug("Failed to validate transfer error protobuf", logz.Error(err))
		http.Error(w, errTextSomethingWentWrong, http.StatusBadRequest)
		return
	}
	if !valid {
		h.log.Debug("Failed to validate transfer error cookie value", logz.Error(err))
		http.Error(w, errTextSomethingWentWrong, http.StatusBadRequest)
		return
	}
	http.Error(w, te.Message, int(te.HttpStatusCode))
}

// handleUserSessionCookie processes user session cookie and serves the next http handler if valid.
// If user session is not valid, redirect back to OAuth.
func (h *authHandler) handleUserSessionCookie(w http.ResponseWriter, r *http.Request, next http.HandlerFunc, value string) {
	us, err := h.store.GetUserSession(r.Context(), value)
	if err != nil {
		h.log.Error("Failed to read user session from redis", logz.Error(err))
		http.Error(w, errTextSomethingWentWrong, http.StatusInternalServerError)
		return
	}
	if us == nil || !us.IsValid(userSessionTTL, r.Host) {
		h.redirectToOAuth(w, r)
		return
	}
	authCtx := context.WithValue(r.Context(), userCtxKey{}, us.User)
	authCtx = context.WithValue(authCtx, workspaceCtxKey{}, us.Workspace)
	next.ServeHTTP(w, r.WithContext(authCtx))
}

// handleTransferError processes last transfer error and sets error cookie and redirects back to the original URL.
func (h *authHandler) handleTransferError(w http.ResponseWriter, r *http.Request, transferErrs []string) {
	transferErr := transferErrs[len(transferErrs)-1]
	teData, err := base64.URLEncoding.DecodeString(transferErr)
	if err != nil {
		// This is a coding-mistake or there is a transfer error query param in the user's original request which is not
		// properly encoded. Hence, we redirect to OAuth without modifying the request.
		h.log.Debug("Failed to decode transfer error protobuf", logz.Error(err))
		h.redirectToOAuth(w, r)
		return
	}
	var te TransferError
	err = proto.Unmarshal(teData, &te)
	if err != nil {
		// This is a coding-mistake or there is a transfer error query param in the user's original request.
		// Hence, we redirect to OAuth without modifying the request.
		h.log.Debug("Failed to unmarshal transfer error protobuf", logz.Error(err))
		h.redirectToOAuth(w, r)
		return
	}
	if err = h.validator.Validate(&te); err != nil {
		// This is a coding-mistake or there is a transfer error query param in the user's original request.
		// Hence, we redirect to OAuth without modifying the request.
		h.log.Debug("Failed to validate transfer error protobuf", logz.Error(err))
		h.redirectToOAuth(w, r)
		return
	}
	if valid, err := te.IsValid(r.Host); !valid || err != nil {
		// There is an expired transfer error query param in the user's original request.
		// Hence, we redirect to OAuth without modifying the request.
		h.log.Debug("Failed to validate transfer error protobuf", logz.Error(err))
		h.redirectToOAuth(w, r)
		return
	}
	setCookie(w, r, errorCookieName, transferErr, errorCookieTTL)
	http.Redirect(w, r, te.OriginalUrl, http.StatusTemporaryRedirect)
}

// handleTransferToken processes last transfer token and sets user session cookie and redirects back to the original URL.
func (h *authHandler) handleTransferToken(w http.ResponseWriter, r *http.Request, transferTokens []string) {
	transferToken := transferTokens[len(transferTokens)-1]
	ctx := r.Context()
	host := r.Host

	// Get and validate transfer token
	tt, err := h.store.GetTransferToken(ctx, transferToken)
	if err != nil {
		h.log.Error("Failed to read transfer token from redis", logz.Error(err))
		// Since the transfer token is not valid, we cannot extract the original URL.
		// Hence, we try to construct it from the request.
		originalURL := buildOriginalURL(r)
		// There is a chance that if the user's original request contained the transfer token param, then it will be
		// removed. Given that this will only happen if they are doing an OAuth and Redis is throwing errors, this is
		// an acceptable compromise.
		removeLastValueForParam(originalURL, transferTokenParamName)
		h.handleErrorOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrong, http.StatusInternalServerError)
		return
	}
	if tt == nil {
		// This is either an expired(in redis) or the user's original request actually contains the transfer token.
		// So we must preserve it as is and perform a redirect.
		h.redirectToOAuth(w, r)
		return
	}

	originalURL, err := url.Parse(tt.OriginalUrl)
	if err != nil {
		// Since the original URL parsed is not valid, we try to construct it from the request.
		originalURL = buildOriginalURL(r)
		// There is a chance that if the user's original request contained the transfer token param, then it will be
		// removed. Given that this will only happen if they are doing an OAuth and Redis is throwing errors, this is
		// an acceptable compromise.
		removeLastValueForParam(originalURL, transferTokenParamName)
		h.log.Error("Failed to read transfer token from redis", logz.Error(err))
		h.handleErrorOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrong, http.StatusInternalServerError)
		return
	}
	if valid, err := tt.IsValid(transferTokenTTL, host); !valid || err != nil { //nolint:govet
		// Transfer token not found or expired or host mismatch
		r.URL = originalURL
		h.redirectToOAuth(w, r)
		return
	}

	// Delete transfer token (single use)
	if err = h.store.DeleteTransferToken(ctx, transferToken); err != nil {
		h.log.Error("Failed to delete transfer token from redis", logz.Error(err))
		// The end user's request is technically not affected by this, and we could continue processing the request.
		// However, since the transfer token is actually used to set the user session cookie, out of abundance of
		// caution, we want to abort the request and return an error to prevent any replay attack scenarios.
		h.handleErrorOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrong, http.StatusInternalServerError)
		return
	}

	// Generate new session ID and store user session
	sessionID := generateRandomToken()

	us := &UserSession{
		Host:      host,
		User:      tt.User,
		CreatedAt: timestamppb.Now(),
		Workspace: tt.Workspace,
	}

	if err = h.store.StoreUserSession(ctx, sessionID, us); err != nil {
		h.log.Error("Failed to store user session", logz.Error(err))
		h.handleErrorOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrong, http.StatusInternalServerError)
		return
	}

	// Set user session cookie and redirect
	setCookie(w, r, userSessionCookieName, sessionID, userSessionTTL)
	http.Redirect(w, r, originalURL.String(), http.StatusTemporaryRedirect)
}

// redirectToOAuth initiates the OAuth flow.
func (h *authHandler) redirectToOAuth(w http.ResponseWriter, r *http.Request) {
	// Create and encode OAuth state
	originalURL := buildOriginalURL(r)
	stateValue, err := h.encodeOAuthState(originalURL)
	if err != nil {
		h.handleProcessingError("Failed to encode OAuth state", err)
		http.Error(w, "Something went wrong while redirecting to OAuth provider", http.StatusInternalServerError)
		return
	}

	// Store OAuth state
	codeVerifier := oauth2.GenerateVerifier()
	oauthState := &OAuthState{
		OriginalUrl:  originalURL.String(),
		CodeVerifier: codeVerifier,
		StateValue:   stateValue,
		CreatedAt:    timestamppb.Now(),
	}
	if err = h.store.StoreOAuthState(r.Context(), stateValue, oauthState); err != nil {
		h.log.Error("Failed to store oauth state", logz.Error(err))
		http.Error(w, "Something went wrong while redirecting to OAuth provider", http.StatusInternalServerError)
		return
	}

	// Build and redirect to OAuth URL with PKCE parameters
	authURL := h.oauthConfig.AuthCodeURL(stateValue, oauth2.S256ChallengeOption(codeVerifier))
	http.Redirect(w, r, authURL, http.StatusTemporaryRedirect)
}

// oauthRedirectCallback handles OAuth callback after authentication.
func (h *authHandler) oauthRedirectCallback(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		// This is either a coding mistake or someone probing the server.
		// So we should return immediately and not try to redirect back to the user's workspace URL.
		http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
		return
	}

	query := r.URL.Query()
	code := query.Get(oauthCodeParamName)
	stateValue := query.Get(oauthStateParamName)

	if code == "" || stateValue == "" {
		// This is either a coding mistake or someone trying a replay attack.
		// So we should return immediately and not try to redirect back to the user's workspace URL.
		http.Error(w, errTextMissingOrCodeState, http.StatusBadRequest)
		return
	}

	// Validate OAuth state
	decodedState, err := h.decodeOAuthState(stateValue)
	if err != nil {
		// This is either a coding mistake or someone trying a replay attack.
		// So we should return immediately and not try to redirect back to the user's workspace URL.
		http.Error(w, errTextInvalidOrMissingState, http.StatusBadRequest)
		return
	}

	originalURL, err := url.Parse(decodedState.OriginalUrl)
	if err != nil {
		// This is either a coding mistake or someone trying a replay attack.
		// So we should return immediately and not try to redirect back to the user's workspace URL.
		h.log.Error("Failed to parse url from decoded state", logz.Error(err))
		http.Error(w, errTextSomethingWentWrong, http.StatusInternalServerError)
		return
	}

	ctx := r.Context()
	oauthState, err := h.store.GetOAuthState(ctx, stateValue)
	if err != nil {
		h.log.Error("Failed to read oauth state from redis", logz.Error(err))
		h.handleErrorNotOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrongOAuthCallback, http.StatusInternalServerError)
		return
	}
	if oauthState == nil || !oauthState.IsValid(oauthStateTTL, stateValue) {
		h.handleErrorNotOnWorkspaceHost(w, r, originalURL, errTextInvalidSession, http.StatusBadRequest)
		return
	}

	// Exchange code for user
	u, err := h.exchangeCodeForUser(ctx, code, oauthState.CodeVerifier)
	if err != nil {
		h.handleErrorNotOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrongOAuthCallback, http.StatusInternalServerError)
		return
	}

	// Check if user is authorized to access this workspace
	authzResp, err := gapi.AuthorizeUserWorkspaceAccess(ctx, h.gitLabClient, originalURL.Host, u.Id)
	if err != nil {
		h.log.Error("Failed to authorize user's workspace access", logz.Error(err))
		h.handleErrorNotOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrongOAuthCallback, http.StatusInternalServerError)
		return
	}
	if authzResp.Status != gapi.AuthorizeUserWorkspaceAccessStatus_AUTHORIZED {
		h.handleErrorNotOnWorkspaceHost(w, r, originalURL, http.StatusText(http.StatusNotFound), http.StatusNotFound)
		return
	}

	// Generate and store transfer token
	token := generateRandomToken()

	transferTokenObj := &TransferToken{
		User:        u,
		OriginalUrl: decodedState.OriginalUrl,
		CreatedAt:   timestamppb.Now(),
		Workspace:   &Workspace{Id: authzResp.Info.Id, Port: authzResp.Info.Port},
	}

	if err = h.store.StoreTransferToken(ctx, token, transferTokenObj); err != nil {
		h.log.Error("Failed to store transfer token in redis", logz.Error(err))
		h.handleErrorNotOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrongOAuthCallback, http.StatusInternalServerError)
		return
	}

	// Cleanup OAuth state
	if err = h.store.DeleteOAuthState(ctx, stateValue); err != nil {
		h.log.Error("Failed to delete oauth state", logz.Error(err))
		// The end user's request is technically not affected by this, and we could continue processing the request.
		// However, since the oauth state token is actually used to set the transfer token, out of abundance of
		// caution, we want to abort the request and return an error to prevent any replay attack scenarios.
		h.handleErrorNotOnWorkspaceHost(w, r, originalURL, errTextSomethingWentWrongOAuthCallback, http.StatusInternalServerError)
		return
	}

	// Redirect to original URL with transfer token in query param
	appendValueToParam(originalURL, transferTokenParamName, token)
	http.Redirect(w, r, originalURL.String(), http.StatusTemporaryRedirect)
}

// handleErrorNotOnWorkspaceHost redirects to the request with an additional transfer error in query param if there are
// any errors on hosts which are not on user's originally accessed URL e.g. oauth redirect callback.
func (h *authHandler) handleErrorNotOnWorkspaceHost(w http.ResponseWriter, r *http.Request, originalURL *url.URL, msg string, httpStatusCode uint32) {
	te := &TransferError{
		HttpStatusCode: httpStatusCode,
		Message:        msg,
	}
	teData, err := proto.Marshal(te)
	if err != nil {
		// This is a coding-mistake.
		h.handleProcessingError("Failed to proto.Marshal transfer error", err)
		http.Error(w, "Something went wrong. Please report an issue at https://gitlab.com/gitlab-org/gitlab/-/issues/new", http.StatusInternalServerError)
		return
	}
	encodedErr := base64.URLEncoding.EncodeToString(teData)
	// Redirect to original URL with transfer error in query param
	appendValueToParam(originalURL, transferErrorParamName, encodedErr)
	http.Redirect(w, r, originalURL.String(), http.StatusTemporaryRedirect)
}

// handleErrorOnWorkspaceHost sets the error cookie if there are any errors on hosts which are on user's originally
// accessed URL e.g. error while handling transfer token.
func (h *authHandler) handleErrorOnWorkspaceHost(w http.ResponseWriter, r *http.Request, originalURL *url.URL, msg string, httpStatusCode uint32) { //nolint: unparam
	te := &TransferError{
		HttpStatusCode: httpStatusCode,
		Message:        msg,
	}
	teData, err := proto.Marshal(te)
	if err != nil {
		// This is a coding-mistake.
		h.handleProcessingError("Failed to proto.Marshal transfer error", err)
		http.Error(w, "Something went wrong. Please report an issue at https://gitlab.com/gitlab-org/gitlab/-/issues/new", http.StatusInternalServerError)
		return
	}
	encodedErr := base64.URLEncoding.EncodeToString(teData)
	// Set error cookie and redirects back to the original URL.
	setCookie(w, r, errorCookieName, encodedErr, errorCookieTTL)
	http.Redirect(w, r, originalURL.String(), http.StatusTemporaryRedirect)
}

// exchangeCodeForUser exchanges OAuth code for user information.
func (h *authHandler) exchangeCodeForUser(ctx context.Context, code, codeVerifier string) (*User, error) {
	// Exchange code for token with PKCE
	token, err := h.oauthConfig.Exchange(
		ctx,
		code,
		oauth2.VerifierOption(codeVerifier),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to exchange code for token: %w", err)
	}

	// Get user info
	userInfo, err := h.oidcProvider.UserInfo(ctx, h.oauthConfig.TokenSource(ctx, token))
	if err != nil {
		return nil, fmt.Errorf("failed to get userinfo: %w", err)
	}

	userID, err := strconv.ParseInt(userInfo.Subject, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("failed to parse user ID as integer: %w", err)
	}

	return &User{Id: userID}, nil
}

func (h *authHandler) encodeOAuthState(originalURL *url.URL) (string, error) {
	stateParam := &OAuthStateParam{
		OriginalUrl: originalURL.String(),
		Nonce:       generateRandomToken(), // Generate random nonce for CSRF protection
		CreatedAt:   timestamppb.Now(),
	}

	protoData, err := proto.Marshal(stateParam)
	if err != nil {
		return "", err
	}

	return base64.URLEncoding.EncodeToString(protoData), nil
}

func (h *authHandler) decodeOAuthState(stateValue string) (*OAuthStateParam, error) {
	jsonData, err := base64.URLEncoding.DecodeString(stateValue)
	if err != nil {
		return nil, fmt.Errorf("invalid state encoding: %w", err)
	}

	var stateParam OAuthStateParam
	if err = proto.Unmarshal(jsonData, &stateParam); err != nil {
		return nil, fmt.Errorf("invalid state format: %w", err)
	}
	if err = h.validator.Validate(&stateParam); err != nil {
		return nil, fmt.Errorf("invalid state - protobuf validations failed: %w", err)
	}

	if !stateParam.IsValid(oauthStateParamTTL) {
		return nil, fmt.Errorf("state expired")
	}

	return &stateParam, nil
}

// OAuth helpers
func generateRandomToken() string {
	b := make([]byte, 32)
	// rand.Read never returns an error and would instead crash the program if it fails.
	_, _ = rand.Read(b)
	return base64.RawURLEncoding.EncodeToString(b)
}

// URL manipulation helpers
// appendValueToParam safely appends transfer token/error to URL preserving existing query params.
func appendValueToParam(u *url.URL, name, value string) {
	query := u.Query()
	query.Add(name, value)
	u.RawQuery = query.Encode()
}

// removeLastValueForParam removes the last occurrence of transfer token/error from URL.
func removeLastValueForParam(u *url.URL, name string) { //nolint: unparam
	query := u.Query()
	values, ok := query[name]
	if !ok {
		// Parameter doesn't exist, nothing to do
		return
	}
	switch len(values) {
	case 0:
		return
	case 1:
		query.Del(name)
	default:
		query[name] = values[:len(values)-1] // Remove only the last one
	}
	u.RawQuery = query.Encode()
}

// buildOriginalURL constructs the original URL from request.
func buildOriginalURL(r *http.Request) *url.URL {
	// Make a copy to not modify the original request.
	originalURL := *r.URL
	originalURL.Scheme = getRequestScheme(r)
	originalURL.Host = r.Host
	return &originalURL
}

// Request helpers
func getCookie(cookies []*http.Cookie, name string) *http.Cookie {
	for _, cookie := range cookies {
		if cookie.Name == name {
			return cookie
		}
	}
	return nil
}

func getRequestScheme(r *http.Request) string {
	switch {
	case r.Header.Get(httpz.XForwardedProtoHeader) == "https": // When behind a proxy on TLS
		return "https"
	case r.TLS != nil: // When exposed directly on TLS
		return "https"
	default:
		return "http"
	}
}

func setCookie(w http.ResponseWriter, r *http.Request, name, value string, ttl time.Duration) {
	cookie := &http.Cookie{
		Name:     name,
		Value:    value,
		Path:     "/",
		MaxAge:   int(ttl / time.Second),
		HttpOnly: true,
		Secure:   getRequestScheme(r) == "https",
		SameSite: http.SameSiteLaxMode,
	}
	http.SetCookie(w, cookie)
}

// Context helpers
// getUserFromContext is used by other handlers to extract the authenticated user details from the context.
func getUserFromContext(ctx context.Context) (*User, error) {
	u, ok := ctx.Value(userCtxKey{}).(*User)
	if !ok {
		return nil, fmt.Errorf("failed to get user from context")
	}
	return u, nil
}

// getWorkspaceFromContext is used by other handlers to extract the workspace being accessed by the authenticated user
// from the context.
func getWorkspaceFromContext(ctx context.Context) (*Workspace, error) {
	w, ok := ctx.Value(workspaceCtxKey{}).(*Workspace)
	if !ok {
		return nil, fmt.Errorf("failed to get workspace from context")
	}
	return w, nil
}
