package com.sunyard.cisd.device.tool;

import org.bouncycastle.asn1.*;
import org.bouncycastle.asn1.cms.*;
import org.bouncycastle.cms.*;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
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[] sm3Digest = calculateSM3Digest(plaintext);

        ASN1EncodableVector signedAttrsVector = new ASN1EncodableVector();
        Attribute messageDigestAttr = new Attribute(CMSAttributes.messageDigest,
            new DERSet(new DEROctetString(sm3Digest)));
        signedAttrsVector.add(messageDigestAttr);

        Attribute contentTypeAttr = new Attribute(CMSAttributes.contentType,
            new DERSet(CMSObjectIdentifiers.data));
        signedAttrsVector.add(contentTypeAttr);

        DERSet signedAttrs = new DERSet(signedAttrsVector);
        byte[] signedAttrsEncoded = signedAttrs.getEncoded(ASN1Encoding.DER);

        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM, PROVIDER);
        signature.initSign(signingKey);
        signature.update(signedAttrsEncoded);
        byte[] signatureValue = signature.sign();

        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(plaintext);

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

        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        baos.write((byte)(encryptedKey.length >> 8));
        baos.write((byte)encryptedKey.length);
        baos.write(encryptedKey);
        baos.write(iv);
        baos.write(encryptedData);

        byte[] encryptedContent = baos.toByteArray();

        byte[] finalData = new byte[4 + signatureValue.length + encryptedContent.length];
        byte[] sigLen = new byte[]{
            (byte)(signatureValue.length >> 24),
            (byte)(signatureValue.length >> 16),
            (byte)(signatureValue.length >> 8),
            (byte)signatureValue.length
        };
        System.arraycopy(sigLen, 0, finalData, 0, 4);
        System.arraycopy(signatureValue, 0, finalData, 4, signatureValue.length);
        System.arraycopy(encryptedContent, 0, finalData, 4 + signatureValue.length, encryptedContent.length);

        return finalData;
    }

    private static byte[] calculateSM3Digest(byte[] data) throws Exception {
        MessageDigest digest = MessageDigest.getInstance("SM3", PROVIDER);
        return digest.digest(data);
    }

    public static class VerifyResult {
        public final boolean signatureValid;
        public final byte[] messageDigest;
        public final String digestAlgorithm;

        public VerifyResult(boolean signatureValid, byte[] messageDigest, String digestAlgorithm) {
            this.signatureValid = signatureValid;
            this.messageDigest = messageDigest;
            this.digestAlgorithm = digestAlgorithm;
        }
    }

    public static boolean verifySignatureWithAttrs(byte[] signatureValue, byte[] signedAttrsEncoded, PublicKey publicKey) throws Exception {
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM, PROVIDER);
        signature.initVerify(publicKey);
        signature.update(signedAttrsEncoded);
        return signature.verify(signatureValue);
    }

    public static byte[] buildSignedAttrs(byte[] messageDigest) {
        ASN1EncodableVector signedAttrsVector = new ASN1EncodableVector();
        
        Attribute contentTypeAttr = new Attribute(CMSAttributes.contentType,
            new DERSet(CMSObjectIdentifiers.data));
        signedAttrsVector.add(contentTypeAttr);
        
        Attribute messageDigestAttr = new Attribute(CMSAttributes.messageDigest,
            new DERSet(new DEROctetString(messageDigest)));
        signedAttrsVector.add(messageDigestAttr);

        DERSet signedAttrs = new DERSet(signedAttrsVector);
        try {
            return signedAttrs.getEncoded(ASN1Encoding.DER);
        } catch (Exception e) {
            return null;
        }
    }

    public static byte[] extractMessageDigestFromAttrs(byte[] signedAttrsEncoded) throws Exception {
        ASN1InputStream asn1InputStream = new ASN1InputStream(signedAttrsEncoded);
        ASN1Set signedAttrs = ASN1Set.getInstance(asn1InputStream.readObject());

        for (int i = 0; i < signedAttrs.size(); i++) {
            Attribute attr = Attribute.getInstance(signedAttrs.getObjectAt(i));
            if (attr.getAttrType().equals(CMSAttributes.messageDigest)) {
                return ASN1OctetString.getInstance(
                    attr.getAttrValues().getObjectAt(0)).getOctets();
            }
        }
        return null;
    }

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

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

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

        byte[] signatureValue = new byte[sigLen];
        System.arraycopy(envelopeData, 4, signatureValue, 0, sigLen);

        byte[] encryptedContent = new byte[envelopeData.length - 4 - sigLen];
        System.arraycopy(envelopeData, 4 + sigLen, encryptedContent, 0, encryptedContent.length);

        int keyLen = ((encryptedContent[0] & 0xFF) << 8) | (encryptedContent[1] & 0xFF);
        byte[] encryptedKey = new byte[keyLen];
        System.arraycopy(encryptedContent, 2, encryptedKey, 0, keyLen);

        byte[] iv = new byte[12];
        System.arraycopy(encryptedContent, 2 + keyLen, iv, 0, 12);

        byte[] encryptedData = new byte[encryptedContent.length - 2 - keyLen - 12];
        System.arraycopy(encryptedContent, 2 + keyLen + 12, 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[] plaintext = sm4Cipher.doFinal(encryptedData);

        byte[] messageDigest = calculateSM3Digest(plaintext);
        byte[] signedAttrsEncoded = buildSignedAttrs(messageDigest);

        boolean signatureValid = verifySignatureWithAttrs(signatureValue, signedAttrsEncoded, signingPublicKey);

        return new DecryptedResult(plaintext, signatureValue, signatureValid, messageDigest);
    }

    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("原文SM3哈希: " + bytesToHex(calculateSM3Digest(plaintext)));
        System.out.println();

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

        System.out.println("=== 不解密验证签名 (使用 signedAttrs) ===");
        byte[] messageDigest = calculateSM3Digest(plaintext);
        byte[] signedAttrsEncoded = buildSignedAttrs(messageDigest);
        
        int sigLen = ((envelope[0] & 0xFF) << 24) |
                    ((envelope[1] & 0xFF) << 16) |
                    ((envelope[2] & 0xFF) << 8) |
                    (envelope[3] & 0xFF);
        byte[] signatureValue = new byte[sigLen];
        System.arraycopy(envelope, 4, signatureValue, 0, sigLen);
        
        boolean sigValid = verifySignatureWithAttrs(signatureValue, signedAttrsEncoded, signingKeyPair.publicKey);
        System.out.println("签名有效: " + sigValid);
        System.out.println("提取的消息摘要: " + bytesToHex(extractMessageDigestFromAttrs(signedAttrsEncoded)));
        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.messageDigest));
        System.out.println();

        if (sigValid && result.signatureValid) {
            System.out.println("=== 全部验证通过 ===");
        } else {
            System.out.println("=== 验证失败 ===");
        }
    }
}