yao/openapi/oauth/signing.go

package oauth

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v4"
	"github.com/yaoapp/yao/openapi/oauth/types"
	"github.com/yaoapp/yao/share"
)

// SigningCertificates holds the loaded signing certificates and keys
type SigningCertificates struct {
	// Primary signing certificate and private key
	SigningCert    *x509.Certificate `json:"-"`
	SigningKey     interface{}       `json:"-"` // *rsa.PrivateKey, *ecdsa.PrivateKey, etc.
	SigningKeyPair *tls.Certificate  `json:"-"`

	// Verification certificates for token validation
	VerificationCerts []*x509.Certificate `json:"-"`

	// mTLS CA certificate for client validation
	MTLSClientCACert *x509.Certificate `json:"-"`

	// Signing algorithm
	Algorithm string `json:"algorithm"`

	// Certificate paths for reference
	SigningCertPath string `json:"signing_cert_path"`
	SigningKeyPath  string `json:"signing_key_path"`

	// Auto-generated flag
	IsAutoGenerated bool `json:"is_auto_generated"`
}

// LoadSigningCertificates loads or generates signing certificates based on configuration
func LoadSigningCertificates(config *types.SigningConfig) (*SigningCertificates, error) {
	certs := &SigningCertificates{
		Algorithm: config.SigningAlgorithm,
	}

	// Check if certificate files exist
	signingCertExists := fileExists(config.SigningCertPath)
	signingKeyExists := fileExists(config.SigningKeyPath)

	// If both files exist, try to load them
	if signingCertExists && signingKeyExists {
		err := loadExistingCertificates(certs, config)
		if err != nil {
			// If loading fails, log warning and generate new certificates
			fmt.Printf("Warning: Failed to load existing certificates (%v), generating new temporary certificates\n", err)
			return generateTemporaryCertificates(config)
		}
		return certs, nil
	}

	// If certificates don't exist, generate temporary ones
	return generateTemporaryCertificates(config)
}

// loadExistingCertificates loads certificates from the configured paths
func loadExistingCertificates(certs *SigningCertificates, config *types.SigningConfig) error {
	// Load signing certificate
	certPEM, err := os.ReadFile(config.SigningCertPath)
	if err != nil {
		return fmt.Errorf("failed to read signing certificate: %w", err)
	}

	certBlock, _ := pem.Decode(certPEM)
	if certBlock == nil {
		return fmt.Errorf("failed to decode signing certificate PEM")
	}

	signingCert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return fmt.Errorf("failed to parse signing certificate: %w", err)
	}

	// Load signing key
	keyPEM, err := os.ReadFile(config.SigningKeyPath)
	if err != nil {
		return fmt.Errorf("failed to read signing key: %w", err)
	}

	keyBlock, _ := pem.Decode(keyPEM)
	if keyBlock == nil {
		return fmt.Errorf("failed to decode signing key PEM")
	}

	var signingKey interface{}
	if config.SigningKeyPassword != "" {
		// Decrypt encrypted key
		keyBytes, err := x509.DecryptPEMBlock(keyBlock, []byte(config.SigningKeyPassword))
		if err != nil {
			return fmt.Errorf("failed to decrypt signing key: %w", err)
		}
		signingKey, err = parsePrivateKey(keyBytes)
		if err != nil {
			return fmt.Errorf("failed to parse decrypted signing key: %w", err)
		}
	} else {
		// Parse unencrypted key
		var err error
		signingKey, err = parsePrivateKey(keyBlock.Bytes)
		if err != nil {
			return fmt.Errorf("failed to parse signing key: %w", err)
		}
	}

	// Create TLS certificate pair
	keyPair, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return fmt.Errorf("failed to create key pair: %w", err)
	}

	certs.SigningCert = signingCert
	certs.SigningKey = signingKey
	certs.SigningKeyPair = &keyPair
	certs.SigningCertPath = config.SigningCertPath
	certs.SigningKeyPath = config.SigningKeyPath
	certs.IsAutoGenerated = false

	// Load verification certificates if configured
	if len(config.VerificationCerts) > 0 {
		verificationCerts, err := loadVerificationCertificates(config.VerificationCerts)
		if err != nil {
			return fmt.Errorf("failed to load verification certificates: %w", err)
		}
		certs.VerificationCerts = verificationCerts
	}

	// Load mTLS CA certificate if configured
	if config.MTLSClientCACertPath != "" {
		mtlsCACert, err := loadCertificateFromFile(config.MTLSClientCACertPath)
		if err != nil {
			return fmt.Errorf("failed to load mTLS CA certificate: %w", err)
		}
		certs.MTLSClientCACert = mtlsCACert
	}

	return nil
}

