package jwt_utils

import (
	"time"

	"github.com/awnumar/memguard"
	jwt "github.com/golang-jwt/jwt/v5"
)

// GenerateJWTToken Generate the `access token` for the secret key.
// SECURITY: HMAC secret is wiped from memory after signing to prevent memory dump attacks.
func GenerateJWTToken(hmacSecret []byte, uuid string, ad time.Duration) (string, time.Time, error) {
	// SECURITY: Create a copy of the secret and wipe the copy after use
	// Note: The original hmacSecret is owned by the caller
	secretCopy := make([]byte, len(hmacSecret))
	copy(secretCopy, hmacSecret)
	defer memguard.WipeBytes(secretCopy) // SECURITY: Wipe secret copy after signing

	token := jwt.New(jwt.SigningMethodHS256)
	expiresIn := time.Now().Add(ad)

	// CWE-391: Safe type assertion even though we just created the token
	// Defensive programming to prevent future panics if jwt library changes
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return "", expiresIn, jwt.ErrTokenInvalidClaims
	}

	claims["session_uuid"] = uuid
	claims["exp"] = expiresIn.Unix()

	tokenString, err := token.SignedString(secretCopy)
	if err != nil {
		return "", expiresIn, err
	}

	return tokenString, expiresIn, nil
}

// GenerateJWTTokenPair Generate the `access token` and `refresh token` for the secret key.
// SECURITY: HMAC secret is wiped from memory after signing to prevent memory dump attacks.
func GenerateJWTTokenPair(hmacSecret []byte, uuid string, ad time.Duration, rd time.Duration) (string, time.Time, string, time.Time, error) {
	// SECURITY: Create a copy of the secret and wipe the copy after use
	secretCopy := make([]byte, len(hmacSecret))
	copy(secretCopy, hmacSecret)
	defer memguard.WipeBytes(secretCopy) // SECURITY: Wipe secret copy after signing

	//
	// Generate token.
	//
	token := jwt.New(jwt.SigningMethodHS256)
	expiresIn := time.Now().Add(ad)

	// CWE-391: Safe type assertion even though we just created the token
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return "", time.Now(), "", time.Now(), jwt.ErrTokenInvalidClaims
	}

	claims["session_uuid"] = uuid
	claims["exp"] = expiresIn.Unix()

	tokenString, err := token.SignedString(secretCopy)
	if err != nil {
		return "", time.Now(), "", time.Now(), err
	}

	//
	// Generate refresh token.
	//
	refreshToken := jwt.New(jwt.SigningMethodHS256)
	refreshExpiresIn := time.Now().Add(rd)

	// CWE-391: Safe type assertion for refresh token
	rtClaims, ok := refreshToken.Claims.(jwt.MapClaims)
	if !ok {
		return "", time.Now(), "", time.Now(), jwt.ErrTokenInvalidClaims
	}

	rtClaims["session_uuid"] = uuid
	rtClaims["exp"] = refreshExpiresIn.Unix()

	refreshTokenString, err := refreshToken.SignedString(secretCopy)
	if err != nil {
		return "", time.Now(), "", time.Now(), err
	}

	return tokenString, expiresIn, refreshTokenString, refreshExpiresIn, nil
}

// ProcessJWTToken validates either the `access token` or `refresh token` and returns either the `uuid` if success or error on failure.
// CWE-347: Implements proper algorithm validation to prevent JWT algorithm confusion attacks
// OWASP A02:2021: Cryptographic Failures - Prevents token forgery through algorithm switching
// SECURITY: HMAC secret copy is wiped from memory after validation.
func ProcessJWTToken(hmacSecret []byte, reqToken string) (string, error) {
	// SECURITY: Create a copy of the secret and wipe the copy after use
	secretCopy := make([]byte, len(hmacSecret))
	copy(secretCopy, hmacSecret)
	defer memguard.WipeBytes(secretCopy) // SECURITY: Wipe secret copy after validation

	token, err := jwt.Parse(reqToken, func(t *jwt.Token) (any, error) {
		// CRITICAL SECURITY FIX: Validate signing method to prevent algorithm confusion attacks
		// Protects against:
		// 1. "none" algorithm bypass (CVE-2015-9235)
		// 2. HS256/RS256 algorithm confusion (CVE-2016-5431)
		// 3. Token forgery through algorithm switching
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, jwt.ErrTokenSignatureInvalid
		}

		// Additional check: Ensure it's specifically HS256
		if t.Method.Alg() != "HS256" {
			return nil, jwt.ErrTokenSignatureInvalid
		}

		return secretCopy, nil
	})
	if err == nil && token.Valid {
		if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
			// Safe type assertion with validation
			sessionUUID, ok := claims["session_uuid"].(string)
			if !ok {
				return "", jwt.ErrTokenInvalidClaims
			}
			return sessionUUID, nil
		}
		return "", err
	}
	return "", err
}