/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.products.ov.mixnet.commons.homomorphic.impl;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.CiphertextImpl;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.Ciphertext;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.cryptolib.mathematical.groups.impl.ZpSubgroup;
import com.scytl.products.ov.mixnet.commons.homomorphic.Cryptosystem;
import com.scytl.products.ov.mixnet.commons.homomorphic.Plaintext;
import com.scytl.products.ov.mixnet.commons.homomorphic.Randomness;
import com.scytl.products.ov.mixnet.commons.tools.ExponentTools;

/**
 * Implementation of the Gjosteen ElGamal cryptosystem.
 */
public class GjosteenElGamal implements Cryptosystem {

    private final Exponent[] privateKey;

    private final ZpGroupElement[] publicKey;

    private final ZpGroupElement generator;

    private final ZpSubgroup group;

    /**
     * Create a GjosteenElGamal cryptosystem that will operate over the received mathematical group.
     * <P>
     * This constructor calculates the public key (which can be derived from the private key), and sets both the public
     * key and the private key within the created cryptosystem. Therefore, a cryptosystem that has been created using
     * this constructor is fully initialized and ready to be used for encrypting and decrypting.
     *
     * @param group
     *            the mathematical group over which this cryptosystem operates.
     * @param privKey
     *            the private key.
     */
    public GjosteenElGamal(final ZpSubgroup group, final Exponent[] privKey) {
        this.group = group;
        this.privateKey = privKey;
        this.generator = group.getGenerator();
        int numKeyElements = privKey.length;
        this.publicKey = new ZpGroupElement[numKeyElements];
        for (int i = 0; i < numKeyElements; i++) {
            try {
            	this.publicKey[i] = this.generator.exponentiate(this.privateKey[i]);
            } catch (GeneralCryptoLibException e) {
                throw new CryptoLibException(e);
            }
        }
    }

    /**
     * Create a GjosteenElGamal cryptosystem that will operate over the received mathematical group.
     * <P>
     * This constructor ONLY sets both the public key within the created cryptosystem, it DOES NOT SET A PRIVATE KEY.
     * Therefore, a cryptosystem that has been created using this constructor can only be used for decrypting if a
     * private key is passed to the decrypt method at the time of performing the decryption.
     *
     * @param group
     *            the mathematical group over which this cryptosystem operates.
     * @param pubKey
     *            the public key.
     */
    public GjosteenElGamal(final ZpSubgroup group, final ElGamalPublicKey pubKey) {
    	this.group = group;
    	this.generator = group.getGenerator();
    	this.privateKey = null;
        List<ZpGroupElement> publicKeyAsListGroupElements = pubKey.getKeys();
        this.publicKey = publicKeyAsListGroupElements.toArray(new ZpGroupElement[publicKeyAsListGroupElements.size()]);
    }

    /**
     * Create a GjosteenElGamal cryptosystem that will operate over the received mathematical group.
     * <P>
     * This constructor creates a random private key, containing {@code numOpts} elements, it then derives the
     * corresponding public key from the private key. It sets both the public key and the private key within the created
     * cryptosystem. Therefore, a cryptosystem that has been created using this constructor is fully initialized and
     * ready to be used for encrypting and decrypting.
     *
     * @param group
     *            the mathematical group over which this cryptosystem operates.
     * @param numOpts
     *            the number of elements that should exist in the created pair of keys.
     */
    public GjosteenElGamal(final ZpSubgroup group, final int numOpts) {
        this(group, ExponentTools.getVectorRandomExponent(numOpts, group.getQ()));
    }

    @Override
    public Randomness getFreshRandomness() {
        return new GjosteenElGamalRandomness(ExponentTools.getRandomExponent(group.getQ()));
    }

    @Override
    public Randomness[] getVectorFreshRandomness(final int length) {
        Randomness[] result = new Randomness[length];
        for (int i = 0; i < result.length; i++) {
            result[i] = getFreshRandomness();
        }
        return result;
    }

    @Override
    public Ciphertext encrypt(final Plaintext plaintext, final Randomness randomness) {

        try {
            final Exponent randomExponent = ((GjosteenElGamalRandomness) randomness).getRandomnessValue();

            final ZpGroupElement gamma = generator.exponentiate(randomExponent);

            final List<ZpGroupElement> phis = new ArrayList<>();

            for (int i = 0, numPhis = publicKey.length; i < numPhis; i++) {
                phis.add(publicKey[i].exponentiate(randomExponent)
                    .multiply(((GjosteenElGamalPlaintext) plaintext).getValue(i)));
            }

            return new CiphertextImpl(gamma, phis);

        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(e);
        }
    }