// generateTemporaryCertificates generates temporary self-signed certificates
func generateTemporaryCertificates(config *types.SigningConfig) (*SigningCertificates, error) {
	// Generate RSA private key
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("failed to generate private key: %w", err)
	}

	// Create certificate template
	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization:  []string{share.App.Name},
			Country:       []string{"US"},
			Province:      []string{""},
			Locality:      []string{""},
			StreetAddress: []string{""},
			PostalCode:    []string{""},
			CommonName:    fmt.Sprintf("%s OAuth Signing Certificate", share.App.Name),
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(365 * 24 * time.Hour), // Valid for 1 year
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	// Generate certificate
	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create certificate: %w", err)
	}

	// Parse the generated certificate
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, fmt.Errorf("failed to parse generated certificate: %w", err)
	}

	// Create PEM blocks
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	})

	keyDER, err := x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal private key: %w", err)
	}

	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: keyDER,
	})

	// Create TLS certificate pair
	keyPair, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return nil, fmt.Errorf("failed to create key pair: %w", err)
	}

	// Determine system directory for storing temporary certificates
	systemDir := getSystemCertificateDirectory()

	// Create directory if it doesn't exist
	if err := os.MkdirAll(systemDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create system certificate directory: %w", err)
	}

	// Generate unique filenames
	timestamp := time.Now().Format("20060102150405")
	certPath := filepath.Join(systemDir, fmt.Sprintf("oauth_signing_cert_%s.pem", timestamp))
	keyPath := filepath.Join(systemDir, fmt.Sprintf("oauth_signing_key_%s.pem", timestamp))

	// Save certificate and key to system directory
	if err := os.WriteFile(certPath, certPEM, 0644); err != nil {
		return nil, fmt.Errorf("failed to write certificate file: %w", err)
	}

	if err := os.WriteFile(keyPath, keyPEM, 0600); err != nil {
		return nil, fmt.Errorf("failed to write key file: %w", err)
	}

	fmt.Printf("Generated temporary OAuth signing certificate at: %s\n", certPath)
	fmt.Printf("Generated temporary OAuth signing key at: %s\n", keyPath)

	return &SigningCertificates{
		SigningCert:     cert,
		SigningKey:      privateKey,
		SigningKeyPair:  &keyPair,
		Algorithm:       config.SigningAlgorithm,
		SigningCertPath: certPath,
		SigningKeyPath:  keyPath,
		IsAutoGenerated: true,
	}, nil
}

// loadVerificationCertificates loads additional verification certificates
func loadVerificationCertificates(certPaths []string) ([]*x509.Certificate, error) {
	var certs []*x509.Certificate

	for _, certPath := range certPaths {
		cert, err := loadCertificateFromFile(certPath)
		if err != nil {
			return nil, fmt.Errorf("failed to load verification certificate %s: %w", certPath, err)
		}
		certs = append(certs, cert)
	}

	return certs, nil
}

// loadCertificateFromFile loads a certificate from a PEM file
func loadCertificateFromFile(certPath string) (*x509.Certificate, error) {
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read certificate file: %w", err)
	}

	certBlock, _ := pem.Decode(certPEM)
	if certBlock == nil {
		return nil, fmt.Errorf("failed to decode certificate PEM")
	}

	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse certificate: %w", err)
	}

	return cert, nil
}

// parsePrivateKey parses a private key from DER bytes
func parsePrivateKey(der []byte) (interface{}, error) {
	// Try PKCS#8 first
	if key, err := x509.ParsePKCS8PrivateKey(der); err == nil {
		return key, nil
	}

	// Try PKCS#1 RSA
	if key, err := x509.ParsePKCS1PrivateKey(der); err == nil {
		return key, nil
	}

	// Try EC private key
	if key, err := x509.ParseECPrivateKey(der); err == nil {
		return key, nil
	}

	return nil, fmt.Errorf("unable to parse private key")
}

