/*
 * 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.cryptolib.elgamal.factory;

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.commons.utils.validations.Validate;
import com.scytl.cryptolib.elgamal.bean.ElGamalComputationsValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncrypterValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.configuration.ConfigGroupType;
import com.scytl.cryptolib.elgamal.configuration.Group;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalEncrypter;
import com.scytl.cryptolib.mathematical.groups.activity.GroupElementsCompressor;
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.cryptolib.mathematical.groups.utils.GroupUtils;
import com.scytl.cryptolib.primitives.securerandom.factory.CryptoRandomInteger;

/**
 * Implementation of an ElGamal encrypter.
 */
public final class CryptoElGamalEncrypter implements
        CryptoAPIElGamalEncrypter {

    private final ElGamalPublicKey _elGamalPublicKey;

    private final int _elGamalPublicKeyLength;

    private final CryptoRandomInteger _cryptoRandomInteger;

    private final Group _group;

    /**
     * Creates an ElGamal encrypter.
     *
     * @param elGamalPublicKey
     *            the ElGamal public key.
     * @param cryptoRandomInteger
     *            a generator of random integers.
     * @param configGroupType
     *            a group type that is defined in the cryptolib policy.
     * @throws GeneralCryptoLibException
     *             if not all of the public keys are belong to the one group.
     */
    CryptoElGamalEncrypter(final ElGamalPublicKey elGamalPublicKey,
            final CryptoRandomInteger cryptoRandomInteger,
            final ConfigGroupType configGroupType)
            throws GeneralCryptoLibException {

        _elGamalPublicKey = elGamalPublicKey;
        _elGamalPublicKeyLength = _elGamalPublicKey.getKeys().size();
        _cryptoRandomInteger = cryptoRandomInteger;
        _group = configGroupType.getGroup();

    }

    @Override
    public ElGamalEncrypterValues encryptGroupElements(
            final List<ZpGroupElement> messages)
            throws GeneralCryptoLibException {

        Validate.notNullOrEmptyAndNoNulls(messages,
            "List of mathematical group elements to ElGamal encrypt");
        Validate.notGreaterThan(messages.size(), _elGamalPublicKeyLength,
            "Number of mathematical group elements to ElGamal encrypt",
            "encrypter public key length");

        return encryptGroupElements(messages, false);
    }

    @Override
    public ElGamalEncrypterValues encryptGroupElementsWithShortExponent(
            final List<ZpGroupElement> messages)
            throws GeneralCryptoLibException {

        Validate.notNullOrEmptyAndNoNulls(messages,
            "List of mathematical group elements to ElGamal encrypt");
        Validate.notGreaterThan(messages.size(), _elGamalPublicKeyLength,
            "Number of mathematical group elements to ElGamal encrypt",
            "encrypter public key length");

        return encryptGroupElements(messages, true);
    }

    private ElGamalEncrypterValues encryptGroupElements(
            final List<ZpGroupElement> messages,
            final boolean useShortExponent)
            throws GeneralCryptoLibException {

        ElGamalEncrypterValues preComputeValues =
            preCompute(useShortExponent);

        ElGamalComputationsValues computationsValues =
            compute(messages, preComputeValues.getComputationValues());

        return new ElGamalEncrypterValues(preComputeValues.getR(),
            computationsValues);
    }

    @Override
    public ElGamalEncrypterValues encryptGroupElements(
            final List<ZpGroupElement> messages,
            final ElGamalEncrypterValues preComputedValues)
            throws GeneralCryptoLibException {

        Validate.notNullOrEmptyAndNoNulls(messages,
            "List of mathematical group elements to ElGamal encrypt");
        Validate.notGreaterThan(messages.size(), _elGamalPublicKeyLength,
            "Number of mathematical group elements to ElGamal encrypt",
            "encrypter public key length");
        Validate.notNull(preComputedValues,
            "ElGamal pre-computed values object");

        ElGamalComputationsValues preComputationsValues =
            preComputedValues.getComputationValues();

        Validate.notGreaterThan(preComputationsValues.getPhis().size(),
            _elGamalPublicKey.getKeys().size(),
            "Number of ElGamal pre-computed phi elements ",
            "encrypter public key length");

        ElGamalComputationsValues computationsValues =
            compute(messages, preComputationsValues);

        return new ElGamalEncrypterValues(preComputedValues.getR(),
            computationsValues);
    }

    @Override
    public ElGamalEncrypterValues encryptStrings(
            final List<String> messages) throws GeneralCryptoLibException {

        return encryptGroupElements(getListAsZpGroupElements(messages));
    }

    @Override
    public ElGamalEncrypterValues encryptStrings(
            final List<String> messages,
            final ElGamalEncrypterValues preComputedValues)
            throws GeneralCryptoLibException {

        Validate.notNull(preComputedValues,
            "ElGamal pre-computed values object");

        return encryptGroupElements(getListAsZpGroupElements(messages),
            preComputedValues);
    }

    @Override
    public ElGamalEncrypterValues preCompute() {
        return preCompute(false);
    }

    @Override
    public ElGamalEncrypterValues preComputeWithShortExponent() {
        return preCompute(true);
    }

    private ElGamalEncrypterValues preCompute(
            final boolean useShortExponent) {
        if (useShortExponent) {
            checkShortExponentUsage();
        }

        List<ZpGroupElement> publicKeys = _elGamalPublicKey.getKeys();
        ZpSubgroup group = _elGamalPublicKey.getGroup();

        Exponent randomExponent;
        ZpGroupElement gamma;
        List<ZpGroupElement> prePhis;

        try {
            // Generate a random exponent
            if (useShortExponent) {
                randomExponent =
                    new GroupUtils().getRandomShortExponent(group,
                        _cryptoRandomInteger);
            } else {
                randomExponent =
                    new GroupUtils().getRandomExponent(group,
                        _cryptoRandomInteger);
            }

            // Generate the first element of the ciphertext, gamma=g^random
            gamma = group.getGenerator().exponentiate(randomExponent);

            // For each element in the array of public keys, compute the
            // following:
            // element:prephi[i]=pubKey[i]^(random)
            ZpGroupElement pubKeyRaised;
            prePhis = new ArrayList<>(publicKeys.size());

            for (ZpGroupElement publicKey : publicKeys) {
                pubKeyRaised = publicKey.exponentiate(randomExponent);
                prePhis.add(pubKeyRaised);
            }

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

        try {
            return new ElGamalEncrypterValues(randomExponent,
                new ElGamalComputationsValues(gamma, prePhis));
        } catch (GeneralCryptoLibException e) {
            throw new CryptoLibException(
                "Could not perform pre-computation of ElGamal encryption values.",
                e);
        }
    }

    private void checkShortExponentUsage() throws CryptoLibException {
        if (Group.QR != _group) {
            throw new CryptoLibException(
                "Short Exponents can only be used with Quadratic Residue "
                    + "Groups. Please check policy settings.");
        }
    }

    private ElGamalComputationsValues compute(
            final List<ZpGroupElement> messages,
            final ElGamalComputationsValues computationValues)
            throws GeneralCryptoLibException {

        ElGamalComputationsValues compressedComputationValues =
            getCompressedList(messages.size(), computationValues);

        List<ZpGroupElement> phis = new ArrayList<>(messages.size());

        // For each element in the array of messages, compute the
        // following:
        // element:phi[i]=message[i]*prePhi[i]
        for (int i = 0; i < compressedComputationValues.getPhis().size(); i++) {
            phis.add(messages.get(i).multiply(
                compressedComputationValues.getPhis().get(i)));
        }

        return new ElGamalComputationsValues(
            compressedComputationValues.getGamma(), phis);
    }

    private ElGamalComputationsValues getCompressedList(
            final int numMessages,
            final ElGamalComputationsValues preComputationValues)
            throws GeneralCryptoLibException {

        List<ZpGroupElement> phis = preComputationValues.getPhis();

        if (phis.size() <= numMessages) {
            return preComputationValues;
        }

        GroupElementsCompressor<ZpGroupElement> compressor =
            new GroupElementsCompressor<>();

        List<ZpGroupElement> compressedList =
            compressor.buildListWithCompressedFinalElement(numMessages,
                phis);

        return new ElGamalComputationsValues(
            preComputationValues.getGamma(), compressedList);
    }

    private List<ZpGroupElement> getListAsZpGroupElements(
            final List<String> messages) throws GeneralCryptoLibException {

        Validate
            .notNullOrEmpty(messages,
                "List of stringified mathematical group elements to ElGamal encrypt");
        for (String elementStr : messages) {
            Validate
                .notNullOrBlank(elementStr,
                    "A stringified mathematical group element to ElGamal encrypt");
        }
        Validate.notGreaterThan(messages.size(), _elGamalPublicKeyLength,
            "Number of mathematical group elements to ElGamal encrypt",
            "encrypter public key length");

        List<ZpGroupElement> valuesAsZpGroupElements = new ArrayList<>();
        for (String elementStr : messages) {
            BigInteger parsedBigInteger = parseAsBigInteger(elementStr);
            valuesAsZpGroupElements.add(new ZpGroupElement(
                parsedBigInteger, _elGamalPublicKey.getGroup()));
        }

        return valuesAsZpGroupElements;
    }

    private BigInteger parseAsBigInteger(final String stringToParse) {

        try {
            return new BigInteger(stringToParse);
        } catch (NumberFormatException nfe) {
            throw new CryptoLibException("'" + stringToParse
                + " 'is not a valid representation of a BigInteger.", nfe);
        }
    }
}
