package meerkat.crypto.concrete;

import meerkat.protobuf.ConcreteCrypto;
import meerkat.protobuf.Crypto;
import meerkat.protobuf.Voting;
import org.spongycastle.math.ec.ECPoint;
import org.junit.Before;
import org.junit.Test;
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.Pair;

import java.math.BigInteger;
import java.util.Random;

import static org.junit.Assert.*;

/**
 * Test class for {@link ECElGamalEncryption}
 */
public class ECElGamalEncryptionTest {
    final Logger logger = LoggerFactory.getLogger(getClass());
    /**
     * Number of times to repeat probabilistic tests.
     */
    public final static int CONFIDENCE = 10;

    Random rand = new Random(0); // Insecure deterministic random for testing.

    ECElGamal.SK key;
    ECGroup group;
    ECElGamalEncryption enc;
    ConcreteCrypto.ElGamalPublicKey serializedPk;


    @Before
    public void setup() throws Exception {
        group = new ECGroup("secp256k1");
        BigInteger sk = ECElGamal.generateSecretKey(group, rand);
        key = new ECElGamal.SK(group, sk);
        serializedPk = ECElGamalUtils.serializePk(group, key);


        enc = new ECElGamalEncryption();

        enc.init(serializedPk);
    }


    Voting.PlaintextBallot genRandomBallot(int numQuestions, int numAnswers, int maxAnswer) {
        Voting.PlaintextBallot.Builder ballot = Voting.PlaintextBallot.newBuilder();
        ballot.setSerialNumber(rand.nextInt(1000000));
        for (int i = 0; i < numQuestions; ++i) {
            Voting.BallotAnswer.Builder answers = ballot.addAnswersBuilder();
            for (int j = 0; j < numAnswers; ++j) {
                answers.addAnswer(rand.nextInt(maxAnswer));
            }
        }
        return ballot.build();
    }

    /**
     * Testing just the key management
     * @throws Exception
     */
    @Test
    public void testPkSerialization() throws Exception {
        ECElGamal.PK pk = enc.getElGamalPK();

        ECPoint point = enc.getGroup().sample(rand);
        Pair<ECPoint, ECPoint> cipher = pk.encrypt(point, pk.getRandom(rand));

        ECPoint decrypted = key.decrypt(cipher);

        assertEquals("Decrypted value not equal to encrypted value!", point, decrypted);
    }

    @Test
    public void testEncryption() throws Exception {
        for (int i = 0; i < CONFIDENCE; ++i) {
            Voting.PlaintextBallot msg = genRandomBallot(2,3,16); // 2 questions with 3 answers each, in range 0-15.
            if (msg.getSerializedSize() > enc.getGroup().getInjectiveEncodeMsgLength()) {
                logger.error("Test Message too big (|msg|={} > max={}), expect failure.",
                        msg.getSerializedSize(), enc.getGroup().getInjectiveEncodeMsgLength());
            }

            Crypto.RerandomizableEncryptedMessage cipherText = enc.encrypt(msg, enc.generateRandomness(rand));

            Voting.PlaintextBallot decrypted = ECElGamalUtils.decrypt(Voting.PlaintextBallot.class, key, group, cipherText);

            assertEquals("Decrypted value differs from encrypted value (i="+i+")!", msg, decrypted);
        }
    }

    @Test
    public void testRerandomizeModifiesCiphertext() throws Exception {
        Voting.PlaintextBallot msg = genRandomBallot(2,3,16); // 2 questions with 3 answers each, in range 0-15.
        Crypto.RerandomizableEncryptedMessage cipher1 = enc.encrypt(msg, enc.generateRandomness(rand));
        Crypto.RerandomizableEncryptedMessage cipher2 = enc.rerandomize(cipher1, enc.generateRandomness(rand));
        assertNotEquals("Rerandomized cipher identical to original!", cipher1, cipher2);
    }

    @Test
    public void testRerandomizePreservesPlaintext() throws Exception {
        for (int i = 0; i < CONFIDENCE; ++i) {
            Voting.PlaintextBallot msg = genRandomBallot(2,3,16); // 2 questions with 3 answers each, in range 0-15.

            Crypto.RerandomizableEncryptedMessage cipher = enc.encrypt(msg, enc.generateRandomness(rand));
            Crypto.RerandomizableEncryptedMessage cipher2 = cipher;
            for (int j = 0; j < CONFIDENCE; ++j)
                cipher2 = enc.rerandomize(cipher2, enc.generateRandomness(rand));

            Voting.PlaintextBallot decrypted = ECElGamalUtils.decrypt(Voting.PlaintextBallot.class, key, group,
                    cipher2);

            assertEquals("Decrypted value differs from original encrypted value (i="+i+")!", msg, decrypted);
        }
    }
}

