package oauth2

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/jkroepke/openvpn-auth-oauth2/internal/config"
	"github.com/jkroepke/openvpn-auth-oauth2/internal/oauth2/idtoken"
	"github.com/jkroepke/openvpn-auth-oauth2/internal/oauth2/types"
	"github.com/jkroepke/openvpn-auth-oauth2/internal/state"
	"github.com/zitadel/logging"
	"github.com/zitadel/oidc/v3/pkg/client/rp"
)

const (
	title   = "title"
	message = "message"
)

type openvpnManagementClient interface {
	AcceptClient(ctx context.Context, logger *slog.Logger, client state.ClientIdentifier, username string, clientConfigName string)
	DenyClient(ctx context.Context, logger *slog.Logger, client state.ClientIdentifier, reason string)
}

// OAuth2Start returns a http.Handler that starts the OAuth2 authorization flow.
// It checks if the request has a valid state GET parameter generated by state.New.
// Optionally, it checks the HTTP client IP address against the VPN IP address.
// After the checks, the request is delegated to [rp.AuthURLHandler].
func (c *Client) OAuth2Start() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()

		// check if request has a state GET parameter generated state.New.
		encryptedOIDCState := r.URL.Query().Get("state")
		if encryptedOIDCState == "" {
			w.WriteHeader(http.StatusBadRequest)

			return
		}

		// decode the state GET parameter
		session, err := state.Decrypt(c.stateCrypto, encryptedOIDCState)
		if err != nil {
			c.logger.LogAttrs(ctx, slog.LevelWarn, "invalid state: "+err.Error())
			w.WriteHeader(http.StatusBadRequest)

			return
		}

		logger := c.createSessionLogger(session)

		if c.conf.HTTP.Check.IPAddr {
			if err := checkClientIPAddr(r, c.conf, session); err != nil {
				logger.LogAttrs(ctx, slog.LevelWarn, err.Error())

				if !errors.Is(err, ErrClientRejected) {
					c.openvpn.DenyClient(ctx, logger, session.Client, "client rejected")
					w.WriteHeader(http.StatusInternalServerError)

					return
				}

				c.openvpn.DenyClient(ctx, logger, session.Client, err.Error())
				w.WriteHeader(http.StatusForbidden)

				return
			}
		}

		logger.LogAttrs(ctx, slog.LevelInfo, "initialize authorization via oauth2")

		authorizeParams := c.authorizeParams

		if c.conf.OAuth2.Nonce {
			authorizeParams = append(authorizeParams, rp.WithURLParam("nonce", c.getNonce(c.getClientID(session))))
		}

		rp.AuthURLHandler(func() string {
			return encryptedOIDCState
		}, c.relyingParty, authorizeParams...).ServeHTTP(w, r)
	})
}

// OAuth2Callback returns an http.Handler that handles the OAuth2 callback.
func (c *Client) OAuth2Callback() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
		defer cancel()

		encryptedOAuth2State := r.URL.Query().Get("state")
		if encryptedOAuth2State == "" {
			c.writeHTTPError(ctx, w, c.logger, http.StatusBadRequest, "Bad Request", "state is empty")

			return
		}

		oAuth2State, err := state.Decrypt(c.stateCrypto, encryptedOAuth2State)
		if err != nil {
			c.writeHTTPError(ctx, w, c.logger, http.StatusBadRequest, "Invalid State", err.Error())

			return
		}

		logger := c.createSessionLoggerWithState(oAuth2State)

		ctx = logging.ToContext(ctx, logger)

		clientID := c.getClientID(oAuth2State)

		if c.conf.OAuth2.Nonce {
			ctx = context.WithValue(ctx, types.CtxNonce{}, c.getNonce(clientID))
			r = r.WithContext(ctx)
		}

		var codeExchangeHandler rp.CodeExchangeCallback[*idtoken.Claims]

		if c.conf.OAuth2.UserInfo {
			codeExchangeHandler = rp.UserinfoCallback(c.postCodeExchangeHandler(logger, encryptedOAuth2State, oAuth2State, clientID))
		} else {
			codeExchangeHandler = func(
				w http.ResponseWriter, r *http.Request,
				tokens idtoken.IDToken, encryptedOAuth2State state.EncryptedState, provider rp.RelyingParty,
			) {
				c.postCodeExchangeHandler(logger, encryptedOAuth2State, oAuth2State, clientID)(w, r, tokens, encryptedOAuth2State, provider, nil)
			}
		}

		rp.CodeExchangeHandler(
			codeExchangeHandler,
			c.relyingParty,
		).ServeHTTP(w, r)
	})
}

