package meerkat.crypto.concrete;

import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import meerkat.protobuf.ConcreteCrypto;
import meerkat.protobuf.Crypto;
import org.spongycastle.jce.spec.ECParameterSpec;
import org.spongycastle.jce.spec.ECPublicKeySpec;
import org.spongycastle.math.ec.ECPoint;
import org.factcenter.qilin.primitives.concrete.ECElGamal;
import org.factcenter.qilin.primitives.concrete.ECGroup;
import org.factcenter.qilin.primitives.generic.ElGamal;
import org.factcenter.qilin.util.Pair;

import java.io.ByteArrayInputStream;
import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;

/**
 * Created by talm on 02/11/16.
 */
public class Util {

    public final static String ENCRYPTION_KEY_ALGORITHM = ECElGamalEncryption.KEY_ALGORITHM;

    /**
     * Decode a BigInteger from a protobuf
     * @param i
     * @return
     */
    public static BigInteger decodeBigInteger(Crypto.BigInteger i) {
        return new BigInteger(i.getData().toByteArray());
    }

    /**
     * Encode a BigInteger in a protobuf.
     * @param i
     * @return
     */
    public static Crypto.BigInteger encodeBigInteger(BigInteger i) {
        return Crypto.BigInteger.newBuilder().setData(ByteString.copyFrom(i.toByteArray())).build();
    }

    /**
     * Serialize an El-Gamal public key into a form acceptable by {@link ECElGamalEncryption}
     * @param pk
     * @return
     */
    public static ConcreteCrypto.ElGamalPublicKey encodePK(ECGroup group, ElGamal.PK<ECPoint> pk) {
        ECPoint pkPoint = pk.getPK();
        ECParameterSpec params = group.getCurveParams();

        ECPublicKeySpec pubKeySpec = new ECPublicKeySpec(pkPoint, params);

        try {
            KeyFactory fact = KeyFactory.getInstance(ENCRYPTION_KEY_ALGORITHM,
                    GlobalCryptoSetup.getInstance().getBouncyCastleProvider());
            PublicKey javaPk = fact.generatePublic(pubKeySpec);
            ConcreteCrypto.ElGamalPublicKey serializedPk = ConcreteCrypto.ElGamalPublicKey.newBuilder()
                    .setSubjectPublicKeyInfo(ByteString.copyFrom(javaPk.getEncoded())).build();

            return serializedPk;
        } catch (Exception e) {
            throw new RuntimeException("Error converting public key!", e);
        }
    }

    /**
     * Deserialize an ECElGamal secret key.
     * @param serializedPk
     * @param serializedSk
     * @return
     * @throws InvalidKeySpecException
     */
    public static ECElGamal.SK decodeSK(ConcreteCrypto.ElGamalPublicKey serializedPk, Crypto.BigInteger serializedSk)
            throws InvalidKeySpecException {
        ECElGamalEncryption enc = new ECElGamalEncryption();
        enc.init(serializedPk);
        ECElGamal.SK sk = new ECElGamal.SK(enc.getGroup(), decodeBigInteger(serializedSk));
        return sk;
    }

    /**
     * Standard (non-threshold) decryption for testing purposes.
     * @param secretKey
     * @return
     */
    public static <T extends Message>  T decrypt(Class<T> plaintextMessageType, ECElGamal.SK secretKey, ECGroup group, Crypto.RerandomizableEncryptedMessage opaqueCipher)
            throws InvalidProtocolBufferException {
        ConcreteCrypto.ElGamalCiphertext cipherText = ConcreteCrypto.ElGamalCiphertext.parseFrom(opaqueCipher.getData());
        ConcreteCrypto.GroupElement c1encoded = cipherText.getC1();
        ConcreteCrypto.GroupElement c2encoded = cipherText.getC2();

        ECPoint c1 = ECElGamalEncryption.decodeElement(group, c1encoded);
        ECPoint c2 = ECElGamalEncryption.decodeElement(group, c2encoded);

        assert (group.contains(c1));
        assert (group.contains(c2));

        ECPoint plaintextEncoded = secretKey.decrypt(new Pair<ECPoint, ECPoint>(c1, c2));

        byte[] plaintext = group.injectiveDecode(plaintextEncoded);

        ByteArrayInputStream in = new ByteArrayInputStream(plaintext);

        try {
            java.lang.reflect.Method newBuilder = plaintextMessageType.getMethod("newBuilder");
            Message.Builder builder = (Message.Builder) newBuilder.invoke(plaintextMessageType);
            builder.mergeDelimitedFrom(in);
            return plaintextMessageType.cast(builder.build());
        } catch (Exception e) {
            throw new InvalidProtocolBufferException("Plaintext protobuf error");
        }
    }


}
