/*
 * 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.shares.handler;

import java.nio.charset.StandardCharsets;
import java.security.KeyException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import com.scytl.cryptolib.api.asymmetric.AsymmetricServiceAPI;
import com.scytl.cryptolib.api.elgamal.ElGamalServiceAPI;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.elgamal.bean.ElGamalEncrypterValues;
import com.scytl.cryptolib.elgamal.bean.ElGamalPrivateKey;
import com.scytl.cryptolib.elgamal.bean.ElGamalPublicKey;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalDecrypter;
import com.scytl.cryptolib.elgamal.cryptoapi.CryptoAPIElGamalEncrypter;
import com.scytl.cryptolib.mathematical.groups.impl.ZpGroupElement;
import com.scytl.products.ov.config.multishare.ElGamalPrivateKeySharesService;
import com.scytl.products.ov.config.multishare.MultipleSharesContainer;
import com.scytl.products.ov.config.shares.exception.SharesException;
import com.scytl.products.ov.config.shares.exception.SmartcardException;
import com.scytl.products.ov.config.shares.keys.PrivateKeySerializer;
import com.scytl.products.ov.config.shares.keys.elgamal.ElGamalPrivateKeyAdapter;
import com.scytl.products.ov.config.shares.keys.elgamal.ElGamalPublicKeyAdapter;
import com.scytl.products.ov.config.shares.service.SmartCardService;
import com.scytl.shares.ShamirAlgorithm;
import com.scytl.shares.Share;
import com.scytl.shares.ShareException;

/**
 * Stateless implementation of a read shares handler.
 */
public class StatelessReadSharesHandler {

    private final PrivateKeySerializer privateKeySerializer;

    private final SmartCardService smartcardService;

    private final AsymmetricServiceAPI asymmetricServiceAPI;

    private final ElGamalServiceAPI elGamalServiceAPI;

    public StatelessReadSharesHandler(final PrivateKeySerializer privateKeySerializer,
            final SmartCardService smartcardService, AsymmetricServiceAPI asymmetricService,
            ElGamalServiceAPI elGamalService) {
        this.privateKeySerializer = privateKeySerializer;
        this.smartcardService = smartcardService;
        this.asymmetricServiceAPI = asymmetricService;
        this.elGamalServiceAPI = elGamalService;
    }

    public Share readShare(final String pin, final PublicKey issuerPublicKey) throws SharesException {

        Share share;
        try {
            share = smartcardService.read(pin, issuerPublicKey);
        } catch (SmartcardException e) {
            throw new SharesException("An error occurred while reading the smartcard", e);
        }

        return share;
    }

    public Share readShareElGamal(final String pin, final PublicKey issuerPublicKey) throws SharesException {

        Share share;
        try {
            share = smartcardService.readElGamal(pin, issuerPublicKey);
        } catch (SmartcardException e) {
            throw new SharesException("An error occurred while reading the smartcard", e);
        }

        return share;
    }

    /**
     * Returns a base64 encoded string, which contains the serialized share. The issuer public key is used to verify the
     * signature of the smartcard.
     */
    public String readShareAndStringify(final String pin, final PublicKey issuerPublicKey) throws SharesException {

        Share share = readShare(pin, issuerPublicKey);
        byte[] serializedShare = share.serialize();
        return new String(Base64.getEncoder().encode(serializedShare), StandardCharsets.UTF_8);
    }

    public String readShareAndStringifyElGamal(final String pin, final PublicKey issuerPublicKey)
            throws SharesException {

        Share share = readShareElGamal(pin, issuerPublicKey);
        return Base64.getEncoder().encodeToString(share.serialize());
    }

    public PrivateKey getPrivateKey(final Set<Share> shares, final PublicKey subjectPublicKey) throws SharesException {

        ShamirAlgorithm sa = new ShamirAlgorithm();
        PrivateKey privateKey;
        try {

            byte[] recovered = sa.recover(shares);

            privateKey = privateKeySerializer.reconstruct(recovered, subjectPublicKey);

            validateKeyPair(subjectPublicKey, privateKey);

        } catch (ShareException e) {
            throw new SharesException("There was an error reconstructing the secret from the shares", e);
        } catch (KeyException e) {
            throw new SharesException("There was an error reconstructing the private key from the shares", e);
        }

        return privateKey;
    }

