package meerkat.crypto.concrete;

import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import meerkat.crypto.Encryption;
import meerkat.protobuf.ConcreteCrypto;
import meerkat.protobuf.Crypto;
import org.spongycastle.crypto.params.AsymmetricKeyParameter;
import org.spongycastle.crypto.params.ECDomainParameters;
import org.spongycastle.crypto.params.ECKeyParameters;
import org.spongycastle.crypto.params.ECPublicKeyParameters;
import org.spongycastle.crypto.util.PublicKeyFactory;
import org.spongycastle.jce.spec.ECParameterSpec;
import org.spongycastle.math.ec.ECCurve;
import org.spongycastle.math.ec.ECPoint;
import org.spongycastle.util.BigIntegers;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.factcenter.qilin.primitives.concrete.ECElGamal;
import org.factcenter.qilin.primitives.concrete.ECGroup;
import org.factcenter.qilin.util.PRGRandom;
import org.factcenter.qilin.util.Pair;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.spec.*;
import java.util.Random;

/**
 * Created by talm on 17/11/15.
 */
public class ECElGamalEncryption implements Encryption {

    private static GlobalCryptoSetup globalCryptoSetup = GlobalCryptoSetup.getInstance();

    final Logger logger = LoggerFactory.getLogger(getClass());

    public final static String KEY_ALGORITHM = "ECDH";

    /**
     * The Qilin format El-Gamal public key
     */
    ECElGamal.PK elGamalPK;

    ECCurve curve;

    ECGroup group;

    public ECGroup getGroup() { return group; }

    public ECElGamal.PK getElGamalPK() {
        return elGamalPK;
    }

    public void init(ConcreteCrypto.ElGamalPublicKey serializedPk) throws InvalidKeySpecException {
        AsymmetricKeyParameter keyParam;

        try {
            keyParam = PublicKeyFactory.createKey(serializedPk.getSubjectPublicKeyInfo().toByteArray());
        } catch (IOException e) {
            // Shouldn't every happen
            logger.error("Invalid Public Key Encoding", e);
            throw new InvalidKeySpecException("Invalid Public Key Encoding", e);
        }

        if (!(keyParam instanceof ECPublicKeyParameters)) {
            logger.error("Public key is a {}, not a valid public EC Key!", keyParam.getClass());
            throw new InvalidKeySpecException("Not a valid EC public key!");
        }

        ECDomainParameters params = ((ECKeyParameters) keyParam).getParameters();
        ECParameterSpec ecParams = new ECParameterSpec(params.getCurve(), params.getG(), params.getN(), params.getH(),
                params.getSeed());

        curve = params.getCurve();
        group = new ECGroup(ecParams);

        elGamalPK = new ECElGamal.PK(group, ((ECPublicKeyParameters) keyParam).getQ());
    }

    /**
     * Encode a group element into a protobuf.
     * @param p
     * @return
     */
    public ConcreteCrypto.GroupElement encodeElement(ECPoint p) {
        return encodeElement(group, p);
    }


    /**
     * Encode a group element into a protobuf.
     *
     * @param group group to use for encoding
     * @param p element to encode.
     * @return
     */
    public static ConcreteCrypto.GroupElement encodeElement(ECGroup group, ECPoint p) {
        return ConcreteCrypto.GroupElement.newBuilder().setData(ByteString.copyFrom(p.getEncoded(true))).build();
    }

    /**
     * Decode from the serialized representation to an {@link ECPoint} object.
     * @param bs
     * @return
     */
    public ECPoint decodeElement(ConcreteCrypto.GroupElement bs) {
        return decodeElement(group, bs);
    }

    /**
     * Decode from the serialized representation to an {@link ECPoint} object.
     * @param group group to use for decoding.
     * @param bs
     * @return
     */
    public static ECPoint decodeElement(ECGroup group, ConcreteCrypto.GroupElement bs) {
        return group.decode(bs.getData().toByteArray());
    }

