package com.sunyard.cisd.device.tool;

import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder;
import org.bouncycastle.cms.*;
import org.bouncycastle.cms.jcajce.*;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;
import org.bouncycastle.util.encoders.Hex;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.ByteArrayOutputStream;
import java.security.*;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.*;

public class SM2SignedEnvelopeUtil {

    private static final String PROVIDER = "BC";
    private static final String CURVE_NAME = "sm2p256v1";
    private static final String KEY_ALGORITHM = "EC";
    private static final String SIGNATURE_ALGORITHM = "SM3withSM2";
    private static final String SM4_ALGORITHM = "SM4/GCM/NoPadding";

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    public static class SM2KeyPair {
        public final PrivateKey privateKey;
        public final PublicKey publicKey;
        public final byte[] privateKeyBytes;
        public final byte[] publicKeyBytes;

        public SM2KeyPair(PrivateKey privateKey, PublicKey publicKey) {
            this.privateKey = privateKey;
            this.publicKey = publicKey;
            this.privateKeyBytes = privateKey.getEncoded();
            this.publicKeyBytes = publicKey.getEncoded();
        }

        public String getPrivateKeyBase64() {
            return Base64.getEncoder().encodeToString(privateKeyBytes);
        }

        public String getPublicKeyBase64() {
            return Base64.getEncoder().encodeToString(publicKeyBytes);
        }

        public String getPrivateKeyHex() {
            return Hex.toHexString(privateKeyBytes);
        }

        public String getPublicKeyHex() {
            return Hex.toHexString(publicKeyBytes);
        }
    }

