package com.sunyard.chsm.sdf.adapter;

import com.sunyard.chsm.sdf.context.AlgId;
import com.sunyard.chsm.sdf.model.DeviceInfo;
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.LangUtils;
import com.sunyard.chsm.utils.gm.BCECUtils;
import com.sunyard.chsm.utils.gm.BCSM2Utils;
import com.sunyard.chsm.utils.gm.BCSM4Utils;
import lombok.SneakyThrows;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.BigIntegers;
import org.springframework.util.Assert;

import java.math.BigInteger;
import java.security.KeyPair;
import java.security.SecureRandom;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author liulu
 * @since 2024/11/5
 */
public class BcSdfApiAdaptor implements SdfApiAdapter {

    protected static final Map<String, byte[]> KEY_HANDLE_CONTEXT = new ConcurrentHashMap<>();

    private static final DeviceInfo deviceInfo;

    static {
        deviceInfo = new DeviceInfo();
        deviceInfo.setIssuerName("BC");
        deviceInfo.setDeviceName("BC-3000");
        deviceInfo.setDeviceSerial("BC202411051037");
        deviceInfo.setDeviceVersion(1);
        deviceInfo.setStandardVersion(1);
        deviceInfo.setAsymAlgAbility(new long[]{7493065891348563935L, 3000543215027029126L});
        deviceInfo.setSymAlgAbility(2782707529341253316L);
        deviceInfo.setHashAlgAbility(7010608454676760881L);
    }

    protected static final BcSdfApiAdaptor INSTANCE = new BcSdfApiAdaptor();

    @Override
    public String openDevice() {
        return "c95a78d9c04a557b7b46dbcb5f36cc66";
    }

    @Override
    public boolean closeDevice(String deviceHandle) {
        return true;
    }

    @Override
    public String openSession(String deviceHandle) {
        return UUID.randomUUID().toString();
    }

    @Override
    public void closeSession(String sessionHandle) {

    }

    @Override
    public DeviceInfo getDeviceInfo(String sessionHandle) {
        return deviceInfo;
    }

    @Override
    public byte[] generateRandom(String sessionHandle, int uiLength) {
        byte[] res = new byte[uiLength];
        new SecureRandom().nextBytes(res);
        return res;
    }

    @Override
    public EccPubKey exportEncPublicKeyECC(String sessionHandle, int uiKeyIndex) {
        BigInteger d = new BigInteger(1, getD());
        ECPoint q = BCSM2Utils.G_POINT.multiply(d).normalize();
        return new EccPubKey(256, q.getXCoord().getEncoded(), q.getYCoord().getEncoded());
    }

    private byte[] getD() {
        return LangUtils.merge(LangUtils.toByteArray(deviceInfo.getAsymAlgAbility()[0]),
                LangUtils.toByteArray(deviceInfo.getAsymAlgAbility()[1]),
                LangUtils.toByteArray(deviceInfo.getSymAlgAbility()),
                LangUtils.toByteArray(deviceInfo.getHashAlgAbility()));
    }

    @SneakyThrows
    @Override
    public EccKey generateKeyPairECC(String sessionHandle, AlgId alg) {
        // 生成密钥对
        KeyPair keyPair = BCSM2Utils.generateKeyPair();
        BCECPublicKey pubKey = (BCECPublicKey) keyPair.getPublic();
        BCECPrivateKey priKey = (BCECPrivateKey) keyPair.getPrivate();
        byte[] x = pubKey.getQ().getXCoord().getEncoded();
        byte[] y = pubKey.getQ().getYCoord().getEncoded();
        byte[] d = BigIntegers.asUnsignedByteArray(32, priKey.getD());
        return new EccKey(LangUtils.merge(x, y), d);
    }

    @Override
    public boolean getPrivateKeyAccessRight(String hSessionHandle, int uiKeyIndex, byte[] pucPassword) {
        return true;
    }

