package com.xiaopeng.common.core.utils;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA 加密解密工具类
 * <p>
 * 支持功能：
 * 1. 公钥加密、私钥解密（OAEP 填充模式）
 * 2. 私钥签名、公钥验签
 * 3. 密钥对生成
 * 4. Base64 编码/解码
 * <p>
 * 注意事项：
 * - RSA 使用 OAEP 填充模式，更安全，默认密钥长度 2048 位，最大加密数据长度为 190 字节
 * - 超长数据需要分段加密或使用对称加密（如 AES）
 * - OAEP 填充使用 SHA-256 和 MGF1 算法
 * - 使用 ThreadLocal 缓存 Cipher 实例，提升高并发性能
 */
public class RSAUtil {

    private static final String RSA_ALGORITHM = "RSA";
    private static final String SIGNATURE_ALGORITHM = "SHA256withRSA";
    private static final String TRANSFORMATION = "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";
    private static final int KEY_SIZE = 2048;

    /**
     * 缓存 Cipher 实例，提高高并发性能
     * Cipher 不是线程安全的，每个线程使用自己的实例
     */
    private static final ThreadLocal<Cipher> ENCRYPT_CIPHER = ThreadLocal.withInitial(() -> {
        try {
            return Cipher.getInstance(TRANSFORMATION);
        } catch (Exception e) {
            throw new RuntimeException("初始化加密 Cipher 失败", e);
        }
    });

    private static final ThreadLocal<Cipher> DECRYPT_CIPHER = ThreadLocal.withInitial(() -> {
        try {
            return Cipher.getInstance(TRANSFORMATION);
        } catch (Exception e) {
            throw new RuntimeException("初始化解密 Cipher 失败", e);
        }
    });

    private static final ThreadLocal<java.security.Signature> SIGNATURE = ThreadLocal.withInitial(() -> {
        try {
            return java.security.Signature.getInstance(SIGNATURE_ALGORITHM);
        } catch (Exception e) {
            throw new RuntimeException("初始化 Signature 失败", e);
        }
    });

