package org.lingniu.idp.utils.jwt;

import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import lombok.Getter;
import org.lingniu.idp.config.JwtProperties;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;
import org.springframework.util.FileCopyUtils;
import org.springframework.util.StringUtils;

import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.text.ParseException;
import java.util.Base64;
import java.util.Date;
import java.util.Map;
import java.util.UUID;

@Component
public class JwtUtil {

    @Getter
    private final RSAPrivateKey privateKey;
    @Getter
    private final RSAPublicKey publicKey;
    private final JWSHeader header;
    private final long expiration;
    private final String tokenPrefix;
    
    // 构造函数
    public JwtUtil(JwtProperties jwtProperties) throws Exception {
        this.expiration = jwtProperties.getExpiration().toMillis();
        this.tokenPrefix = jwtProperties.getPrefix();
        
        JwtProperties.RsaKey rsaConfig = jwtProperties.getRsa();
        
        // 加载或生成密钥
        if (StringUtils.hasText(rsaConfig.getPrivateKey()) && 
            StringUtils.hasText(rsaConfig.getPublicKey())) {
            // 从配置加载密钥
            this.privateKey = loadPrivateKey(rsaConfig.getPrivateKey());
            this.publicKey = loadPublicKey(rsaConfig.getPublicKey());
        } else {
            // 自动生成密钥（仅用于开发）
            KeyPair keyPair = generateKeyPair(rsaConfig.getKeySize());
            this.privateKey = (RSAPrivateKey) keyPair.getPrivate();
            this.publicKey = (RSAPublicKey) keyPair.getPublic();
            
            System.out.println("⚠️ 未配置RSA密钥，已自动生成新的密钥对");
            System.out.println("公钥（保存到配置中）:");
            System.out.println(toPemFormat(publicKey));
            System.out.println("\n私钥（保存到配置中）:");
            System.out.println(toPemFormat(privateKey));
        }
        
        // 创建JWT头部
        this.header = new JWSHeader.Builder(JWSAlgorithm.RS256)
                .type(JOSEObjectType.JWT)
                .keyID("jwt-key-" + System.currentTimeMillis())
                .build();
    }
    
    /**
     * 加载私钥
     */
    private RSAPrivateKey loadPrivateKey(String keyContent) throws Exception {
        // 检查是否是文件路径
        if (keyContent.startsWith("classpath:") || keyContent.startsWith("file:")) {
            keyContent = loadKeyFromFile(keyContent);
        }
        
        // 清理PEM格式
        String cleanedKey = cleanPemKey(keyContent, "PRIVATE");
        
        // Base64解码
        byte[] keyBytes = Base64.getDecoder().decode(cleanedKey);
        
        // 创建密钥规范
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        
        return (RSAPrivateKey) keyFactory.generatePrivate(spec);
    }
    
    /**
     * 加载公钥
     */
    private RSAPublicKey loadPublicKey(String keyContent) throws Exception {
        // 检查是否是文件路径
        if (keyContent.startsWith("classpath:") || keyContent.startsWith("file:")) {
            keyContent = loadKeyFromFile(keyContent);
        }
        
        // 清理PEM格式
        String cleanedKey = cleanPemKey(keyContent, "PUBLIC");
        
        // Base64解码
        byte[] keyBytes = Base64.getDecoder().decode(cleanedKey);
        
        // 创建密钥规范
        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        
        return (RSAPublicKey) keyFactory.generatePublic(spec);
    }
    
    /**
     * 从文件加载密钥
     */
    private String loadKeyFromFile(String filePath) throws Exception {
        Resource resource;
        if (filePath.startsWith("classpath:")) {
            resource = new ClassPathResource(filePath.substring(10));
        } else if (filePath.startsWith("file:")) {
            resource = new FileSystemResource(filePath.substring(5));
        } else {
            resource = new ClassPathResource(filePath);
        }
        
        if (!resource.exists()) {
            throw new IllegalArgumentException("密钥文件不存在: " + filePath);
        }
        
        byte[] bytes = FileCopyUtils.copyToByteArray(resource.getInputStream());
        return new String(bytes);
    }
    
    /**
     * 清理PEM密钥
     */
    private String cleanPemKey(String key, String keyType) {
        if (!StringUtils.hasText(key)) {
            throw new IllegalArgumentException(keyType + "密钥内容为空");
        }
        
        // 移除所有空白字符
        key = key.replaceAll("\\s+", "");
        
        // 如果没有PEM标记，直接返回
        if (!key.contains("-----BEGIN")) {
            return key;
        }
        
        // 提取PEM内容
        String beginMarker = "-----BEGIN" + keyType + "KEY-----";
        String endMarker = "-----END" + keyType + "KEY-----";
        
        int beginIndex = key.indexOf(beginMarker);
        if (beginIndex == -1) {
            beginMarker = "-----BEGIN" + keyType + " KEY-----";
            beginIndex = key.indexOf(beginMarker);
        }
        
        if (beginIndex == -1) {
            throw new IllegalArgumentException("无效的PEM格式: 找不到BEGIN标记");
        }
        
        beginIndex += beginMarker.length();
        int endIndex = key.indexOf(endMarker, beginIndex);
        if (endIndex == -1) {
            endMarker = "-----END" + keyType + " KEY-----";
            endIndex = key.indexOf(endMarker, beginIndex);
        }
        
        if (endIndex == -1) {
            throw new IllegalArgumentException("无效的PEM格式: 找不到END标记");
        }
        
        return key.substring(beginIndex, endIndex).replaceAll("\\s+", "");
    }
    
    /**
     * 生成密钥对
     */
    private KeyPair generateKeyPair(int keySize) throws Exception {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        keyGen.initialize(keySize);
        return keyGen.generateKeyPair();
    }
    