    @Override
    public EccCipher exchangeDigitEnvelopeBaseOnECC(String sessionHandle, int uiKeyIndex, EccPubKey pubKey, EccCipher pucEncDateIn) {

        ECPublicKeyParameters pub = BCECUtils.createECPublicKeyParameters(pubKey.getX(), pubKey.getY());
        ECPrivateKeyParameters pri = BCECUtils.createECPrivateKeyParameters(getD());
        try {
            byte[] k = BCSM2Utils.decrypt(pri, LangUtils.merge(new byte[]{0x04}, pucEncDateIn.getC1C3C2Bytes()));
            byte[] encrypt = BCSM2Utils.encrypt(pub, k);
            return EccCipher.fromBytes(encrypt);
        } catch (InvalidCipherTextException e) {
            throw new IllegalArgumentException(e);
        }
    }

    private static final SecureRandom RANDOM = new SecureRandom();

    @Override
    public EccSignature externalSignECC(String sessionHandle, EccPriKey privateKey, byte[] pucData) {
        ECPrivateKeyParameters pri = BCECUtils.createECPrivateKeyParameters(privateKey.getD());

        BigInteger n = pri.getParameters().getN();
        BigInteger d = pri.getD();

        BigInteger e = new BigInteger(1, pucData);
        BigInteger r, s = null;

        do {
            BigInteger k;
            do {
                k = new BigInteger(n.bitLength(), RANDOM);
            } while (k.compareTo(BigInteger.ONE) < 0 || k.compareTo(n) >= 0);

            ECPoint kG = pri.getParameters().getG().multiply(k).normalize();
            BigInteger x1 = kG.getAffineXCoord().toBigInteger();

            r = e.add(x1).mod(n);

            if (r.equals(BigInteger.ZERO) || r.add(k).equals(n)) {
                continue;
            }

            BigInteger dPlus1Inv = d.add(BigInteger.ONE).modInverse(n);
            s = dPlus1Inv.multiply(k.subtract(r.multiply(d))).mod(n);

        } while (Objects.equals(s, BigInteger.ZERO));
        return new EccSignature(BigIntegers.asUnsignedByteArray(32, r),
                BigIntegers.asUnsignedByteArray(32, s));
    }

    @Override
    public boolean externalVerifyECC(String sessionHandle, EccPubKey publicKey, byte[] pucData, EccSignature pucSignature) {
        ECPublicKeyParameters pub = BCECUtils.createECPublicKeyParameters(publicKey.getX(), publicKey.getY());
        if (pucData.length != 32) {
            throw new IllegalArgumentException("Hash length must be 32 bytes");
        }
        ECDomainParameters domainParams = pub.getParameters();
        ECPoint G = domainParams.getG();
        BigInteger n = domainParams.getN();
        ECPoint Q = pub.getQ();
        BigInteger e = new BigInteger(1, pucData);

        BigInteger r = new BigInteger(1, pucSignature.getR());
        BigInteger s = new BigInteger(1, pucSignature.getS());
        if (r.compareTo(BigIntegers.ONE) < 0 || r.compareTo(n) >= 0 || s.compareTo(BigIntegers.ONE) < 0 || s.compareTo(n) >= 0) {
            return false;
        }

        BigInteger t = r.add(s).mod(n);
        if (t.equals(BigInteger.ZERO)) {
            return false;
        }

        ECPoint point = G.multiply(s).add(Q.multiply(t)).normalize();

        BigInteger R = e.add(point.getAffineXCoord().toBigInteger()).mod(n);

        return R.equals(r);
    }

    @Override
    public EccCipher externalEncryptECC(String sessionHandle, EccPubKey pubKey, byte[] pucData) {
        ECPublicKeyParameters pub = BCECUtils.createECPublicKeyParameters(pubKey.getX(), pubKey.getY());
        try {
            byte[] encrypt = BCSM2Utils.encrypt(pub, pucData);
            return EccCipher.fromBytes(encrypt);
        } catch (InvalidCipherTextException e) {
            throw new IllegalArgumentException(e);
        }
    }


    @Override
    public byte[] externalDecryptECC(String sessionHandle, EccPriKey pucPrivateKeyEcc, EccCipher pucEncData) {
        ECPrivateKeyParameters pri = BCECUtils.createECPrivateKeyParameters(pucPrivateKeyEcc.getD());
        try {
            return BCSM2Utils.decrypt(pri, LangUtils.merge(new byte[]{0x04}, pucEncData.getC1C3C2Bytes()));
        } catch (InvalidCipherTextException e) {
            throw new IllegalArgumentException(e);
        }
    }