// fileExists checks if a file exists
func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	return !os.IsNotExist(err)
}

// getSystemCertificateDirectory returns the appropriate system directory for storing certificates
func getSystemCertificateDirectory() string {
	// Use different directories based on the operating system
	homeDir, err := os.UserHomeDir()
	if err != nil {
		// Fallback to temporary directory
		return filepath.Join(os.TempDir(), "yao-oauth-certs")
	}

	// Create a hidden directory in user's home
	return filepath.Join(homeDir, ".yao", "oauth", "certs")
}

// ValidateCertificate validates a certificate for OAuth signing
func (c *SigningCertificates) ValidateCertificate() error {
	if c.SigningCert == nil {
		return fmt.Errorf("signing certificate is nil")
	}

	// Check if certificate is expired
	now := time.Now()
	if now.Before(c.SigningCert.NotBefore) {
		return fmt.Errorf("signing certificate is not yet valid")
	}

	if now.After(c.SigningCert.NotAfter) {
		return fmt.Errorf("signing certificate has expired")
	}

	// Check if certificate has appropriate key usage
	if c.SigningCert.KeyUsage&x509.KeyUsageDigitalSignature == 0 {
		return fmt.Errorf("signing certificate does not have digital signature key usage")
	}

	return nil
}

// GetPublicKey returns the public key from the signing certificate
func (c *SigningCertificates) GetPublicKey() interface{} {
	if c.SigningCert == nil {
		return nil
	}
	return c.SigningCert.PublicKey
}

// GetKeyID returns a key identifier for the signing certificate
func (c *SigningCertificates) GetKeyID() string {
	if c.SigningCert == nil {
		return ""
	}

	// Use the certificate's serial number as key ID
	return c.SigningCert.SerialNumber.String()
}

// CleanupTemporaryCertificates removes auto-generated temporary certificates
func (c *SigningCertificates) CleanupTemporaryCertificates() error {
	if !c.IsAutoGenerated {
		return nil // Don't delete user-provided certificates
	}

	var errs []error

	if c.SigningCertPath != "" && fileExists(c.SigningCertPath) {
		if err := os.Remove(c.SigningCertPath); err != nil {
			errs = append(errs, fmt.Errorf("failed to remove certificate file %s: %w", c.SigningCertPath, err))
		}
	}

	if c.SigningKeyPath != "" && fileExists(c.SigningKeyPath) {
		if err := os.Remove(c.SigningKeyPath); err != nil {
			errs = append(errs, fmt.Errorf("failed to remove key file %s: %w", c.SigningKeyPath, err))
		}
	}

	if len(errs) > 0 {
		return fmt.Errorf("cleanup errors: %v", errs)
	}

	return nil
}

// Service signing certificate methods

// GetSigningCertificates returns the signing certificates for the service
func (s *Service) GetSigningCertificates() *SigningCertificates {
	return s.signingCerts
}

// GetSigningKey returns the signing private key
func (s *Service) GetSigningKey() interface{} {
	if s.signingCerts == nil {
		return nil
	}
	return s.signingCerts.SigningKey
}

// GetSigningCertificate returns the signing certificate
func (s *Service) GetSigningCertificate() interface{} {
	if s.signingCerts == nil {
		return nil
	}
	return s.signingCerts.SigningCert
}

// GetSigningAlgorithm returns the signing algorithm
func (s *Service) GetSigningAlgorithm() string {
	if s.signingCerts == nil {
		return "RS256" // default
	}
	return s.signingCerts.Algorithm
}

// GetKeyID returns the key identifier for JWT token signing
func (s *Service) GetKeyID() string {
	if s.signingCerts == nil {
		return ""
	}
	return s.signingCerts.GetKeyID()
}

