package com.lesingle.edu.common.util;

import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicLong;

/**
 * RSA 加密工具类
 * 支持密钥对生成、加密解密、密钥版本管理
 */
@Slf4j
@Component
public class RsaEncryptionUtil {

    private static final String ALGORITHM = "RSA";
    private static final int KEY_SIZE = 2048;

    private volatile KeyPair keyPair;
    private volatile String currentPublicKey;
    private final AtomicLong keyVersion = new AtomicLong(0);

    /**
     * 初始化时生成密钥对
     */
    @PostConstruct
    public void init() throws Exception {
        generateNewKeyPair();
        log.info("RSA 加密工具初始化完成，当前密钥版本: {}", keyVersion.get());
    }

    /**
     * 生成新的 RSA 密钥对
     * 每月定时执行，也可手动调用
     */
    public synchronized void generateNewKeyPair() throws Exception {
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM);
            generator.initialize(KEY_SIZE);
            this.keyPair = generator.generateKeyPair();
            this.currentPublicKey = Base64.getEncoder().encodeToString(
                    keyPair.getPublic().getEncoded()
            );
            keyVersion.incrementAndGet();
            log.info("RSA 密钥对已更换，新版本号: {}，密钥长度: {} 位", keyVersion.get(), KEY_SIZE);
        } catch (Exception e) {
            log.error("生成 RSA 密钥对失败", e);
            throw e;
        }
    }

    /**
     * 获取当前公钥（Base64 编码）
     */
    public String getPublicKey() {
        return currentPublicKey;
    }

    /**
     * 获取当前密钥版本号
     */
    public long getKeyVersion() {
        return keyVersion.get();
    }

    /**
     * 使用私钥解密数据
     *
     * @param encryptedData Base64 编码的加密数据
     * @return 解密后的明文
     */
    public String decrypt(String encryptedData) throws Exception {
        if (encryptedData == null || encryptedData.isEmpty()) {
            throw new IllegalArgumentException("加密数据不能为空");
        }

        try {
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, keyPair.getPrivate());
            byte[] decrypted = cipher.doFinal(
                    Base64.getDecoder().decode(encryptedData)
            );
            return new String(decrypted, StandardCharsets.UTF_8);
        } catch (Exception e) {
            log.error("RSA 解密失败", e);
            throw new Exception("密码解密失败，请重试", e);
        }
    }

    /**
     * 使用公钥加密数据（主要用于测试）
     *
     * @param plainText 明文
     * @param publicKeyBase64 Base64 编码的公钥
     * @return Base64 编码的密文
     */
    public String encrypt(String plainText, String publicKeyBase64) throws Exception {
        if (plainText == null || plainText.isEmpty()) {
            throw new IllegalArgumentException("明文不能为空");
        }

        try {
            byte[] keyBytes = Base64.getDecoder().decode(publicKeyBase64);
            java.security.spec.X509EncodedKeySpec spec = new java.security.spec.X509EncodedKeySpec(keyBytes);
            java.security.KeyFactory keyFactory = java.security.KeyFactory.getInstance(ALGORITHM);
            PublicKey publicKey = keyFactory.generatePublic(spec);

            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            byte[] encrypted = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(encrypted);
        } catch (Exception e) {
            log.error("RSA 加密失败", e);
            throw new Exception("密码加密失败", e);
        }
    }

    /**
     * 使用当前公钥加密数据（用于测试）
     */
    public String encryptWithCurrentKey(String plainText) throws Exception {
        return encrypt(plainText, currentPublicKey);
    }
}