    public static SM2KeyPair generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM, PROVIDER);
        ECGenParameterSpec ecSpec = new ECGenParameterSpec(CURVE_NAME);
        keyPairGenerator.initialize(ecSpec, new SecureRandom());
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        return new SM2KeyPair(keyPair.getPrivate(), keyPair.getPublic());
    }

    public static PrivateKey loadPrivateKey(byte[] privateKeyBytes) throws Exception {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM, PROVIDER);
        return keyFactory.generatePrivate(keySpec);
    }

    public static PublicKey loadPublicKey(byte[] publicKeyBytes) throws Exception {
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM, PROVIDER);
        return keyFactory.generatePublic(keySpec);
    }

    public static byte[] createSignedAndEnvelopedData(byte[] plaintext, PrivateKey signingKey, PublicKey encryptionPublicKey) throws Exception {
        byte[] signature = signData(plaintext, signingKey);

        ByteArrayOutputStream signedContentBaos = new ByteArrayOutputStream();
        signedContentBaos.write(signature.length >> 24 & 0xFF);
        signedContentBaos.write(signature.length >> 16 & 0xFF);
        signedContentBaos.write(signature.length >> 8 & 0xFF);
        signedContentBaos.write(signature.length & 0xFF);
        signedContentBaos.write(signature);
        signedContentBaos.write(plaintext);
        byte[] signedContent = signedContentBaos.toByteArray();

        byte[] encryptedContent = encryptData(signedContent, encryptionPublicKey);

        return encryptedContent;
    }

    private static byte[] signData(byte[] data, PrivateKey privateKey) throws Exception {
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM, PROVIDER);
        signature.initSign(privateKey);
        signature.update(data);
        return signature.sign();
    }

    private static byte[] encryptData(byte[] data, PublicKey publicKey) throws Exception {
        SecureRandom random = new SecureRandom();
        byte[] iv = new byte[12];
        random.nextBytes(iv);

        KeyGenerator keyGen = KeyGenerator.getInstance("SM4", PROVIDER);
        keyGen.init(128, random);
        SecretKey sm4Key = keyGen.generateKey();

        Cipher sm4Cipher = Cipher.getInstance(SM4_ALGORITHM, PROVIDER);
        GCMParameterSpec gcmSpec = new GCMParameterSpec(128, iv);
        sm4Cipher.init(Cipher.ENCRYPT_MODE, sm4Key, gcmSpec);
        byte[] encryptedData = sm4Cipher.doFinal(data);

        Cipher ecCipher = Cipher.getInstance("ECIES", PROVIDER);
        ecCipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encryptedKey = ecCipher.doFinal(sm4Key.getEncoded());

        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        byte[] keyLenBytes = new byte[]{
            (byte)(encryptedKey.length >> 8),
            (byte)encryptedKey.length
        };
        baos.write(keyLenBytes[0]);
        baos.write(keyLenBytes[1]);
        baos.write(encryptedKey);
        baos.write(iv);
        baos.write(encryptedData);

        return baos.toByteArray();
    }

    public static class DecryptedResult {
        public final byte[] plaintext;
        public final byte[] signature;
        public final boolean signatureValid;

        public DecryptedResult(byte[] plaintext, byte[] signature, boolean signatureValid) {
            this.plaintext = plaintext;
            this.signature = signature;
            this.signatureValid = signatureValid;
        }
    }

    public static DecryptedResult openAndVerifySignedEnvelopedData(byte[] envelopeData, PrivateKey decryptionKey, PublicKey signingPublicKey) throws Exception {
        int offset = 0;
        int encryptedKeyLen = ((envelopeData[offset] & 0xFF) << 8) | (envelopeData[offset + 1] & 0xFF);
        offset += 2;

        byte[] encryptedKey = new byte[encryptedKeyLen];
        System.arraycopy(envelopeData, offset, encryptedKey, 0, encryptedKeyLen);
        offset += encryptedKeyLen;

        byte[] iv = new byte[12];
        System.arraycopy(envelopeData, offset, iv, 0, 12);
        offset += 12;

        byte[] encryptedData = new byte[envelopeData.length - offset];
        System.arraycopy(envelopeData, offset, encryptedData, 0, encryptedData.length);

        Cipher ecCipher = Cipher.getInstance("ECIES", PROVIDER);
        ecCipher.init(Cipher.DECRYPT_MODE, decryptionKey);
        byte[] sm4KeyBytes = ecCipher.doFinal(encryptedKey);

        SecretKeySpec sm4Key = new SecretKeySpec(sm4KeyBytes, "SM4");

        Cipher sm4Cipher = Cipher.getInstance(SM4_ALGORITHM, PROVIDER);
        GCMParameterSpec gcmSpec = new GCMParameterSpec(128, iv);
        sm4Cipher.init(Cipher.DECRYPT_MODE, sm4Key, gcmSpec);
        byte[] signedContent = sm4Cipher.doFinal(encryptedData);

        int signatureLength = ((signedContent[0] & 0xFF) << 24) |
                             ((signedContent[1] & 0xFF) << 16) |
                             ((signedContent[2] & 0xFF) << 8) |
                             (signedContent[3] & 0xFF);

        byte[] signature = new byte[signatureLength];
        byte[] plaintext = new byte[signedContent.length - 4 - signatureLength];

        System.arraycopy(signedContent, 4, signature, 0, signatureLength);
        System.arraycopy(signedContent, 4 + signatureLength, plaintext, 0, plaintext.length);

        boolean signatureValid = verifySignature(plaintext, signature, signingPublicKey);

        return new DecryptedResult(plaintext, signature, signatureValid);
    }

    private static boolean verifySignature(byte[] data, byte[] signature, PublicKey publicKey) throws Exception {
        Signature sig = Signature.getInstance(SIGNATURE_ALGORITHM, PROVIDER);
        sig.initVerify(publicKey);
        sig.update(data);
        return sig.verify(signature);
    }

    public static String toBase64(byte[] data) {
        return Base64.getEncoder().encodeToString(data);
    }

    public static byte[] fromBase64(String base64) {
        return Base64.getDecoder().decode(base64);
    }

    public static String bytesToHex(byte[] data) {
        return Hex.toHexString(data);
    }

    public static byte[] hexToBytes(String hex) {
        return Hex.decode(hex);
    }

    public static void main(String[] args) throws Exception {
        System.out.println("=== SM2 带签名数字信封测试 ===");
        System.out.println();

        SM2KeyPair signingKeyPair = generateKeyPair();
        SM2KeyPair encryptionKeyPair = generateKeyPair();

        System.out.println("签名密钥对已生成");
        System.out.println("加密密钥对已生成");
        System.out.println();

        String testData = "这是测试数据 - Hello World! 你好世界！";
        byte[] plaintext = testData.getBytes("UTF-8");

        System.out.println("原文: " + testData);
        System.out.println("原文长度: " + plaintext.length + " bytes");
        System.out.println();

        System.out.println("创建带签名的数字信封...");
        byte[] envelope = createSignedAndEnvelopedData(plaintext, signingKeyPair.privateKey, encryptionKeyPair.publicKey);
        System.out.println("信封已创建，长度: " + envelope.length + " bytes");
        System.out.println("信封Base64: " + toBase64(envelope).substring(0, 80) + "...");
        System.out.println();

        System.out.println("打开并验证数字信封...");
        DecryptedResult result = openAndVerifySignedEnvelopedData(envelope, encryptionKeyPair.privateKey, signingKeyPair.publicKey);

        System.out.println("解密后原文: " + new String(result.plaintext, "UTF-8"));
        System.out.println("签名有效: " + result.signatureValid);
        System.out.println("签名: " + bytesToHex(result.signature));
        System.out.println();

        if (result.signatureValid) {
            System.out.println("=== 验签成功 ===");
        } else {
            System.out.println("=== 验签失败 ===");
        }
    }
}