// SignToken signs a token based on the configured format (jwt or opaque)
// extraClaims: optional extra claims to add to the token (e.g., team_id, tenant_id)
func (s *Service) SignToken(tokenType, clientID, scope, subject string, expiresIn int, extraClaims ...map[string]interface{}) (string, error) {
	var claims map[string]interface{}
	if len(extraClaims) > 0 {
		claims = extraClaims[0]
	}

	switch s.config.Token.AccessTokenFormat {
	case "jwt":
		return s.signJWTToken(tokenType, clientID, scope, subject, expiresIn, claims)
	case "opaque":
		return s.signOpaqueToken(tokenType, clientID, scope, subject)
	default:
		// Default to JWT if format is not specified or unknown
		return s.signJWTToken(tokenType, clientID, scope, subject, expiresIn, claims)
	}
}

// VerifyToken verifies a token based on its format and returns token claims
func (s *Service) VerifyToken(token string) (*types.TokenClaims, error) {
	if strings.Contains(token, ".") {
		return s.verifyJWTToken(token)
	}
	return s.verifyOpaqueToken(token)
}

// VerifyTokenAllowExpired verifies token signature but allows expired tokens.
// Used by Guard to parse expired access tokens before attempting refresh.
func (s *Service) VerifyTokenAllowExpired(token string) (*types.TokenClaims, error) {
	if strings.Contains(token, ".") {
		return s.verifyJWTTokenAllowExpired(token)
	}
	return s.verifyOpaqueToken(token)
}

// VerifyRefreshToken verifies a refresh token based on its format.
// For opaque tokens it looks up the refresh token store (not the access token store).
func (s *Service) VerifyRefreshToken(token string) (*types.TokenClaims, error) {
	if strings.Contains(token, ".") {
		// JWT refresh tokens can be verified with the same JWT logic
		return s.verifyJWTToken(token)
	}
	return s.verifyOpaqueRefreshToken(token)
}

// verifyOpaqueRefreshToken verifies an opaque refresh token using the refresh token store.
func (s *Service) verifyOpaqueRefreshToken(token string) (*types.TokenClaims, error) {
	tokenInfo, err := s.getRefreshTokenData(token)
	if err != nil {
		return nil, fmt.Errorf("refresh token not found or invalid: %w", err)
	}

	clientID, _ := tokenInfo["client_id"].(string)
	scope, _ := tokenInfo["scope"].(string)
	subject, _ := tokenInfo["subject"].(string)

	claims := &types.TokenClaims{
		Subject:   subject,
		ClientID:  clientID,
		Scope:     scope,
		TokenType: "refresh_token",
		Issuer:    s.config.IssuerURL,
	}

	if issuedAt, ok := tokenInfo["issued_at"].(int64); ok {
		claims.IssuedAt = time.Unix(issuedAt, 0)
	}

	if expiresAt, ok := tokenInfo["expires_at"].(int64); ok {
		claims.ExpiresAt = time.Unix(expiresAt, 0)
		if time.Now().After(claims.ExpiresAt) {
			return nil, fmt.Errorf("refresh token expired")
		}
	}

	return claims, nil
}