// OAuth2ProfileSubmit returns an http.Handler that handles the profile selection form submission.
func (c *Client) OAuth2ProfileSubmit() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
		defer cancel()

		// G120: Parsing form data without limiting request body size can allow memory exhaustion
		r.Body = http.MaxBytesReader(w, r.Body, 1<<20)

		encryptedToken := r.FormValue("token")
		if encryptedToken == "" {
			c.writeHTTPError(ctx, w, c.logger, http.StatusBadRequest, "Bad Request", "token is empty")

			return
		}

		tokenBytes, err := c.stateCrypto.DecryptBytesWithTime([]byte(encryptedToken))
		if err != nil {
			c.writeHTTPError(ctx, w, c.logger, http.StatusBadRequest, "Invalid token", err.Error())

			return
		}

		var token clientConfigToken

		if err = json.Unmarshal(tokenBytes, &token); err != nil {
			c.writeHTTPError(ctx, w, c.logger, http.StatusBadRequest, "Invalid token", "unable to parse token")

			return
		}

		oAuth2State, err := state.Decrypt(c.stateCrypto, token.EncryptedOAuth2State)
		if err != nil {
			c.writeHTTPError(ctx, w, c.logger, http.StatusBadRequest, "Invalid State", err.Error())

			return
		}

		logger := c.createSessionLoggerWithState(oAuth2State)

		clientID := c.getClientID(oAuth2State)
		stateHash := sha256.Sum256([]byte(encryptedToken))
		storageKey := fmt.Sprintf("%s-token-%x", clientID, stateHash[:8])

		encryptedStoredToken, err := c.storage.Get(ctx, storageKey)
		if err != nil {
			c.writeHTTPError(ctx, w, logger, http.StatusBadRequest, "unable to retrieve refresh token from storage", err.Error())

			return
		}

		_ = c.storage.Delete(ctx, storageKey)

		if encryptedStoredToken != encryptedToken {
			c.writeHTTPError(ctx, w, logger, http.StatusBadRequest, "Bad Request", "token mismatch from profile selector")

			return
		}

		username := token.Username
		if username == config.CommonNameModeOmitValue {
			username = ""
		}

		clientConfigName := r.FormValue("profile")
		logger = logger.With(slog.String("selected_profile", clientConfigName))

		// Validate that the selected profile is in the allowed list
		if !slices.Contains(token.Profiles, clientConfigName) {
			c.openvpn.DenyClient(ctx, logger, oAuth2State.Client, "invalid profile selection")
			c.writeHTTPError(ctx, w, logger, http.StatusForbidden, "Invalid Profile", "The selected profile is not allowed")

			return
		}

		logger.LogAttrs(ctx, slog.LevelInfo, "successful authorization via oauth2 with profile selection")

		c.openvpn.AcceptClient(ctx, logger, oAuth2State.Client, username, clientConfigName)
		c.writeHTTPSuccess(ctx, w, logger)
	})
}

// getClientID returns the token storage key for a session based on the refresh-token configuration.
func (c *Client) getClientID(session state.State) string {
	if c.conf.OAuth2.Refresh.UseSessionID && session.Client.SessionID != "" {
		return session.Client.SessionID
	}

	return strconv.FormatUint(session.Client.CID, 10)
}