    /**
     * Reconstructs key given a set of serialized shares in base64 encoded format, and the corresponding public key.
     */
    public PrivateKey getPrivateKeyWithSerializedShares(final Set<String> serializedShares,
            final PublicKey subjectPublicKey) throws SharesException {

        final Set<Share> shares = new HashSet<>();

        for (String serializedShare : serializedShares) {

            byte[] shareBytes = Base64.getDecoder().decode(serializedShare);
            try {
                Share share = new Share(shareBytes);
                shares.add(share);
            } catch (ShareException e) {
                throw new SharesException("There was an error while deserializing the shares", e);
            }
        }

        return getPrivateKey(shares, subjectPublicKey);
    }

    public ElGamalPrivateKey getPrivateKeyWithSerializedSharesElGamal(final Set<String> serializedShares,
            final ElGamalPublicKey elGamalPublicKey) throws SharesException {

        final List<MultipleSharesContainer> multipleSharesContainers = new ArrayList<>();

        for (String serializedShare : serializedShares) {

            byte[] shareBytes = Base64.getDecoder().decode(serializedShare.getBytes(StandardCharsets.UTF_8));
            MultipleSharesContainer multipleSharesContainer = null;
            try {
                multipleSharesContainer = new MultipleSharesContainer(shareBytes);
            } catch (ShareException e) {
                throw new SharesException("There was an error while deserializing the shares", e);
            }
            multipleSharesContainers.add(multipleSharesContainer);
        }

        ElGamalPrivateKeySharesService service = new ElGamalPrivateKeySharesService();
        try {
            return service.recover(multipleSharesContainers, elGamalPublicKey);
        } catch (ShareException e) {
            throw new SharesException("Exception while trying to recover share", e);
        }
    }

    public String getSmartcardLabel() throws SharesException {

        try {
            return smartcardService.readSmartcardLabel();
        } catch (SmartcardException e) {
            throw new SharesException("Error while trying to read the smartcard label", e);
        }
    }

    private void validateKeyPair(final PublicKey subjectPublicKey, final PrivateKey privateKey) throws SharesException {

        if ("SCYTL_EL_GAMAL".equals(subjectPublicKey.getAlgorithm())) {

            // The only number that will be member of any given group
            String valueToEncrypt = "1";

            ElGamalPublicKeyAdapter elGamalPublicKeyAdapter = (ElGamalPublicKeyAdapter) subjectPublicKey;
            ElGamalPublicKey elGamalPublicKey = elGamalPublicKeyAdapter.getPublicKey();

            ElGamalPrivateKeyAdapter elGamalPrivateKeyAdapter = (ElGamalPrivateKeyAdapter) privateKey;
            ElGamalPrivateKey elGamalPrivateKey = elGamalPrivateKeyAdapter.getPrivateKey();

            CryptoAPIElGamalEncrypter encrypter = null;
            List<ZpGroupElement> decrypted = null;

            try {
                encrypter = elGamalServiceAPI.createEncrypter(elGamalPublicKey);
                List<String> toEncrypt = new ArrayList<>();
                toEncrypt.add(valueToEncrypt);
                ElGamalEncrypterValues elGamalEncrypterValues = encrypter.encryptStrings(toEncrypt);

                CryptoAPIElGamalDecrypter decrypter = elGamalServiceAPI.createDecrypter(elGamalPrivateKey);
                decrypted = decrypter.decrypt(elGamalEncrypterValues.getComputationValues(), true);

            } catch (GeneralCryptoLibException e) {
                throw new SharesException(
                    "There was an error while validating the reconstructed private key with the given public key", e);
            }
            String decryptedValue = decrypted.get(0).getValue().toString();

            if (!valueToEncrypt.equals(decryptedValue)) {
                throw new SharesException(
                    "There was an error validating the reconstructed private key with the given public key");
            }

        } else {

            final String testString = "foobar";

            byte[] encryptedTestString = null;
            byte[] decrypted = null;

            try {
                encryptedTestString =
                    asymmetricServiceAPI.encrypt(subjectPublicKey, testString.getBytes(StandardCharsets.UTF_8));
                decrypted = asymmetricServiceAPI.decrypt(privateKey, encryptedTestString);
            } catch (GeneralCryptoLibException e) {
                throw new SharesException(
                    "There was an error while validating the reconstructed private key with the given public key", e);
            }
            String decryptedTestString = new String(decrypted, StandardCharsets.UTF_8);

            if (!testString.equals(decryptedTestString)) {
                throw new SharesException(
                    "There was an error validating the reconstructed private key with the given public key");
            }
        }

    }

}