// SignIDToken signs an ID token with specific parameters and stores it
func (s *Service) SignIDToken(clientID, scope string, expiresIn int, userdata *types.OIDCUserInfo, extraClaims ...map[string]interface{}) (string, error) {
	if s.signingCerts == nil || s.signingCerts.SigningKey == nil {
		return "", fmt.Errorf("signing certificates not initialized")
	}

	// Check if userdata and subject are provided
	if userdata == nil || userdata.Sub == "" {
		return "", fmt.Errorf("userdata or userdata.Sub is required for ID token")
	}

	tokenSubject := userdata.Sub

	now := time.Now()

	// Create OIDC ID Token claims
	idTokenClaims := &types.OIDCIDToken{
		// Required ID Token claims
		Iss: s.config.IssuerURL,
		Sub: tokenSubject,
		Aud: clientID,
		Exp: now.Add(time.Duration(expiresIn) * time.Second).Unix(),
		Iat: now.Unix(),
	}

	// Create a map for custom claims that includes both standard and user info
	claims := jwt.MapClaims{
		// Standard OIDC ID Token claims
		"iss": idTokenClaims.Iss,
		"sub": idTokenClaims.Sub,
		"aud": idTokenClaims.Aud,
		"exp": idTokenClaims.Exp,
		"iat": idTokenClaims.Iat,
		"jti": generateJTI(),
	}

	// Add extra claims if provided (e.g., team_id, tenant_id)
	if len(extraClaims) > 0 {
		for key, value := range extraClaims[0] {
			claims[key] = value
		}
	}

	// Add user information from userdata
	// Add standard OIDC user claims if they exist
	if userdata.Name != "" {
		claims["name"] = userdata.Name
	}
	if userdata.GivenName != "" {
		claims["given_name"] = userdata.GivenName
	}
	if userdata.FamilyName != "" {
		claims["family_name"] = userdata.FamilyName
	}
	if userdata.MiddleName != "" {
		claims["middle_name"] = userdata.MiddleName
	}
	if userdata.Nickname != "" {
		claims["nickname"] = userdata.Nickname
	}
	if userdata.PreferredUsername != "" {
		claims["preferred_username"] = userdata.PreferredUsername
	}
	if userdata.Profile != "" {
		claims["profile"] = userdata.Profile
	}
	if userdata.Picture != "" {
		claims["picture"] = userdata.Picture
	}
	if userdata.Website != "" {
		claims["website"] = userdata.Website
	}
	if userdata.Email != "" {
		claims["email"] = userdata.Email
	}
	if userdata.EmailVerified != nil {
		claims["email_verified"] = *userdata.EmailVerified
	}
	if userdata.Gender != "" {
		claims["gender"] = userdata.Gender
	}
	if userdata.Birthdate != "" {
		claims["birthdate"] = userdata.Birthdate
	}
	if userdata.Zoneinfo != "" {
		claims["zoneinfo"] = userdata.Zoneinfo
	}
	if userdata.Locale != "" {
		claims["locale"] = userdata.Locale
	}
	if userdata.PhoneNumber != "" {
		claims["phone_number"] = userdata.PhoneNumber
	}
	if userdata.PhoneNumberVerified != nil {
		claims["phone_number_verified"] = *userdata.PhoneNumberVerified
	}
	if userdata.Address != nil {
		claims["address"] = userdata.Address
	}
	if userdata.UpdatedAt != nil {
		claims["updated_at"] = *userdata.UpdatedAt
	}

	// Add Yao custom fields with namespace
	if userdata.YaoUserID != "" {
		claims["yao:user_id"] = userdata.YaoUserID
	}
	if userdata.YaoTenantID != "" {
		claims["yao:tenant_id"] = userdata.YaoTenantID
	}
	if userdata.YaoTeamID != "" {
		claims["yao:team_id"] = userdata.YaoTeamID
	}
	if userdata.YaoIsOwner != nil {
		claims["yao:is_owner"] = *userdata.YaoIsOwner
	}
	if userdata.YaoTypeID != "" {
		claims["yao:type_id"] = userdata.YaoTypeID
	}
	if userdata.YaoAuthSource != "" {
		claims["yao:auth_source"] = userdata.YaoAuthSource
	}
	// Add Yao team info if present
	if userdata.YaoTeam != nil {
		teamMap := make(map[string]interface{})
		if userdata.YaoTeam.TeamID != "" {
			teamMap["team_id"] = userdata.YaoTeam.TeamID
		}
		if userdata.YaoTeam.Logo != "" {
			teamMap["logo"] = userdata.YaoTeam.Logo
		}
		if userdata.YaoTeam.Name != "" {
			teamMap["name"] = userdata.YaoTeam.Name
		}
		if userdata.YaoTeam.OwnerID != "" {
			teamMap["owner_id"] = userdata.YaoTeam.OwnerID
		}
		if userdata.YaoTeam.Description != "" {
			teamMap["description"] = userdata.YaoTeam.Description
		}
		if userdata.YaoTeam.UpdatedAt != nil {
			teamMap["updated_at"] = *userdata.YaoTeam.UpdatedAt
		}
		if len(teamMap) > 0 {
			claims["yao:team"] = teamMap
		}
	}
	// Add Yao member info if present (for team context)
	if userdata.YaoMember != nil {
		memberMap := make(map[string]interface{})
		if userdata.YaoMember.MemberID != "" {
			memberMap["member_id"] = userdata.YaoMember.MemberID
		}
		if userdata.YaoMember.DisplayName != "" {
			memberMap["display_name"] = userdata.YaoMember.DisplayName
		}
		if userdata.YaoMember.Bio != "" {
			memberMap["bio"] = userdata.YaoMember.Bio
		}
		if userdata.YaoMember.Avatar != "" {
			memberMap["avatar"] = userdata.YaoMember.Avatar
		}
		if userdata.YaoMember.Email != "" {
			memberMap["email"] = userdata.YaoMember.Email
		}
		if len(memberMap) > 0 {
			claims["yao:member"] = memberMap
		}
	}
	// Add Yao type info if present
	if userdata.YaoType != nil {
		typeMap := make(map[string]interface{})
		if userdata.YaoType.TypeID != "" {
			typeMap["type_id"] = userdata.YaoType.TypeID
		}
		if userdata.YaoType.Name != "" {
			typeMap["name"] = userdata.YaoType.Name
		}
		if userdata.YaoType.Locale != "" {
			typeMap["locale"] = userdata.YaoType.Locale
		}
		if len(typeMap) > 0 {
			claims["yao:type"] = typeMap
		}
	}

	// Add scope if provided (useful for determining which claims to include)
	if scope != "" {
		claims["scope"] = scope
	}

	// Create token with claims
	token := jwt.NewWithClaims(getSigningMethod(s.config.Token.AccessTokenSigningAlg), claims)

	// Set key ID in header
	token.Header["kid"] = s.GetKeyID()

	// Set token type in header
	token.Header["typ"] = "JWT"

	// Sign token with private key
	signedToken, err := token.SignedString(s.signingCerts.SigningKey)
	if err != nil {
		return "", fmt.Errorf("failed to sign ID token: %w", err)
	}

	return signedToken, nil
}