// createSessionLogger creates a logger with common session information.
func (c *Client) createSessionLogger(session state.State) *slog.Logger {
	return c.logger.With(
		slog.String("ip", fmt.Sprintf("%s:%s", session.IPAddr, session.IPPort)),
		slog.Uint64("cid", session.Client.CID),
		slog.Uint64("kid", session.Client.KID),
		slog.String("common_name", session.Client.CommonName),
	)
}

// createSessionLoggerWithState creates a logger with session information including session_id and session_state.
func (c *Client) createSessionLoggerWithState(session state.State) *slog.Logger {
	return c.logger.With(
		slog.String("ip", fmt.Sprintf("%s:%s", session.IPAddr, session.IPPort)),
		slog.Uint64("cid", session.Client.CID),
		slog.Uint64("kid", session.Client.KID),
		slog.String("common_name", session.Client.CommonName),
		slog.String("session_id", session.Client.SessionID),
		slog.String("session_state", session.SessionState),
	)
}

// postCodeExchangeHandler returns the callback that validates exchanged tokens and accepts or denies the OpenVPN client.
func (c *Client) postCodeExchangeHandler(
	logger *slog.Logger, encryptedState state.EncryptedState, session state.State, clientID string,
) rp.CodeExchangeUserinfoCallback[*idtoken.Claims, *types.UserInfo] {
	return func(
		w http.ResponseWriter, r *http.Request, tokens idtoken.IDToken, _ string,
		_ rp.RelyingParty, userInfo *types.UserInfo,
	) {
		ctx := r.Context()

		logger = withIDTokenClaimsLogger(ctx, logger, tokens)

		user, err := c.fetchUser(ctx, logger, tokens, userInfo)
		if err != nil {
			c.openvpn.DenyClient(ctx, logger, session.Client, err.openvpnReason)
			c.writeHTTPError(ctx, w, logger, err.httpStatus, err.errorType, err.err.Error())

			return
		}

		logger = logger.With(
			slog.String("user_subject", user.Subject),
			slog.String("user_username", user.Username),
		)

		if err = c.validateInteractiveUser(ctx, session, user, tokens); err != nil {
			c.openvpn.DenyClient(ctx, logger, session.Client, err.openvpnReason)
			c.writeHTTPError(ctx, w, logger, err.httpStatus, err.errorType, err.err.Error())

			return
		}

		logger.LogAttrs(ctx, slog.LevelInfo, "successful authorization via oauth2")

		username := user.Username
		if username == "" {
			username = session.Client.CommonName
		}

		codeExchange := codeExchangeRequest{
			logger:         logger,
			encryptedState: encryptedState,
			session:        session,
			clientID:       clientID,
			tokens:         tokens,
			username:       username,
		}

		if c.handleUserSelectedClientConfig(ctx, w, codeExchange) {
			return
		}

		clientConfigName := c.clientConfigNameFromToken(tokens, username)

		c.acceptOAuth2Client(ctx, w, codeExchange, clientConfigName)
	}
}

// withIDTokenClaimsLogger enriches the logger with ID token claim fields when claims are available.
func withIDTokenClaimsLogger(ctx context.Context, logger *slog.Logger, tokens idtoken.IDToken) *slog.Logger {
	if tokens.IDTokenClaims == nil {
		return logger
	}

	logger = logger.With(
		slog.String("idtoken_subject", tokens.IDTokenClaims.Subject),
		slog.String("idtoken_email", tokens.IDTokenClaims.EMail),
		slog.String("idtoken_preferred_username", tokens.IDTokenClaims.PreferredUsername),
	)

	logger.LogAttrs(ctx, slog.LevelDebug, "claims", slog.Any("claims", tokens.IDTokenClaims.Claims))

	return logger
}

type userValidationError struct {
	err           error
	errorType     string
	openvpnReason string
	httpStatus    int
}