    @Override
    public Crypto.RerandomizableEncryptedMessage encrypt(Message plaintext, Crypto.EncryptionRandomness rnd) {

        // We write the message using writeDelimited to so the length gets prepended.
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        try {
            plaintext.writeDelimitedTo(out);
        } catch (IOException e) {
            logger.error("Should never happen!", e);
            throw new RuntimeException("Error in ByteArrayOutputStream!", e);
        }
        byte[] msg = out.toByteArray();
        ECPoint encodedMsg = group.injectiveEncode(msg, new PRGRandom(msg));

        BigInteger rndInt = extractRandomness(rnd);
        Pair<ECPoint,ECPoint> cipherText = elGamalPK.encrypt(encodedMsg, rndInt);
        ConcreteCrypto.ElGamalCiphertext encodedCipherText = ConcreteCrypto.ElGamalCiphertext.newBuilder()
                .setC1(encodeElement(cipherText.a))
                .setC2(encodeElement(cipherText.b))
                .build();

        return Crypto.RerandomizableEncryptedMessage.newBuilder()
                .setData(encodedCipherText.toByteString())
                .build();
    }

    public static ConcreteCrypto.ElGamalCiphertext RerandomizableEncryptedMessage2ElGamalCiphertext(Crypto.RerandomizableEncryptedMessage msg) throws InvalidProtocolBufferException {
        return  ConcreteCrypto.ElGamalCiphertext.parseFrom(msg.getData());
    }

    public  static Crypto.RerandomizableEncryptedMessage elGamalCiphertext2RerandomizableEncryptedMessage(ConcreteCrypto.ElGamalCiphertext msg) {
        return Crypto.RerandomizableEncryptedMessage.newBuilder()
                .setData(msg.toByteString()).build();
    }

    @Override
    public Crypto.RerandomizableEncryptedMessage rerandomize(Crypto.RerandomizableEncryptedMessage msg, Crypto.EncryptionRandomness rnd) throws InvalidProtocolBufferException {
        BigInteger rndInt = extractRandomness(rnd);
        Pair<ECPoint,ECPoint> randomizer = elGamalPK.encrypt(curve.getInfinity(), rndInt);
        ConcreteCrypto.ElGamalCiphertext originalEncodedCipher= ConcreteCrypto.ElGamalCiphertext.parseFrom(msg.getData());

        Pair<ECPoint,ECPoint> originalCipher = new Pair<ECPoint, ECPoint>(
                decodeElement(originalEncodedCipher.getC1()),
                decodeElement(originalEncodedCipher.getC2()));
        Pair<ECPoint,ECPoint> newCipher = elGamalPK.add(originalCipher, randomizer);

        return Crypto.RerandomizableEncryptedMessage.newBuilder()
                .setData(
                    ConcreteCrypto.ElGamalCiphertext.newBuilder()
                    .setC1(encodeElement(newCipher.a))
                    .setC2(encodeElement(newCipher.b))
                    .build().toByteString()
                ).build();
    }

    @Override
    public Crypto.EncryptionRandomness generateRandomness(Random rand) {
        BigInteger randomInt = generateRandomExponent(rand);
        Crypto.EncryptionRandomness retval = Crypto.EncryptionRandomness.newBuilder()
                .setData(ByteString.copyFrom(BigIntegers.asUnsignedByteArray(randomInt))).build();

        return retval;
    }

    /**
     * Return the discrete log of a random element in the group (i.e., a random value in Z_{groupOrder})
     * @param rand
     * @return
     */
    public BigInteger generateRandomExponent(Random rand) {
        return new BigInteger(group.getCurveParams().getN().bitLength() - 1, rand);
    }

    /**
     * Extract (deserialize) the random exponent from a {@link Crypto.EncryptionRandomness} protobuf.
     *
     * @param encryptionRandomness
     * @return
     */
    public BigInteger extractRandomness(Crypto.EncryptionRandomness encryptionRandomness){
        return BigIntegers.fromUnsignedByteArray(encryptionRandomness.getData().toByteArray());
    }
}