// signJWTToken signs a JWT token using the configured signing algorithm
func (s *Service) signJWTToken(tokenType, clientID, scope, subject string, expiresIn int, extraClaims map[string]interface{}) (string, error) {
	if s.signingCerts == nil || s.signingCerts.SigningKey == nil {
		return "", fmt.Errorf("signing certificates not initialized")
	}

	now := time.Now()

	// Use MapClaims to support extra claims
	claims := jwt.MapClaims{
		// Standard JWT claims
		"iss": s.config.IssuerURL,
		"sub": subject,
		"aud": clientID,
		"exp": now.Add(time.Duration(expiresIn) * time.Second).Unix(),
		"nbf": now.Unix(),
		"iat": now.Unix(),
		"jti": generateJTI(),
		// Custom OAuth claims
		"client_id":  clientID,
		"scope":      scope,
		"token_type": tokenType,
	}

	// Add extra claims if provided (e.g., team_id, tenant_id)
	for key, value := range extraClaims {
		claims[key] = value
	}

	// Create token with claims
	token := jwt.NewWithClaims(getSigningMethod(s.config.Token.AccessTokenSigningAlg), claims)

	// Set key ID in header
	token.Header["kid"] = s.GetKeyID()

	// Sign token with private key
	return token.SignedString(s.signingCerts.SigningKey)
}

// verifyJWTToken verifies a JWT token and returns its claims
func (s *Service) verifyJWTToken(tokenString string) (*types.TokenClaims, error) {
	return s.parseJWTToken(tokenString, false)
}

// verifyJWTTokenAllowExpired parses a JWT token, verifying signature but allowing expiration.
// Returns claims even if the token is expired (signature must still be valid).
func (s *Service) verifyJWTTokenAllowExpired(tokenString string) (*types.TokenClaims, error) {
	return s.parseJWTToken(tokenString, true)
}

// parseJWTToken is the shared JWT parsing logic.
// When allowExpired is true, expired tokens are still parsed (signature-only verification).
func (s *Service) parseJWTToken(tokenString string, allowExpired bool) (*types.TokenClaims, error) {
	if s.signingCerts == nil || s.signingCerts.SigningCert == nil {
		return nil, fmt.Errorf("signing certificates not initialized")
	}

	parserOpts := []jwt.ParserOption{}
	if allowExpired {
		parserOpts = append(parserOpts, jwt.WithoutClaimsValidation())
	}

	token, err := jwt.ParseWithClaims(tokenString, jwt.MapClaims{}, func(token *jwt.Token) (interface{}, error) {
		expectedMethod := getSigningMethod(s.config.Token.AccessTokenSigningAlg)
		if token.Method != expectedMethod {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return s.signingCerts.GetPublicKey(), nil
	}, parserOpts...)

	if err != nil {
		return nil, fmt.Errorf("failed to parse JWT token: %w", err)
	}

	if !allowExpired && !token.Valid {
		return nil, fmt.Errorf("invalid JWT token")
	}

	mapClaims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return nil, fmt.Errorf("invalid JWT claims type")
	}

	return s.extractTokenClaims(mapClaims), nil
}