// fetchUser resolves the authenticated user from token claims or UserInfo data.
func (c *Client) fetchUser(
	ctx context.Context,
	logger *slog.Logger,
	tokens idtoken.IDToken,
	userInfo *types.UserInfo,
) (types.UserInfo, *userValidationError) {
	user, err := c.provider.GetUser(ctx, logger, tokens, userInfo)
	if err != nil {
		return types.UserInfo{}, &userValidationError{
			err:           err,
			httpStatus:    http.StatusInternalServerError,
			errorType:     "fetchUser",
			openvpnReason: "unable to fetch user data",
		}
	}

	return user, nil
}

// validateInteractiveUser applies provider-specific and CEL checks for an interactive OAuth2 login.
func (c *Client) validateInteractiveUser(
	ctx context.Context,
	session state.State,
	user types.UserInfo,
	tokens idtoken.IDToken,
) *userValidationError {
	if err := c.provider.CheckUser(ctx, session, user, tokens); err != nil {
		return failedUserValidation(err)
	}

	if err := c.CheckTokenCEL(CELAuthModeInteractive, session, tokens); err != nil {
		return failedUserValidation(err)
	}

	return nil
}

// failedUserValidation converts a validation error into the common denial response metadata.
func failedUserValidation(err error) *userValidationError {
	return &userValidationError{
		err:           err,
		httpStatus:    http.StatusForbidden,
		errorType:     "user validation",
		openvpnReason: "client rejected",
	}
}

type codeExchangeRequest struct {
	logger         *slog.Logger
	tokens         idtoken.IDToken
	encryptedState state.EncryptedState
	clientID       string
	username       string
	session        state.State
}

// handleUserSelectedClientConfig handles client configuration profile selection when it is enabled.
func (c *Client) handleUserSelectedClientConfig(ctx context.Context, w http.ResponseWriter, req codeExchangeRequest) bool {
	if !c.conf.OpenVPN.ClientConfig.UserSelector.Enabled {
		return false
	}

	clientConfigProfiles := c.extractConfigProfilesFromIDToken(req.tokens)

	switch len(clientConfigProfiles) {
	case 0:
		return false
	case 1:
		c.acceptOAuth2Client(ctx, w, req, clientConfigProfiles[0])

		return true
	default:
		c.postCodeExchangeHandlerStoreRefreshToken(ctx, req.logger, req.session, req.clientID, req.tokens)

		if err := c.renderClientConfigProfileSelector(ctx, w, req, clientConfigProfiles); err != nil {
			c.openvpn.DenyClient(ctx, req.logger, req.session.Client, err.openvpnReason)
			c.writeHTTPError(ctx, w, req.logger, http.StatusInternalServerError, err.errorType, err.err.Error())
		}

		return true
	}
}

type profileSelectorError struct {
	err           error
	errorType     string
	openvpnReason string
}

// renderClientConfigProfileSelector stores a selector token and renders the profile selection page.
func (c *Client) renderClientConfigProfileSelector(
	ctx context.Context,
	w http.ResponseWriter,
	req codeExchangeRequest,
	clientConfigProfiles []string,
) *profileSelectorError {
	clientConfigSelectorToken, err := c.createProfileSelectorToken(req.encryptedState, req.username, clientConfigProfiles)
	if err != nil {
		return &profileSelectorError{
			err:           err,
			errorType:     "profile selector token",
			openvpnReason: "unable to create profile selector token",
		}
	}

	err = c.storeProfileSelectorToken(ctx, clientConfigSelectorToken, req.clientID)
	if err != nil {
		return &profileSelectorError{
			err:           err,
			errorType:     "store profile selector token",
			openvpnReason: "unable to store profile selector token",
		}
	}

	err = c.conf.HTTP.Template.Execute(w, map[string]any{
		title:                  "Select Profile",
		message:                "Please select your client configuration profile.",
		"token":                clientConfigSelectorToken,
		"clientConfigProfiles": clientConfigProfiles,
	})
	if err != nil {
		req.logger.LogAttrs(ctx, slog.LevelError, "template error", slog.Any(
			"err", fmt.Errorf("executing template: %w", err),
		))
		w.WriteHeader(http.StatusInternalServerError)

		return nil
	}

	req.logger.LogAttrs(
		ctx, slog.LevelInfo, "presented client configuration profile selector",
		slog.String("profiles", strings.Join(clientConfigProfiles, ", ")),
	)

	return nil
}