    @Override
    public Ciphertext encrypt(final Plaintext plaintext) {

        return encrypt(plaintext, new GjosteenElGamalRandomness(ExponentTools.getRandomExponent(group.getQ())));
    }

    public Ciphertext encrypt(final ZpGroupElement[] plaintext, final Randomness randomness) {

        return encrypt(new GjosteenElGamalPlaintext(plaintext), randomness);
    }

    public Ciphertext encrypt(final ZpGroupElement[] plaintext) {

        return encrypt(plaintext, getFreshRandomness());
    }

    /**
     * Decrypts the received ciphertext using the private key that is set within this cryptosystem.
     *
     * @param ciphertext
     *            the ciphertext to be decrypted.
     * @return the decrypted plaintext.
     */
    public ZpGroupElement[] decrypt(final ZpGroupElement[] ciphertext) {

        validateInputToSingleArgDecrypt(ciphertext);

        return decrypt(ciphertext, privateKey);
    }

    /**
     * Decrypts the received ciphertext using the received private key.
     *
     * @param ciphertext
     *            the ciphertext to be decrypted.
     * @param privateKeyExponents
     *            the private key to be used to decrypt the ciphertext (as an array of Exponents).
     * @return the decrypted plaintext.
     */
    public ZpGroupElement[] decrypt(final ZpGroupElement[] ciphertext, final Exponent[] privateKeyExponents) {

        try {

            validateInputToDoubleArgDecrypt(ciphertext, privateKeyExponents);

            ZpGroupElement gamma = ciphertext[0];

            int numPhis = ciphertext.length - 1;
            ZpGroupElement[] phis = new ZpGroupElement[numPhis];
            System.arraycopy(ciphertext, 1, phis, 0, numPhis);

            Exponent negatedExponent;
            ZpGroupElement[] plaintext = new ZpGroupElement[numPhis];

            for (int i = 0; i < numPhis; i++) {

                // Compute the e = negate (-) of _privateKey[i]
                negatedExponent = privateKeyExponents[i].negate();

                // Compute dm[i]= gamma^(e) * phi[i]
                plaintext[i] = gamma.exponentiate(negatedExponent).multiply(phis[i]);
            }

            return plaintext;

        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(e);
        }
    }

    @Override
    public Ciphertext encryptRaisingToRandom(final Exponent b, final Randomness tau) {

        try {

            ZpGroupElement[] aux = new ZpGroupElement[publicKey.length];
            ZpGroupElement aux2 = group.getGenerator().exponentiate(b);
            for (int i = 0; i < aux.length; i++) {
                aux[i] = aux2;
            }

            Plaintext p = new GjosteenElGamalPlaintext(aux);
            return encrypt(p, tau);

        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(e);
        }
    }

    @Override
    public Ciphertext getEncryptionOf1() {
        ZpGroupElement[] aux = new ZpGroupElement[publicKey.length];
        for (int i = 0; i < aux.length; i++) {
            aux[i] = group.getIdentity();
        }
        return encrypt(aux);
    }

    @Override
    public Ciphertext getEncryptionOf1(final Randomness r) {
        ZpGroupElement[] aux = new ZpGroupElement[publicKey.length];
        for (int i = 0; i < aux.length; i++) {
            aux[i] = group.getIdentity();
        }
        return encrypt(aux, r);
    }

    @Override
    public Randomness get0Randomness() {

        try {
            return new GjosteenElGamalRandomness(new Exponent(group.getQ(), BigInteger.ZERO));

        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(e);
        }
    }

    @Override
    public int getNumberOfMessages() {
        return publicKey.length;
    }

    private void validateInputToSingleArgDecrypt(final ZpGroupElement[] ciphertext) {

        if (ciphertext == null) {
            throw new IllegalArgumentException("The received ciphertext was null");
        } else if (privateKey == null) {
            throw new IllegalArgumentException(
                "Cannot decrypt using this method because a private key is not set within this cryptosystem");
        } else if (ciphertext.length != privateKey.length + 1) {
            throw new IllegalArgumentException(
                "The ciphertext must be the same length as the private key that is set within this cryptosystem");
        }
    }

    private static void validateInputToDoubleArgDecrypt(final ZpGroupElement[] ciphertext,
            final Exponent[] privateKeyExponents) {

        if (ciphertext == null) {
            throw new IllegalArgumentException("The received ciphertext was null");
        } else if (privateKeyExponents == null) {
            throw new IllegalArgumentException("The received private key was null");
        } else if (ciphertext.length != privateKeyExponents.length + 1) {
            throw new IllegalArgumentException("The ciphertext must be the same length as the private key");
        }
    }
}