// extractTokenClaims converts jwt.MapClaims to types.TokenClaims
func (s *Service) extractTokenClaims(mapClaims jwt.MapClaims) *types.TokenClaims {
	tokenClaims := &types.TokenClaims{
		Extra: make(map[string]interface{}),
	}

	if sub, ok := mapClaims["sub"].(string); ok {
		tokenClaims.Subject = sub
	}
	if clientID, ok := mapClaims["client_id"].(string); ok {
		tokenClaims.ClientID = clientID
	}
	if scope, ok := mapClaims["scope"].(string); ok {
		tokenClaims.Scope = scope
	}
	if tokenType, ok := mapClaims["token_type"].(string); ok {
		tokenClaims.TokenType = tokenType
	}
	if iss, ok := mapClaims["iss"].(string); ok {
		tokenClaims.Issuer = iss
	}
	if jti, ok := mapClaims["jti"].(string); ok {
		tokenClaims.JTI = jti
	}

	if exp, ok := mapClaims["exp"].(float64); ok {
		tokenClaims.ExpiresAt = time.Unix(int64(exp), 0)
	}
	if iat, ok := mapClaims["iat"].(float64); ok {
		tokenClaims.IssuedAt = time.Unix(int64(iat), 0)
	}

	if aud, ok := mapClaims["aud"].(string); ok {
		tokenClaims.Audience = []string{aud}
	} else if audArray, ok := mapClaims["aud"].([]interface{}); ok {
		audience := make([]string, 0, len(audArray))
		for _, a := range audArray {
			if audStr, ok := a.(string); ok {
				audience = append(audience, audStr)
			}
		}
		tokenClaims.Audience = audience
	}

	if teamID, ok := mapClaims["team_id"].(string); ok {
		tokenClaims.TeamID = teamID
	}
	if tenantID, ok := mapClaims["tenant_id"].(string); ok {
		tokenClaims.TenantID = tenantID
	}

	standardClaims := map[string]bool{
		"sub": true, "client_id": true, "scope": true, "token_type": true,
		"exp": true, "iat": true, "nbf": true, "iss": true, "aud": true, "jti": true,
		"team_id": true, "tenant_id": true,
	}
	for key, value := range mapClaims {
		if !standardClaims[key] {
			tokenClaims.Extra[key] = value
		}
	}

	return tokenClaims
}

// signOpaqueToken signs an opaque token using HMAC or RSA signature
func (s *Service) signOpaqueToken(tokenType, clientID, scope, subject string) (string, error) {
	// Generate base opaque token
	baseToken, err := s.generateOpaqueTokenBase(tokenType, clientID)
	if err != nil {
		return "", fmt.Errorf("failed to generate base opaque token: %w", err)
	}

	// Create token metadata for signature
	tokenData := fmt.Sprintf("%s.%s.%s.%s.%d", baseToken, clientID, scope, subject, time.Now().Unix())

	// Sign the token data
	signature, err := s.signData([]byte(tokenData))
	if err != nil {
		return "", fmt.Errorf("failed to sign opaque token: %w", err)
	}

	// Combine base token with signature
	signedToken := fmt.Sprintf("%s.%s", baseToken, base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(signature))

	return signedToken, nil
}