// clientConfigNameFromToken returns the client configuration name from the configured token claim when present.
func (c *Client) clientConfigNameFromToken(tokens idtoken.IDToken, defaultName string) string {
	clientConfigClaim := c.conf.OpenVPN.ClientConfig.TokenClaim
	if clientConfigClaim == "" || tokens.IDTokenClaims == nil || defaultName == "" {
		return defaultName
	}

	claimValue, ok := tokens.IDTokenClaims.Claims[clientConfigClaim]
	if !ok {
		return defaultName
	}

	clientConfigName, ok := claimValue.(string)
	if !ok {
		return defaultName
	}

	return clientConfigName
}

// acceptOAuth2Client accepts the OpenVPN client, stores refresh state, and renders the success page.
func (c *Client) acceptOAuth2Client(
	ctx context.Context,
	w http.ResponseWriter,
	req codeExchangeRequest,
	clientConfigName string,
) {
	c.openvpn.AcceptClient(ctx, req.logger, req.session.Client, req.username, clientConfigName)
	c.postCodeExchangeHandlerStoreRefreshToken(ctx, req.logger, req.session, req.clientID, req.tokens)
	c.writeHTTPSuccess(ctx, w, req.logger)
}

// postCodeExchangeHandlerStoreRefreshToken stores refresh data for future non-interactive authentication.
func (c *Client) postCodeExchangeHandlerStoreRefreshToken(
	ctx context.Context, logger *slog.Logger, session state.State, clientID string, tokens idtoken.IDToken,
) {
	if !c.conf.OAuth2.Refresh.Enabled {
		return
	}

	if !c.conf.OAuth2.Refresh.ValidateUser {
		if err := c.storage.Set(ctx, clientID, types.EmptyToken); err != nil {
			logger.LogAttrs(ctx, slog.LevelWarn, err.Error())
		} else {
			logger.LogAttrs(ctx, slog.LevelDebug, "empty token for non-interactive re-authentication stored")
		}

		return
	}

	refreshToken, err := c.provider.GetRefreshToken(tokens)
	if err != nil {
		logLevel := slog.LevelWarn

		if errors.Is(err, ErrNoRefreshToken) {
			if session.SessionState == "AuthenticatedEmptyUser" || session.SessionState == "Authenticated" {
				logLevel = slog.LevelDebug
			}
		}

		logger.LogAttrs(ctx, logLevel, fmt.Errorf("oauth2.refresh is enabled, but %w", err).Error())
	}

	if refreshToken == "" {
		logger.LogAttrs(ctx, slog.LevelWarn, "refresh token is empty")
	} else if err = c.storage.Set(ctx, clientID, refreshToken); err != nil {
		logger.LogAttrs(
			ctx, slog.LevelWarn, "unable to store refresh token",
			slog.Any("err", err),
		)
	} else {
		logger.LogAttrs(ctx, slog.LevelDebug, "refresh token for non-interactive re-authentication stored")
	}
}

// extractConfigProfilesFromIDToken extracts the list of available profiles from static configuration and token claims.
func (c *Client) extractConfigProfilesFromIDToken(tokens idtoken.IDToken) []string {
	profiles := c.conf.OpenVPN.ClientConfig.UserSelector.StaticValues

	clientConfigClaim := c.conf.OpenVPN.ClientConfig.TokenClaim
	if clientConfigClaim == "" || tokens.IDTokenClaims == nil {
		return profiles
	}

	iClaimValue, ok := tokens.IDTokenClaims.Claims[clientConfigClaim]
	if !ok {
		return profiles
	}

	switch claimValue := iClaimValue.(type) {
	case []any:
		profiles = slices.Grow(profiles, len(claimValue))

		for _, profile := range claimValue {
			if profileStr, ok := profile.(string); ok {
				profiles = append(profiles, profileStr)
			}
		}
	case string:
		profiles = append(profiles, claimValue)
	}

	return profiles
}

