package com.sunyard.cisd.device.tool;

import org.bouncycastle.jce.ECNamedCurveTable;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.jce.spec.ECParameterSpec;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class SM2Util {
    private static final String ALGORITHM = "EC";
    private static final String PROVIDER = "BC";
    private static final String CURVE_NAME = "sm2p256v1";
    private static final String SIGNATURE_ALGORITHM = "SM3withSM2";
    
    static {
        Security.addProvider(new BouncyCastleProvider());
    }
    
    public static byte[] sign(byte[] data, PrivateKey privateKey) throws Exception {
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM, PROVIDER);
        signature.initSign(privateKey);
        signature.update(data);
        return signature.sign();
    }
    
    public static boolean verify(byte[] data, byte[] sign, PublicKey publicKey) throws Exception {
        Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM, PROVIDER);
        signature.initVerify(publicKey);
        signature.update(data);
        return signature.verify(sign);
    }
    
    public static byte[] sign(byte[] data, String privateKeyFilePath) throws Exception {
        PrivateKey privateKey = loadPrivateKey(privateKeyFilePath);
        return sign(data, privateKey);
    }
    
    public static boolean verify(byte[] data, byte[] sign, String publicKeyFilePath) throws Exception {
        PublicKey publicKey = loadPublicKey(publicKeyFilePath);
        return verify(data, sign, publicKey);
    }
    
    public static PrivateKey loadPrivateKey(String filePath) throws Exception {
        InputStream inputStream = SM2Util.class.getResourceAsStream(filePath);
        if (inputStream == null) {
            throw new IllegalArgumentException("Private key file not found: " + filePath);
        }
        
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
        StringBuilder keyBuilder = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            if (!line.startsWith("-----")) {
                keyBuilder.append(line.trim());
            }
        }
        reader.close();
        
        String keyBase64 = keyBuilder.toString();
        byte[] keyBytes = Base64.getDecoder().decode(keyBase64);
        
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM, PROVIDER);
        return keyFactory.generatePrivate(keySpec);
    }
    
    public static PublicKey loadPublicKey(String filePath) throws Exception {
        InputStream inputStream = SM2Util.class.getResourceAsStream(filePath);
        if (inputStream == null) {
            throw new IllegalArgumentException("Public key file not found: " + filePath);
        }
        
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
        StringBuilder keyBuilder = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            if (!line.startsWith("-----")) {
                keyBuilder.append(line.trim());
            }
        }
        reader.close();
        
        String keyBase64 = keyBuilder.toString();
        byte[] keyBytes = Base64.getDecoder().decode(keyBase64);
        
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM, PROVIDER);
        return keyFactory.generatePublic(keySpec);
    }
    
    public static KeyPair generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM, PROVIDER);
        ECParameterSpec ecSpec = ECNamedCurveTable.getParameterSpec(CURVE_NAME);
        keyPairGenerator.initialize(ecSpec, new SecureRandom());
        return keyPairGenerator.generateKeyPair();
    }
    
    public static String privateKeyToPEM(PrivateKey privateKey) {
        byte[] encoded = privateKey.getEncoded();
        String base64 = Base64.getEncoder().encodeToString(encoded);
        StringBuilder pem = new StringBuilder();
        pem.append("-----BEGIN PRIVATE KEY-----\n");
        int index = 0;
        while (index < base64.length()) {
            pem.append(base64.substring(index, Math.min(index + 64, base64.length())));
            pem.append("\n");
            index += 64;
        }
        pem.append("-----END PRIVATE KEY-----\n");
        return pem.toString();
    }
    
    public static String publicKeyToPEM(PublicKey publicKey) {
        byte[] encoded = publicKey.getEncoded();
        String base64 = Base64.getEncoder().encodeToString(encoded);
        StringBuilder pem = new StringBuilder();
        pem.append("-----BEGIN PUBLIC KEY-----\n");
        int index = 0;
        while (index < base64.length()) {
            pem.append(base64.substring(index, Math.min(index + 64, base64.length())));
            pem.append("\n");
            index += 64;
        }
        pem.append("-----END PUBLIC KEY-----\n");
        return pem.toString();
    }
}