package com.sunyard.chsm.pool;

import com.sunyard.chsm.sdf.SdfApiService;
import com.sunyard.chsm.sdf.context.AlgId;
import com.sunyard.chsm.sdf.model.EccKey;
import com.sunyard.chsm.sdf.util.PaddingUtil;
import com.sunyard.chsm.utils.gm.BCSM3Utils;
import lombok.extern.slf4j.Slf4j;
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, byte[] key, byte[] iv, byte[] data) {
        return apply(s -> {
            String keyHandle = s.getSdfApiAdapter().importKey(s.getSessionHandle(), key);
            byte[] encrypt = s.getSdfApiAdapter().symEncrypt(s.getSessionHandle(), keyHandle, alg, iv, PaddingUtil.PKCS7Padding(data));
            s.getSdfApiAdapter().destroyKey(s.getSessionHandle(), keyHandle);
            return encrypt;
        });
    }

    @Override
    public byte[] symDecrypt(AlgId alg, 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.PKCS7Unpadding(decrypt);
    }

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

    @Override
    public byte[] calculateMAC(byte[] symKey, byte[] pucIv, byte[] pucData) {
        return apply(s -> {
            String keyHandle = s.getSdfApiAdapter().importKey(s.getSessionHandle(), symKey);
            byte[] mac = s.getSdfApiAdapter().calculateMAC(s.getSessionHandle(), keyHandle, AlgId.SGD_SM4_MAC, pucIv, pucData);
            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);
    }
}
