package com.project.zcustom.tools;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.HexUtil;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.SmUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.hutool.crypto.asymmetric.SM2;
import cn.hutool.crypto.digest.SM3;
import cn.hutool.crypto.symmetric.SM4;
import com.project.zcustom.controller.core.StringUtils;
import org.bouncycastle.crypto.engines.SM2Engine;

import java.security.KeyPair;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

public class EncryptionUtil {

    /**
     * 缓存map，用于缓存各个国密对象
     */
    private static Map<String, Object> SMMap = new HashMap<>();

    private static final int cacheLimit = 100;

    /**
     * 缓存SM对象
     * @param key   键
     * @param obj   对象
     */
    private static void putSMMap(String key, Object obj) {
        if(StringUtils.isEmpty(key)) {
            return;
        }
        if(SMMap.size() > cacheLimit) {
            SMMap.clear();
        } else if(SMMap.containsKey(key)) {
            return;
        }
        SMMap.put(key, obj);
    }

    /**
     * 取出SM对象
     * @param key   键
     * @param type  类型
     * @return      对象
     */
    private static Object getSMMap(String key, int type) {
        if(StringUtils.isEmpty(key)) {
            return null;
        }
        if(!SMMap.containsKey(key)) {
            return null;
        }
        if(Objects.isNull(SMMap.get(key))) {
            return null;
        }
        Object obj = SMMap.get(key);
        if(2 == type && obj instanceof SM2) {
            return obj;
        }
        if(3 == type && obj instanceof SM3) {
            return obj;
        }
        if(4 == type && obj instanceof SM4) {
            return obj;
        }
        return null;
    }

    /**
     * 将十六进制字符串解码为byte[]
     * @param str   字符串
     * @return      字节数组
     */
    private static byte[] decodeHex(String str) {
        return HexUtil.decodeHex(str);
    }

    /**
     * 将字节数组转换为十六进制字符串
     * @param bytes     字节数组
     * @return          字符串
     */
    private static String encodeHexStr(byte[] bytes) {
        return HexUtil.encodeHexStr(bytes);
    }

    /**
     * 将字符串转换为十六进制字符串，结果为小写，默认编码是UTF-8
     * @param str       字符串
     * @return          十六进制字符串
     */
    private static String encodeHexStr(String str) {
        return HexUtil.encodeHexStr(str);
    }

    /**
     * 获取SM2对象
     * @param privateKey    私钥
     * @param publicKey     公钥
     * @return      非对称加密对象
     */
    private static SM2 getSM2(String publicKey, String privateKey) {
        if(StringUtils.isEmpty(publicKey) && StringUtils.isEmpty(privateKey)) {
            throw new RuntimeException("公钥私钥不可都为空");
        }
        Object cacheObj = getSMMap(privateKey, 2);
        if(!Objects.isNull(cacheObj)) {
            return (SM2)cacheObj;
        }
        cacheObj = getSMMap(publicKey, 2);
        if(!Objects.isNull(cacheObj)) {
            return (SM2)cacheObj;
        }
        SM2 sm2 = SmUtil.sm2(
                StringUtils.isEmpty(privateKey) ? null : decodeHex(privateKey),
                StringUtils.isEmpty(publicKey) ? null: decodeHex(publicKey));
        sm2.usePlainEncoding();
        sm2.setMode(SM2Engine.Mode.C1C2C3);

        //尝试缓存SM对象
        putSMMap(privateKey, sm2);
        putSMMap(publicKey, sm2);

        return sm2;
    }

    /**
     * 获取SM3对象
     * @param salt  盐值
     * @return      摘要对象
     */
    private static SM3 getSM3(String salt) {
        Object cacheObj = getSMMap(salt, 3);
        if(!Objects.isNull(cacheObj)) {
            return (SM3)cacheObj;
        }
        SM3 sm3 = SmUtil.sm3WithSalt(decodeHex(salt));
        putSMMap(salt, sm3);
        return sm3;
    }

