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)
	})
}

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),
	)
}

//nolint:cyclop,gocognit,nestif
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()

		if tokens.IDTokenClaims != nil {
			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))
		}

		user, err := c.provider.GetUser(ctx, logger, tokens, userInfo)
		if err != nil {
			c.openvpn.DenyClient(ctx, logger, session.Client, "unable to fetch user data")
			c.writeHTTPError(ctx, w, logger, http.StatusInternalServerError, "fetchUser", err.Error())

			return
		}

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

		if err = c.provider.CheckUser(ctx, session, user, tokens); err != nil {
			c.openvpn.DenyClient(ctx, logger, session.Client, "client rejected")
			c.writeHTTPError(ctx, w, logger, http.StatusForbidden, "user validation", err.Error())

			return
		}

		if err = c.CheckTokenCEL(CELAuthModeInteractive, session, tokens); err != nil {
			c.openvpn.DenyClient(ctx, logger, session.Client, "client rejected")
			c.writeHTTPError(ctx, w, logger, http.StatusForbidden, "user validation", err.Error())

			return
		}

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

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

		if c.conf.OpenVPN.ClientConfig.UserSelector.Enabled {
			clientConfigProfiles := c.extractConfigProfilesFromIDToken(tokens)

			// If there's exactly one profile, use it directly
			if len(clientConfigProfiles) == 1 {
				c.openvpn.AcceptClient(ctx, logger, session.Client, username, clientConfigProfiles[0])
				c.postCodeExchangeHandlerStoreRefreshToken(ctx, logger, session, clientID, tokens)
				c.writeHTTPSuccess(ctx, w, logger)

				return
			}

			// Only show profile selector if there's more than one profile to choose from
			if len(clientConfigProfiles) > 1 {
				c.postCodeExchangeHandlerStoreRefreshToken(ctx, logger, session, clientID, tokens)

				clientConfigSelectorToken, err := c.createProfileSelectorToken(encryptedState, username, clientConfigProfiles)
				if err != nil {
					c.openvpn.DenyClient(ctx, logger, session.Client, "unable to create profile selector token")
					c.writeHTTPError(ctx, w, logger, http.StatusInternalServerError, "profile selector token", err.Error())

					return
				}

				err = c.storeProfileSelectorToken(ctx, clientConfigSelectorToken, clientID)
				if err != nil {
					c.openvpn.DenyClient(ctx, logger, session.Client, "unable to store profile selector token")
					c.writeHTTPError(ctx, w, logger, http.StatusInternalServerError, "store profile selector token", err.Error())

					return
				}

				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 {
					logger.LogAttrs(ctx, slog.LevelError, "template error", slog.Any(
						"err", fmt.Errorf("executing template: %w", err),
					))
					w.WriteHeader(http.StatusInternalServerError)

					return
				}

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

				return
			}

			// If no profiles, fall through to the default behavior
		}

		clientConfigName := username

		if clientConfigClaim := c.conf.OpenVPN.ClientConfig.TokenClaim; clientConfigClaim != "" && tokens.IDTokenClaims != nil {
			if claimValue, ok := tokens.IDTokenClaims.Claims[clientConfigClaim]; ok {
				if clientConfigTokenValue, ok := claimValue.(string); ok && clientConfigName != "" {
					clientConfigName = clientConfigTokenValue
				}
			}
		}

		c.openvpn.AcceptClient(ctx, logger, session.Client, username, clientConfigName)
		c.postCodeExchangeHandlerStoreRefreshToken(ctx, logger, session, clientID, tokens)
		c.writeHTTPSuccess(ctx, w, logger)
	}
}

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
}

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
}

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
}

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)
}

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, fmt.Errorf("executing template: %w", err).Error())
		w.WriteHeader(http.StatusInternalServerError)
	}
}

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", fmt.Errorf("executing template: %w", err),
		))
		w.WriteHeader(http.StatusInternalServerError)
	}
}
