package org.lingniu.idp.utils.jwt;

import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.jwk.KeyUse;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Date;
import java.util.UUID;

public class Jwks {
    
    /**
     * 从现有的RSA密钥生成JWK
     */
    public static RSAKey generateRsa(RSAPublicKey publicKey, RSAPrivateKey privateKey, String keyId) {
        RSAKey.Builder builder = new RSAKey.Builder(publicKey)
                .keyID(keyId)
                .keyUse(KeyUse.SIGNATURE)
                .algorithm(JWSAlgorithm.RS256);
        
        if (privateKey != null) {
            builder.privateKey(privateKey);
        }
        
        return builder.build();
    }
    
    /**
     * 仅生成公钥JWK（用于验证端）
     */
    public static RSAKey generatePublicJwk(RSAPublicKey publicKey, String keyId) {
        return new RSAKey.Builder(publicKey)
                .keyID(keyId)
                .keyUse(KeyUse.SIGNATURE)
                .algorithm(JWSAlgorithm.RS256)
                .build();
    }
    
    /**
     * 生成新的RSA密钥对（仅用于开发或测试）
     */
    public static RSAKey generateRsa() {
        return generateRsa(UUID.randomUUID().toString());
    }
    
    public static RSAKey generateRsa(String keyId) {
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
            keyPairGenerator.initialize(2048);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            
            RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
            RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
            
            return new RSAKey.Builder(publicKey)
                    .privateKey(privateKey)
                    .keyID(keyId)
                    .keyUse(KeyUse.SIGNATURE)
                    .algorithm(JWSAlgorithm.RS256)
                    .build();
        } catch (Exception e) {
            throw new RuntimeException("生成RSA密钥失败", e);
        }
    }
    
    /**
     * 创建带时间戳的密钥ID
     */
    public static String generateKeyIdWithTimestamp(String prefix) {
        return prefix + "-" + System.currentTimeMillis();
    }
    
    /**
     * 获取默认密钥ID
     */
    public static String getDefaultKeyId() {
        return "jwt-key-" + new Date().getTime();
    }
}