    /**
     * 获取SM4对象
     * @param key   密钥
     * @return      对称加密对象
     */
    private static SM4 getSM4(String key) {
        Object cacheObj = getSMMap(key, 4);
        if(!Objects.isNull(cacheObj)) {
            return (SM4)cacheObj;
        }
        SM4 sm4 = SmUtil.sm4(decodeHex(key));
        putSMMap(key, sm4);
        return sm4;
    }

    /**
     * 生成密钥对
     * @return  密钥对
     */
    public static KeyPair generateKey() {
        return SecureUtil.generateKeyPair("SM2");
    }

    /**
     * 获取私钥
     * @param keyPair   密钥对
     * @return  私钥
     */
    public static String getPrivateKey(KeyPair keyPair) {
        return HexUtil.encodeHexStr(keyPair.getPrivate().getEncoded());
    }

    /**
     * 获取私钥D
     * @param privateKey   私钥
     * @return              D
     */
    public static String getPrivateKeyD(String privateKey) {
        try {
            return encodeHexStr(getSM2(null, privateKey).getD());
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("获取私钥D值失败");
        }
    }

    /**
     * 获取公钥
     * @param keyPair 密钥对
     * @return  公钥
     */
    public static String getPublicKey(KeyPair keyPair) {
        return HexUtil.encodeHexStr(keyPair.getPublic().getEncoded());
    }

    /**
     * 获取公钥Q
     * @param publicKey 公钥
     * @return          Q
     */
    public static String getPublicKeyQ(String publicKey) {
        try {
            return encodeHexStr(getSM2(publicKey, null).getQ(false));
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("获取公钥Q值失败");
        }
    }

    /**
     * 获取非对称加密结果
     * @param publicKey     公钥
     * @param data          原始数据
     * @return              加密数据
     */
    public static String getAsymmetricEncryptionData(String publicKey, String data) {
        try {
            return getSM2(publicKey, null).encryptHex(data, KeyType.PublicKey);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("非对称加密失败");
        }
    }

    /**
     * 获取非对称解密结果
     * @param privateKey    私钥
     * @param data          加密数据
     * @return              原始数据
     */
    public static String getAsymmetricDecryptionData(String privateKey, String data) {
        try {
            return getSM2(null, privateKey).decryptStr(data, KeyType.PrivateKey);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("非对称解密失败");
        }
    }

    /**
     * 获取签名
     * @param privateKey    私钥
     * @param data          原始数据
     * @return              签名
     */
    public static String getSign(String privateKey, String data) {
        try {
            return getSM2(null, privateKey).signHex(encodeHexStr(data));
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("获取签名失败");
        }
    }

    /**
     * 验证签名
     * @param publicKey 公钥
     * @param data      原始数据
     * @return          签名
     */
    public static Boolean verifySign(String publicKey, String data, String sign) {
        try {
            return getSM2(publicKey, null).verifyHex(encodeHexStr(data), sign);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("验证签名失败");
        }
    }

    /**
     * 获取摘要值
     * @param salt  盐值
     * @param data  原始数据
     * @return      摘要值
     */
    public static String getAbstractData(String salt, String data) {
        try {
            return getSM3(salt).digestHex(data);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("获取摘要值失败");
        }
    }

    /**
     * 获取对称加密结果
     * @param key   密钥
     * @param data  原始数据
     * @return      加密数据
     */
    public static String getSymmetricalEncryptionData(String key, String data) {
        try {
            return getSM4(key).encryptHex(data);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("对称加密失败");
        }

    }

    /**
     * 获取对称解密结果
     * @param key   密钥
     * @param data  加密数据
     * @return      原始数据
     */
    public static String getSymmetricalDecryptionData(String key, String data) {
        try {
            return getSM4(key).decryptStr(data, CharsetUtil.CHARSET_UTF_8);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("对称解密失败");
        }
    }
}