// storeProfileSelectorToken stores an encrypted profile selector token under a short-lived lookup key.
func (c *Client) storeProfileSelectorToken(ctx context.Context, encryptedToken, clientID string) error {
	stateHash := sha256.Sum256([]byte(encryptedToken))
	storageKey := fmt.Sprintf("%s-token-%x", clientID, stateHash[:8])

	err := c.storage.Set(ctx, storageKey, encryptedToken)
	if err != nil {
		return fmt.Errorf("unable to store profile selector token: %w", err)
	}

	return nil
}

// createProfileSelectorToken creates an encrypted token containing the available client configuration profiles.
func (c *Client) createProfileSelectorToken(encryptedState state.EncryptedState, username string, profiles []string) (string, error) {
	if username == "" {
		username = config.CommonNameModeOmitValue
	}

	tokenContent := clientConfigToken{
		Username:             username,
		Profiles:             profiles,
		EncryptedOAuth2State: encryptedState,
	}

	tokenBytes, err := json.Marshal(tokenContent)
	if err != nil {
		return "", fmt.Errorf("unable to marshal profile selector token: %w", err)
	}

	encryptedToken, err := c.stateCrypto.EncryptBytesWithTime(tokenBytes)
	if err != nil {
		return "", fmt.Errorf("unable to encrypt profile selector token: %w", err)
	}

	return string(encryptedToken), nil
}

// httpErrorHandler maps OAuth2 relying party errors to OpenVPN denial and an HTTP error page.
func (c *Client) httpErrorHandler(ctx context.Context, w http.ResponseWriter, httpStatus int, errorType, errorDesc, encryptedSession string) {
	logger := c.logger

	session, err := state.Decrypt(c.stateCrypto, encryptedSession)
	if err == nil {
		logger = c.createSessionLogger(session)

		c.openvpn.DenyClient(ctx, logger, session.Client, "client rejected")
	} else {
		logger.LogAttrs(ctx, slog.LevelDebug, "httpErrorHandler: "+err.Error())
	}

	c.writeHTTPError(ctx, w, logger, httpStatus, errorType, errorDesc)
}

// writeHTTPError writes the configured access-denied page and logs a generated error ID.
func (c *Client) writeHTTPError(ctx context.Context, w http.ResponseWriter, logger *slog.Logger, httpCode int, errorType, errorDesc string) {
	if httpCode == http.StatusUnauthorized {
		httpCode = http.StatusForbidden
	}

	h := sha256.New()
	h.Write([]byte(time.Now().String()))

	errorID := hex.EncodeToString(h.Sum(nil))

	logger.LogAttrs(ctx, slog.LevelWarn, fmt.Sprintf("%s: %s", errorType, errorDesc), slog.String("error_id", errorID))
	w.WriteHeader(httpCode)

	err := c.conf.HTTP.Template.Execute(w, map[string]string{
		title:     "Access denied",
		message:   "Please contact your administrator.",
		"errorID": errorID,
	})
	if err != nil {
		logger.LogAttrs(ctx, slog.LevelError, "template error", slog.Any(
			"err", err,
		))
		w.WriteHeader(http.StatusInternalServerError)
	}
}

// writeHTTPSuccess writes the configured access-granted page.
func (c *Client) writeHTTPSuccess(ctx context.Context, w http.ResponseWriter, logger *slog.Logger) {
	err := c.conf.HTTP.Template.Execute(w, map[string]string{
		title:   "Access granted",
		message: "You can close this window now.",
	})
	if err != nil {
		logger.LogAttrs(ctx, slog.LevelError, "template error", slog.Any(
			"err", err,
		))
		w.WriteHeader(http.StatusInternalServerError)
	}
}