    /**
     * 创建JWT令牌
     */
    public String createToken(Map<String, Object> claims) {
        return createToken(claims, this.expiration);
    }
    
    public String createToken(Map<String, Object> claims, long expirationMs) {
        try {
            // 构建声明
            JWTClaimsSet.Builder claimsBuilder = new JWTClaimsSet.Builder();
            
            // 添加自定义声明
            for (Map.Entry<String, Object> entry : claims.entrySet()) {
                claimsBuilder.claim(entry.getKey(), entry.getValue());
            }
            
            Date now = new Date();
            claimsBuilder
                .issueTime(now)
                .expirationTime(new Date(now.getTime() + expirationMs))
                .notBeforeTime(now)
                .jwtID(UUID.randomUUID().toString()); // 唯一的JWT ID
            
            JWTClaimsSet claimsSet = claimsBuilder.build();
            
            // 创建签名JWT
            SignedJWT signedJWT = new SignedJWT(header, claimsSet);
            
            // 使用RSA私钥进行签名
            JWSSigner signer = new RSASSASigner(this.privateKey);
            signedJWT.sign(signer);
            
            return signedJWT.serialize();
            
        } catch (JOSEException e) {
            throw new JwtException("创建JWT令牌失败", e);
        }
    }
    
    /**
     * 解析并验证JWT令牌
     */
    public JWTClaimsSet parseToken(String token) {
        try {
            // 提取实际令牌（去除Bearer前缀）
            String actualToken = extractToken(token);
            
            // 解析JWT
            SignedJWT signedJWT = SignedJWT.parse(actualToken);
            
            // 使用RSA公钥验证签名
            JWSVerifier verifier = new RSASSAVerifier(this.publicKey);
            if (!signedJWT.verify(verifier)) {
                throw new JwtException("JWT签名验证失败");
            }
            
            // 获取声明集
            JWTClaimsSet claims = signedJWT.getJWTClaimsSet();
            
            // 验证过期时间
            if (claims.getExpirationTime() != null && 
                claims.getExpirationTime().before(new Date())) {
                throw new JwtException("JWT令牌已过期");
            }
            
            // 验证生效时间
            if (claims.getNotBeforeTime() != null && 
                claims.getNotBeforeTime().after(new Date())) {
                throw new JwtException("JWT令牌尚未生效");
            }
            
            return claims;
            
        } catch (ParseException e) {
            throw new JwtException("JWT令牌解析失败", e);
        } catch (JOSEException e) {
            throw new JwtException("JWT签名验证异常", e);
        }
    }
    
    /**
     * 提取令牌（去除前缀）
     */
    public String extractToken(String authHeader) {
        if (authHeader != null && authHeader.startsWith(tokenPrefix)) {
            return authHeader.substring(tokenPrefix.length()).trim();
        }
        return authHeader;
    }
    
    /**
     * 获取Base64编码的公钥
     */
    public String getPublicKeyBase64() {
        return Base64.getEncoder().encodeToString(this.publicKey.getEncoded());
    }
    
    /**
     * 获取PEM格式的公钥
     */
    public String toPemFormat(RSAPublicKey publicKey) {
        String base64Key = Base64.getEncoder().encodeToString(publicKey.getEncoded());
        return formatAsPem(base64Key, "PUBLIC");
    }
    
    /**
     * 获取PEM格式的私钥
     */
    public String toPemFormat(RSAPrivateKey privateKey) {
        String base64Key = Base64.getEncoder().encodeToString(privateKey.getEncoded());
        return formatAsPem(base64Key, "PRIVATE");
    }
    
    /**
     * 格式化PEM
     */
    private String formatAsPem(String base64Key, String keyType) {
        StringBuilder pem = new StringBuilder();
        pem.append("-----BEGIN ").append(keyType).append(" KEY-----\n");
        
        // 每64个字符换行
        for (int i = 0; i < base64Key.length(); i += 64) {
            pem.append(base64Key, i, Math.min(base64Key.length(), i + 64)).append("\n");
        }
        
        pem.append("-----END ").append(keyType).append(" KEY-----\n");
        return pem.toString();
    }
    
    /**
     * 验证令牌是否有效
     */
    public boolean validateToken(String token) {
        try {
            parseToken(token);
            return true;
        } catch (JwtException e) {
            return false;
        }
    }
    
    /**
     * 获取令牌中的声明值
     */
    public Object getClaim(String token, String claimName) {
        JWTClaimsSet claims = parseToken(token);
        return claims.getClaim(claimName);
    }
    
    /**
     * 获取令牌中的用户名（假设存储在"username"声明中）
     */
    public String getUsername(String token) {
        try {
            return parseToken(token).getStringClaim("username");
        } catch (ParseException e) {
            throw new JwtException("获取用户名失败", e);
        }
    }
    
    /**
     * 获取令牌过期时间
     */
    public Date getExpiration(String token) {
        return parseToken(token).getExpirationTime();
    }
    
    /**
     * 获取令牌签发时间
     */
    public Date getIssuedAt(String token) {
        return parseToken(token).getIssueTime();
    }
    
    /**
     * 获取JWT头部信息
     */
    public JWSHeader getHeader(String token) {
        try {
            return SignedJWT.parse(extractToken(token)).getHeader();
        } catch (ParseException e) {
            throw new JwtException("获取JWT头部失败", e);
        }
    }
    
    /**
     * 自定义异常类
     */
    public static class JwtException extends RuntimeException {
        public JwtException(String message) {
            super(message);
        }
        public JwtException(String message, Throwable cause) {
            super(message, cause);
        }
    }
}