/*
 * 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.config.multishare;

import java.math.BigInteger;
import java.security.KeyException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.mathematical.groups.impl.Exponent;
import com.scytl.shares.ShamirAlgorithm;
import com.scytl.shares.Share;
import com.scytl.shares.ShareException;

/**
 * Allows for an ElGamal privatekey to be split into a number of shares, and the privatekey to be recovered again later.
 * <P>
 * Internally, it uses the Shamir algorithm to perform the splitting.
 */
public class ElGamalPrivateKeySharesService {

    /**
     * Split an ElGamal privatekey into a number of shares.
     * <P>
     * It is important to realize that an ElGamal key is composed of a number of "subkeys" (these subkeys can also be
     * thought of as keys in their own right).
     * <P>
     * This method performs a split on EACH of the subkeys separately, it then groups together the corresponding shares
     * from each of these split operations. For example, if the received ElGamal privatekey contains 5 subkeys, and if
     * the specified number of shares is 2, then each of the subkeys will be split into two shares, and then all of the
     * shares with index '0' will be grouped into one container, and all of the shares with index '1' will be grouped
     * into a container.
     * 
     * @param privateKey
     *            the key to be split into shares.
     * @param numberShares
     *            the number of shares to be generated. The list returned by this method will contain this number of
     *            elements.
     * @param threshold
     *            the minimum number of shares required to be present to recover the key.
     * @return a List of containers (containing the same number of shares as subkeys in the key).
     */
    public List<MultipleSharesContainer> split(final ElGamalPrivateKey privateKey, final int numberShares,
            final int threshold) throws KeyException, ShareException {

        validateSplitInputs(privateKey, numberShares, threshold);

        List<Exponent> subkeys = privateKey.getKeys();
        ShamirAlgorithm sa = new ShamirAlgorithm();

        // split each subkey into shares
        List<List<Share>> allShares = new ArrayList<List<Share>>();

        for (int i = 0; i < subkeys.size(); i++) {
            byte[] secret = subkeys.get(i).getValue().toByteArray();
            Set<Share> shareForSubkey = sa.split(secret, numberShares, threshold);
            allShares.add(new ArrayList<Share>(shareForSubkey));
        }

        // now arrange data into the containers that will be serialized
        int numSubkeys = subkeys.size();
        List<MultipleSharesContainer> containers = new ArrayList<MultipleSharesContainer>();

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

            List<byte[]> sharesForCard = new ArrayList<byte[]>();
            for (int j = 0; j < numSubkeys; j++) {
                sharesForCard.add(allShares.get(j).get(i).serialize());
            }
            containers.add(new MultipleSharesContainer(numberShares, threshold, sharesForCard));
        }

        return containers;
    }

    /**
     * Recover an ElGamal privatekey.
     * 
     * @param containers
     *            the list of containers (containing one or more shares).
     * @param elGamalPublicKey
     *            the publickey that corresponds to the private key.
     * @return the recovered ElGamal privatekey.
     */
    public ElGamalPrivateKey recover(List<MultipleSharesContainer> containers, ElGamalPublicKey elGamalPublicKey)
            throws ShareException {

        validateRecoverInputs(containers, elGamalPublicKey);

        int numberShares = containers.size();
        int numSubkeys = containers.get(0).getShares().size();
        ShamirAlgorithm sa = new ShamirAlgorithm();

        List<Exponent> recoveredListOfSubkeys = new ArrayList<Exponent>();

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

            Set<Share> sharesForRecoveringParticularSubkey = new HashSet<Share>();
            for (int j = 0; j < numberShares; j++) {

                List<byte[]> allSharesFromContainer = containers.get(j).getShares();
                sharesForRecoveringParticularSubkey.add(new Share(allSharesFromContainer.get(i)));
            }

            byte[] recoveredData = sa.recover(sharesForRecoveringParticularSubkey);
            Exponent recoveredSubkey;
            try {
                recoveredSubkey = new Exponent(elGamalPublicKey.getGroup().getQ(), new BigInteger(recoveredData));
            } catch (GeneralCryptoLibException e) {
                throw new ShareException("Exception while trying to reconstruct Privatekey Exponent", e);
            }
            recoveredListOfSubkeys.add(recoveredSubkey);
        }

        try {
            return new ElGamalPrivateKey(recoveredListOfSubkeys, elGamalPublicKey.getGroup());
        } catch (GeneralCryptoLibException e) {
            throw new ShareException("Exception while trying to create new ElGamalPrivateKey", e);
        }
    }

    private void validateSplitInputs(ElGamalPrivateKey privateKey, int numberShares, int threshold)
            throws ShareException {

        if (privateKey == null) {
            throw new ShareException("The received ElGamal private key was null");
        }

        if (numberShares < 1) {
            throw new ShareException("The number of shares should be at least 1, but it was: " + numberShares);
        }

        if (threshold < 1) {
            throw new ShareException("The threshold should be at least 1, but it was: " + threshold);
        }
    }

    private void validateRecoverInputs(List<MultipleSharesContainer> containers, ElGamalPublicKey elGamalPublicKey)
            throws ShareException {

        if (containers == null) {
            throw new ShareException("The received list of containers was null");
        } else if (containers.isEmpty()) {
            throw new ShareException("The received list of containers was empty");
        }

        if (elGamalPublicKey == null) {
            throw new ShareException("The received ElGamal public key was null");
        }
    }
}