// verifyOpaqueToken verifies an opaque token signature and returns token claims
func (s *Service) verifyOpaqueToken(token string) (*types.TokenClaims, error) {
	parts := strings.Split(token, ".")
	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid opaque token format")
	}

	baseToken := parts[0]
	signaturePart := parts[len(parts)-1]

	// Decode signature
	signature, err := base64.URLEncoding.WithPadding(base64.NoPadding).DecodeString(signaturePart)
	if err != nil {
		return nil, fmt.Errorf("failed to decode token signature: %w", err)
	}

	// Extract token information from store
	tokenInfo, err := s.getAccessTokenData(token)
	if err != nil {
		return nil, fmt.Errorf("token not found or invalid: %w", err)
	}

	// Reconstruct token data for verification
	clientID := tokenInfo["client_id"].(string)
	scope := ""
	if scopeVal, ok := tokenInfo["scope"].(string); ok {
		scope = scopeVal
	}
	subject := ""
	if subjectVal, ok := tokenInfo["subject"].(string); ok {
		subject = subjectVal
	}
	issuedAt := tokenInfo["issued_at"].(int64)

	tokenData := fmt.Sprintf("%s.%s.%s.%s.%d", baseToken, clientID, scope, subject, issuedAt)

	// Verify signature
	if err := s.verifySignature([]byte(tokenData), signature); err != nil {
		return nil, fmt.Errorf("invalid token signature: %w", err)
	}

	// Build token claims
	tokenClaims := &types.TokenClaims{
		Subject:   subject,
		ClientID:  clientID,
		Scope:     scope,
		TokenType: "access_token",
		IssuedAt:  time.Unix(issuedAt, 0),
		Issuer:    s.config.IssuerURL,
	}

	if expiresAt, ok := tokenInfo["expires_at"].(int64); ok {
		tokenClaims.ExpiresAt = time.Unix(expiresAt, 0)
	}

	return tokenClaims, nil
}

// signData signs data using the configured signing key
func (s *Service) signData(data []byte) ([]byte, error) {
	if s.signingCerts == nil || s.signingCerts.SigningKey == nil {
		return nil, fmt.Errorf("signing key not available")
	}

	switch key := s.signingCerts.SigningKey.(type) {
	case *rsa.PrivateKey:
		// Use RSA-PSS for signing
		hash := sha256.Sum256(data)
		signature, err := rsa.SignPSS(rand.Reader, key, crypto.SHA256, hash[:], nil)
		if err != nil {
			return nil, fmt.Errorf("failed to sign with RSA key: %w", err)
		}
		return signature, nil
	default:
		return nil, fmt.Errorf("unsupported signing key type: %T", key)
	}
}

// verifySignature verifies a signature using the configured public key
func (s *Service) verifySignature(data []byte, signature []byte) error {
	if s.signingCerts == nil || s.signingCerts.SigningCert == nil {
		return fmt.Errorf("signing certificate not available")
	}

	switch pubKey := s.signingCerts.GetPublicKey().(type) {
	case *rsa.PublicKey:
		// Use RSA-PSS for verification
		hash := sha256.Sum256(data)
		err := rsa.VerifyPSS(pubKey, crypto.SHA256, hash[:], signature, nil)
		if err != nil {
			return fmt.Errorf("failed to verify RSA signature: %w", err)
		}
		return nil
	default:
		return fmt.Errorf("unsupported public key type: %T", pubKey)
	}
}

// generateOpaqueTokenBase generates the base part of an opaque token
func (s *Service) generateOpaqueTokenBase(tokenType, clientID string) (string, error) {
	// Generate random bytes for token
	randomBytes := make([]byte, 32)
	if _, err := rand.Read(randomBytes); err != nil {
		return "", fmt.Errorf("failed to generate random bytes: %w", err)
	}

	// Create base token with type, client ID, timestamp, and random component
	randomPart := base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(randomBytes)
	timestamp := time.Now().Format("20060102150405")

	return fmt.Sprintf("%s_%s_%s_%s", tokenType, clientID, timestamp, randomPart), nil
}

// getSigningMethod returns the JWT signing method for the given algorithm
func getSigningMethod(algorithm string) jwt.SigningMethod {
	switch algorithm {
	case "RS256":
		return jwt.SigningMethodRS256
	case "RS384":
		return jwt.SigningMethodRS384
	case "RS512":
		return jwt.SigningMethodRS512
	case "PS256":
		return jwt.SigningMethodPS256
	case "PS384":
		return jwt.SigningMethodPS384
	case "PS512":
		return jwt.SigningMethodPS512
	default:
		return jwt.SigningMethodRS256 // Default
	}
}

// generateJTI generates a unique JWT ID
func generateJTI() string {
	randomBytes := make([]byte, 16)
	rand.Read(randomBytes)
	return base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(randomBytes)
}