    @Override
    public String importKeyWithISKECC(String sessionHandle, int uiIskIndex, EccCipher eccCipher) {
        ECPrivateKeyParameters pri = BCECUtils.createECPrivateKeyParameters(getD());

        try {
            byte[] pucKey = BCSM2Utils.decrypt(pri, LangUtils.merge(new byte[]{0x04}, eccCipher.getC1C3C2Bytes()));
            Assert.isTrue(pucKey.length == 16, "密钥长度错误");
            String key = UUID.randomUUID().toString();
            KEY_HANDLE_CONTEXT.put(key, pucKey);
            return key;
        } catch (InvalidCipherTextException e) {
            throw new IllegalArgumentException(e);
        }
    }

    @Override
    public String importKey(String sessionHandle, byte[] pucKey) {
        Assert.isTrue(pucKey.length == 16, "密钥长度错误");
        String key = UUID.randomUUID().toString();
        KEY_HANDLE_CONTEXT.put(key, pucKey);
        return key;
    }

    @Override
    public int destroyKey(String sessionHandle, String hKeyHandle) {
        KEY_HANDLE_CONTEXT.remove(hKeyHandle);
        return 0;
    }

    @Override
    @SneakyThrows
    public byte[] symEncrypt(String sessionHandle, String keyHandle, AlgId alg, byte[] pucIv, byte[] pucData) {
        byte[] symKey = getSymKey(keyHandle);
        switch (alg) {
            case SGD_SM4_ECB:
                return BCSM4Utils.encrypt_ECB_NoPadding(symKey, pucData);
            case SGD_SM4_CBC:
                return BCSM4Utils.encrypt_CBC_NoPadding(symKey, pucIv, pucData);
            default:
                throw new IllegalArgumentException("不支持的算法:" + alg.name());
        }
    }

    @Override
    @SneakyThrows
    public byte[] symDecrypt(String sessionHandle, String keyHandle, AlgId alg, byte[] pucIv, byte[] pucEncData) {
        byte[] symKey = getSymKey(keyHandle);
        switch (alg) {
            case SGD_SM4_ECB:
                return BCSM4Utils.decrypt_ECB_NoPadding(symKey, pucEncData);
            case SGD_SM4_CBC:
                return BCSM4Utils.decrypt_CBC_NoPadding(symKey, pucIv, pucEncData);
            default:
                throw new IllegalArgumentException("不支持的算法:" + alg.name());
        }
    }

    protected byte[] getSymKey(String keyHandle) {
        byte[] key = KEY_HANDLE_CONTEXT.get(keyHandle);
        Assert.notNull(key, "密钥句柄不存在");
        return key;
    }

    @Override
    @SneakyThrows
    public byte[] calculateMAC(String sessionHandle, String keyHandle, AlgId uiAlg, byte[] pucIv, byte[] pucData) {
        byte[] symKey = getSymKey(keyHandle);
        return BCSM4Utils.doCBCMac(symKey, pucIv, null, pucData);
    }

    protected static final Map<String, SM3Digest> HASH_CONTEXT = new ConcurrentHashMap<>();

    @Override
    public int hashInit(String sessionHandle, AlgId alg, EccPubKey pucPublicKey, byte[] pucID) {
        HASH_CONTEXT.put(sessionHandle, new SM3Digest());
        return 0;
    }

    @Override
    public int hashUpdate(String sessionHandle, byte[] pucData) {
        Optional.ofNullable(HASH_CONTEXT.get(sessionHandle))
                .ifPresent(it -> it.update(pucData, 0, pucData.length));
        return 0;
    }

    @Override
    public byte[] hashFinish(String sessionHandle) {
        SM3Digest digest = HASH_CONTEXT.remove(sessionHandle);
        Assert.notNull(digest, "session异常");
        byte[] hash = new byte[digest.getDigestSize()];
        digest.doFinal(hash, 0);
        return hash;
    }

}
