package com.dji.sample.territory.utils.jym;

import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.generators.ECKeyPairGenerator;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECKeyGenerationParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.math.ec.ECFieldElement;
import org.bouncycastle.math.ec.ECFieldElement.Fp;
import org.bouncycastle.math.ec.ECPoint;

import java.math.BigInteger;
import java.security.SecureRandom;

public class SM2Factory {
    private static final BigInteger a = new BigInteger("fffffffeffffffffffffffffffffffffffffffff00000000fffffffffffffffc", 16);
    private static final BigInteger b = new BigInteger("28e9fa9e9d9f5e344d5a9e4bcf6509a7f39789f515ab8f92ddbcbd414d940e93", 16);
    private static final BigInteger gx = new BigInteger("32c4ae2c1f1981195f9904466a39c9948fe30bbff2660be1715a4589334c74c7", 16);
    private static final BigInteger gy = new BigInteger("bc3736a2f4f6779c59bdcee36b692153d0a9877cc62a474002df32e52139f0a0", 16);
    private static final BigInteger n = new BigInteger("fffffffeffffffffffffffffffffffff7203df6b21c6052b53bbf40939d54123", 16);
    private static final BigInteger p = new BigInteger("fffffffeffffffffffffffffffffffffffffffff00000000ffffffffffffffff", 16);
    private static final int h = 1;
    public final ECFieldElement ecc_gx_fieldelement;
    public final ECFieldElement ecc_gy_fieldelement;
    public final ECCurve ecc_curve;
    public final ECPoint ecc_point_g;
    public final ECDomainParameters ecc_bc_spec;
    public final ECKeyPairGenerator ecc_key_pair_generator;

    public static SM2Factory getInstance() {
        return new SM2Factory();
    }

    public SM2Factory() {
        this.ecc_gx_fieldelement = new Fp(p, gx);
        this.ecc_gy_fieldelement = new Fp(p, gy);
        this.ecc_curve = new ECCurve.Fp(p, a, b);
        this.ecc_point_g = new ECPoint.Fp(this.ecc_curve, this.ecc_gx_fieldelement, this.ecc_gy_fieldelement, false);
        this.ecc_bc_spec = new ECDomainParameters(this.ecc_curve, this.ecc_point_g, n);
        ECKeyGenerationParameters ecc_ecgenparam = new ECKeyGenerationParameters(this.ecc_bc_spec, new SecureRandom());
        this.ecc_key_pair_generator = new ECKeyPairGenerator();
        this.ecc_key_pair_generator.init(ecc_ecgenparam);
    }

    public byte[] sm2GetZ(byte[] userId, ECPoint userKey) {
        SM3Digest sm3 = new SM3Digest();
        int len = userId.length * 8;
        sm3.update((byte)(len >> 8 & 255));
        sm3.update((byte)(len & 255));
        sm3.update(userId, 0, userId.length);
        byte[] p = Util.byteConvert32Bytes(a);
        sm3.update(p, 0, p.length);
        p = Util.byteConvert32Bytes(b);
        sm3.update(p, 0, p.length);
        p = Util.byteConvert32Bytes(gx);
        sm3.update(p, 0, p.length);
        p = Util.byteConvert32Bytes(gy);
        sm3.update(p, 0, p.length);
        p = Util.byteConvert32Bytes(userKey.normalize().getXCoord().toBigInteger());
        sm3.update(p, 0, p.length);
        p = Util.byteConvert32Bytes(userKey.normalize().getYCoord().toBigInteger());
        sm3.update(p, 0, p.length);
        byte[] md = new byte[sm3.getDigestSize()];
        sm3.doFinal(md, 0);
        return md;
    }

    public void sm2Sign(byte[] md, BigInteger userD, ECPoint userKey,SM2Result sm2Result) {
        BigInteger e = new BigInteger(1, md);
        BigInteger k = null;
        ECPoint kp = null;
        BigInteger r = null;
        BigInteger s = null;

        while(true) {
            do {
                AsymmetricCipherKeyPair keypair = this.ecc_key_pair_generator.generateKeyPair();
                ECPrivateKeyParameters ecpriv = (ECPrivateKeyParameters)keypair.getPrivate();
                ECPublicKeyParameters ecpub = (ECPublicKeyParameters)keypair.getPublic();
                k = ecpriv.getD();
                kp = ecpub.getQ();
                r = e.add(kp.getXCoord().toBigInteger());
                r = r.mod(n);
            } while(r.equals(BigInteger.ZERO));

            if (!r.add(k).equals(n) && r.toString(16).length() == 64) {
                BigInteger da_1 = userD.add(BigInteger.ONE);
                da_1 = da_1.modInverse(n);
                s = r.multiply(userD);
                s = k.subtract(s).mod(n);
                s = da_1.multiply(s).mod(n);
                if (!s.equals(BigInteger.ZERO) && s.toString(16).length() == 64) {
                    sm2Result.r = r;
                    sm2Result.s = s;
                    return;
                }
            }
        }
    }

    public void sm2Verify(byte[] md, ECPoint userKey, BigInteger r, BigInteger s, SM2Result sm2Result) {
        sm2Result.R = null;
        BigInteger e = new BigInteger(1, md);
        BigInteger t = r.add(s).mod(n);
        if (!t.equals(BigInteger.ZERO)) {
            ECPoint x1y1 = this.ecc_point_g.multiply(sm2Result.s);
            x1y1 = x1y1.add(userKey.multiply(t));
            sm2Result.R = e.add(x1y1.normalize().getXCoord().toBigInteger()).mod(n);
        }
    }
}