    /**
     * 生成 RSA 密钥对
     *
     * @return 密钥对
     */
    public static KeyPair generateKeyPair() {
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(RSA_ALGORITHM);
            keyPairGenerator.initialize(KEY_SIZE);
            return keyPairGenerator.generateKeyPair();
        } catch (Exception e) {
            throw new RuntimeException("生成密钥对失败", e);
        }
    }

    /**
     * 获取公钥 Base64 字符串
     *
     * @param keyPair 密钥对
     * @return Base64 编码的公钥
     */
    public static String getPublicKey(KeyPair keyPair) {
        return Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
    }

    /**
     * 获取私钥 Base64 字符串
     *
     * @param keyPair 密钥对
     * @return Base64 编码的私钥
     */
    public static String getPrivateKey(KeyPair keyPair) {
        return Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
    }

    /**
     * 使用公钥加密
     *
     * @param data       待加密数据
     * @param publicKey  Base64 编码的公钥
     * @return Base64 编码的密文
     */
    public static String encryptByPublicKey(String data, String publicKey) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(publicKey);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PublicKey key = keyFactory.generatePublic(keySpec);

            Cipher cipher = ENCRYPT_CIPHER.get();
            cipher.init(Cipher.ENCRYPT_MODE, key);
            byte[] encrypted = cipher.doFinal(data.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(encrypted);
        } catch (Exception e) {
            throw new RuntimeException("公钥加密失败", e);
        }
    }

    /**
     * 使用私钥解密
     *
     * @param data       Base64 编码的密文
     * @param privateKey Base64 编码的私钥
     * @return 明文
     */
    public static String decryptByPrivateKey(String data, String privateKey) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(privateKey);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PrivateKey key = keyFactory.generatePrivate(keySpec);

            Cipher cipher = DECRYPT_CIPHER.get();
            cipher.init(Cipher.DECRYPT_MODE, key);
            byte[] decrypted = cipher.doFinal(Base64.getDecoder().decode(data));
            return new String(decrypted, StandardCharsets.UTF_8);
        } catch (Exception e) {
            throw new RuntimeException("私钥解密失败", e);
        }
    }

    /**
     * 使用私钥加密
     *
     * @param data       待加密数据
     * @param privateKey Base64 编码的私钥
     * @return Base64 编码的密文
     */
    public static String encryptByPrivateKey(String data, String privateKey) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(privateKey);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PrivateKey key = keyFactory.generatePrivate(keySpec);

            Cipher cipher = ENCRYPT_CIPHER.get();
            cipher.init(Cipher.ENCRYPT_MODE, key);
            byte[] encrypted = cipher.doFinal(data.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(encrypted);
        } catch (Exception e) {
            throw new RuntimeException("私钥加密失败", e);
        }
    }

    /**
     * 使用公钥解密
     *
     * @param data      Base64 编码的密文
     * @param publicKey Base64 编码的公钥
     * @return 明文
     */
    public static String decryptByPublicKey(String data, String publicKey) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(publicKey);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PublicKey key = keyFactory.generatePublic(keySpec);

            Cipher cipher = DECRYPT_CIPHER.get();
            cipher.init(Cipher.DECRYPT_MODE, key);
            byte[] decrypted = cipher.doFinal(Base64.getDecoder().decode(data));
            return new String(decrypted, StandardCharsets.UTF_8);
        } catch (Exception e) {
            throw new RuntimeException("公钥解密失败", e);
        }
    }

    /**
     * 使用私钥签名（SHA256withRSA）
     *
     * @param data       待签名的数据
     * @param privateKey Base64 编码的私钥
     * @return Base64 编码的签名
     */
    public static String sign(String data, String privateKey) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(privateKey);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PrivateKey key = keyFactory.generatePrivate(keySpec);

            java.security.Signature signature = SIGNATURE.get();
            signature.initSign(key);
            signature.update(data.getBytes(StandardCharsets.UTF_8));
            byte[] signBytes = signature.sign();
            return Base64.getEncoder().encodeToString(signBytes);
        } catch (Exception e) {
            throw new RuntimeException("签名失败", e);
        }
    }

    /**
     * 使用公钥验签（SHA256withRSA）
     *
     * @param data      原始数据
     * @param publicKey Base64 编码的公钥
     * @param sign      Base64 编码的签名
     * @return 验签结果
     */
    public static boolean verify(String data, String publicKey, String sign) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(publicKey);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PublicKey key = keyFactory.generatePublic(keySpec);

            java.security.Signature signature = SIGNATURE.get();
            signature.initVerify(key);
            signature.update(data.getBytes(StandardCharsets.UTF_8));
            return signature.verify(Base64.getDecoder().decode(sign));
        } catch (Exception e) {
            throw new RuntimeException("验签失败", e);
        }
    }

    /**
     * 分段加密（用于超长数据）
     * <p>
     * 注意：实际项目中，超长数据建议使用 AES + RSA 混合加密
     *
     * @param data      待加密数据
     * @param publicKey Base64 编码的公钥
     * @return Base64 编码的密文
     */
    public static String encryptByPublicKeySegment(String data, String publicKey) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(publicKey);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PublicKey key = keyFactory.generatePublic(keySpec);

            Cipher cipher = ENCRYPT_CIPHER.get();
            cipher.init(Cipher.ENCRYPT_MODE, key);

            byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8);
            int maxLength = 190; // 2048 位密钥 OAEP 模式最大加密长度
            byte[][] segments = splitArray(dataBytes, maxLength);

            StringBuilder result = new StringBuilder();
            for (int i = 0; i < segments.length; i++) {
                byte[] encrypted = cipher.doFinal(segments[i]);
                if (i > 0) {
                    result.append(",");
                }
                result.append(Base64.getEncoder().encodeToString(encrypted));
            }
            return result.toString();
        } catch (Exception e) {
            throw new RuntimeException("分段加密失败", e);
        }
    }

    /**
     * 分段解密
     *
     * @param data       加密的数据（逗号分隔的 Base64 字符串）
     * @param privateKey Base64 编码的私钥
     * @return 明文
     */
    public static String decryptByPrivateKeySegment(String data, String privateKey) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(privateKey);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance(RSA_ALGORITHM);
            PrivateKey key = keyFactory.generatePrivate(keySpec);

            Cipher cipher = DECRYPT_CIPHER.get();
            cipher.init(Cipher.DECRYPT_MODE, key);

            String[] segments = data.split(",");
            StringBuilder result = new StringBuilder();
            for (String segment : segments) {
                byte[] decrypted = cipher.doFinal(Base64.getDecoder().decode(segment));
                result.append(new String(decrypted, StandardCharsets.UTF_8));
            }
            return result.toString();
        } catch (Exception e) {
            throw new RuntimeException("分段解密失败", e);
        }
    }

    /**
     * 将字节数组分段
     */
    private static byte[][] splitArray(byte[] array, int chunkSize) {
        int chunks = (array.length + chunkSize - 1) / chunkSize;
        byte[][] result = new byte[chunks][];
        for (int i = 0; i < chunks; i++) {
            int start = i * chunkSize;
            int end = Math.min(start + chunkSize, array.length);
            result[i] = new byte[end - start];
            System.arraycopy(array, start, result[i], 0, end - start);
        }
        return result;
    }

    /**
     * 测试示例
     */
    public static void main(String[] args) {
        // 1. 生成密钥对
//        KeyPair keyPair = generateKeyPair();
//        String publicKey = getPublicKey(keyPair);
//        String privateKey = getPrivateKey(keyPair);
//
//        System.out.println("公钥: " + publicKey);
//        System.out.println("私钥: " + privateKey);
//        System.out.println();
//
//        // 2. 测试加密解密
//        String originalData = "Hello RSA 加密测试！123456789";
//        System.out.println("原始数据: " + originalData);
//
//        // 公钥加密，私钥解密
//        String encrypted = encryptByPublicKey(originalData, publicKey);
//        System.out.println("加密后: " + encrypted);
//
//        String decrypted = decryptByPrivateKey(encrypted, privateKey);
//        System.out.println("解密后: " + decrypted);
//        System.out.println();
//
//        // 3. 测试签名验签
//        String signature = sign(originalData, privateKey);
//        System.out.println("签名: " + signature);
//
//        boolean verified = verify(originalData, publicKey, signature);
//        System.out.println("验签结果: " + verified);
//
//        // 4. 测试私钥加密，公钥解密
//        String encryptedByPrivate = encryptByPrivateKey(originalData, privateKey);
//        System.out.println("私钥加密: " + encryptedByPrivate);
//
//        String decryptedByPublic = decryptByPublicKey(encryptedByPrivate, publicKey);
//        System.out.println("公钥解密: " + decryptedByPublic);
//
//        // 5. 测试分段加密（长数据）
//        String longData = "这是一段很长的数据，需要分段加密。".repeat(10);
//        String encryptedSegment = encryptByPublicKeySegment(longData, publicKey);
//        System.out.println("分段加密长度: " + encryptedSegment.length());
//
//        String decryptedSegment = decryptByPrivateKeySegment(encryptedSegment, privateKey);
//        System.out.println("分段解密结果: " + (longData.equals(decryptedSegment) ? "成功" : "失败"));
    }
}