package com.sunyard.chsm.pool;

import com.sunyard.chsm.enums.Padding;
import com.sunyard.chsm.sdf.SdfApiService;
import com.sunyard.chsm.sdf.context.AlgId;
import com.sunyard.chsm.sdf.model.EccCipher;
import com.sunyard.chsm.sdf.model.EccKey;
import com.sunyard.chsm.sdf.model.EccPriKey;
import com.sunyard.chsm.sdf.model.EccPubKey;
import com.sunyard.chsm.sdf.model.EccSignature;
import com.sunyard.chsm.sdf.util.PaddingUtil;
import com.sunyard.chsm.utils.gm.BCECUtils;
import com.sunyard.chsm.utils.gm.BCSM3Utils;
import com.sunyard.chsm.utils.gm.SM2PreprocessSigner;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.params.ParametersWithID;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.function.Function;

/**
 * @author liulu
 * @since 2024/12/10
 */
@Slf4j
@Service
public class LoadBalancedSdfApiService implements SdfApiService {

    @Resource
    private DeviceManager deviceManager;


    private <T> T apply(Function<TMKContext, T> fun) {
        TMKContext context = deviceManager.chooseOne();
        try {
            return fun.apply(context);
        } finally {
            deviceManager.returnContextToPool(context);
        }
    }

    @Override
    public byte[] generateRandom(int len) {
        return apply(s -> s.getSdfApiAdapter().generateRandom(s.getSessionHandle(), len));
    }

    @Override
    public byte[] symEncrypt(AlgId alg, Padding padding, byte[] key, byte[] iv, byte[] data) {
        byte[] paddingData = PaddingUtil.padding(padding, data);
        return apply(s -> {
            String keyHandle = s.getSdfApiAdapter().importKey(s.getSessionHandle(), key);
            byte[] encrypt = s.getSdfApiAdapter().symEncrypt(s.getSessionHandle(), keyHandle, alg, iv, paddingData);
            s.getSdfApiAdapter().destroyKey(s.getSessionHandle(), keyHandle);
            return encrypt;
        });
    }

    @Override
    public byte[] symDecrypt(AlgId alg, Padding padding, byte[] key, byte[] iv, byte[] data) {
        byte[] decrypt = apply(s -> {
            String keyHandle = s.getSdfApiAdapter().importKey(s.getSessionHandle(), key);
            byte[] d = s.getSdfApiAdapter().symDecrypt(s.getSessionHandle(), keyHandle, alg, iv, data);
            s.getSdfApiAdapter().destroyKey(s.getSessionHandle(), keyHandle);
            return d;
        });
        return PaddingUtil.unpadding(padding, decrypt);
    }

    @Override
    public EccKey genKeyPairEcc() {
        return apply(s -> s.getSdfApiAdapter().generateKeyPairECC(s.getSessionHandle(), AlgId.SGD_SM2_1));
    }

    @Override
    public EccSignature externalSignECC(byte[] privateKey, byte[] pucData) {
        return externalSignWithIdECC(privateKey, pucData, new byte[0]);
    }

    @Override
    public EccSignature externalSignWithIdECC(byte[] privateKey, byte[] pucData, byte[] userId) {
        EccPriKey sdfEccPrivateKey = EccPriKey.fromBytes(privateKey);
        SM2PreprocessSigner signer = new SM2PreprocessSigner();
        CipherParameters keyParameters = BCECUtils.createECPrivateKeyParameters(sdfEccPrivateKey.getD());
        if (userId != null && userId.length > 0) {
            keyParameters = new ParametersWithID(keyParameters, userId);
        }
        signer.init(true, keyParameters);
        byte[] preHash = signer.preprocess(pucData);

        return apply(s -> s.getSdfApiAdapter().externalSignECC(s.getSessionHandle(), sdfEccPrivateKey, preHash));
    }

    @Override
    public boolean externalVerifyWithIdECC(byte[] publicKey, byte[] pubData, byte[] signData, byte[] userId) {
        EccPubKey eccPublicKey = EccPubKey.fromBytes(publicKey);
        CipherParameters parameters = BCECUtils.createECPublicKeyParameters(eccPublicKey.getX(), eccPublicKey.getY());
        if (userId != null && userId.length > 0) {
            parameters = new ParametersWithID(parameters, userId);
        }
        SM2PreprocessSigner signer = new SM2PreprocessSigner();
        signer.init(false, parameters);
        byte[] preHash = signer.preprocess(pubData);
        EccSignature eccSignature = EccSignature.fromBytes(signData);

        return apply(s -> s.getSdfApiAdapter().externalVerifyECC(s.getSessionHandle(),
                eccPublicKey, preHash, eccSignature));
    }

    @Override
    public boolean externalVerifyECC(byte[] publicKey, byte[] pubData, byte[] signData) {
        return externalVerifyWithIdECC(publicKey, pubData, signData, new byte[0]);
    }

    @Override
    public EccCipher externalEncryptECC(byte[] publicKey, byte[] pucData) {
        EccPubKey eccPublicKey = EccPubKey.fromBytes(publicKey);
        return apply(s -> s.getSdfApiAdapter().externalEncryptECC(s.getSessionHandle(), eccPublicKey, pucData));
    }

    @Override
    public byte[] externalDecryptECC(byte[] privateKey, byte[] encData) {
        EccPriKey sdfEccPrivateKey = EccPriKey.fromBytes(privateKey);
        EccCipher eccCipher = EccCipher.fromBytes(encData);
        return apply(s -> s.getSdfApiAdapter().externalDecryptECC(s.getSessionHandle(), sdfEccPrivateKey, eccCipher));
    }

    @Override
    public byte[] calculateMAC(AlgId algId, Padding padding, byte[] symKey, byte[] pucIv, byte[] pucData) {
        byte[] pad = PaddingUtil.padding(padding, pucData);
        return apply(s -> {
            String keyHandle = s.getSdfApiAdapter().importKey(s.getSessionHandle(), symKey);
            byte[] mac = s.getSdfApiAdapter().calculateMAC(s.getSessionHandle(), keyHandle, algId, pucIv, pad);
            s.getSdfApiAdapter().destroyKey(s.getSessionHandle(), keyHandle);
            return mac;
        });
    }

    @Override
    public byte[] hmac(byte[] key, byte[] srcData) {
        return BCSM3Utils.hmac(key, srcData);
    }

    @Override
    public byte[] hash(byte[] pucData) {
        return apply(s -> {
            String newSession = s.getSdfApiAdapter().openSession(s.getDeviceHandle());
            s.getSdfApiAdapter().hashInit(newSession, AlgId.SGD_SM3, null, null);
            s.getSdfApiAdapter().hashUpdate(newSession, pucData);
            byte[] hash = s.getSdfApiAdapter().hashFinish(newSession);
            s.getSdfApiAdapter().closeSession(newSession);
            return hash;
        });
    }

    @Override
    public byte[] encryptByTMK(byte[] data) {
        byte[] pad = PaddingUtil.PKCS7Padding(data);
        return apply(s -> s.getSdfApiAdapter().symEncrypt(s.getSessionHandle(), s.getKeyHandle(), AlgId.SGD_SM4_CBC, new byte[8], pad));
    }

    @Override
    public byte[] decryptByTMK(byte[] data) {
        byte[] decrypt = apply(s -> s.getSdfApiAdapter().symDecrypt(s.getSessionHandle(), s.getKeyHandle(), AlgId.SGD_SM4_CBC, new byte[8], data));
        return PaddingUtil.PKCS7Unpadding(decrypt);
    }
}