#include <iostream>

#include "client.hpp"

/********************
 * PUBLIC FUNCTIONS *
 ********************/

/*
 * CONSTRUCTORS
 */

PrsonaClient::PrsonaClient(
    const std::vector<Proof>& generatorProof,
    const Twistpoint& elGamalBlindGenerator,
    const BGNPublicKey& serverPublicKey,
    size_t numServers)
    : serverPublicKey(serverPublicKey),
        numServers(numServers),
        max_checked(0)
{
    set_EG_blind_generator(
        generatorProof, elGamalBlindGenerator, numServers);

    longTermPrivateKey.set_random();
    inversePrivateKey = longTermPrivateKey.curveMultInverse();

    decryption_memoizer[elGamalBlindGenerator * max_checked] = max_checked;
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

Twistpoint PrsonaClient::get_short_term_public_key() const
{
    return currentFreshGenerator * longTermPrivateKey;
}

Twistpoint PrsonaClient::get_short_term_public_key(Proof &pi) const
{
    pi = generate_ownership_proof();
    return currentFreshGenerator * longTermPrivateKey;
}

/*
 * SERVER INTERACTIONS
 */

/* Generate a new vote vector to give to the servers
 * @replaces controls which votes are actually being updated and which are not
 *
 * You may really want to make currentEncryptedVotes a member variable, 
 * but it doesn't behave correctly when adding new clients after this one. */
std::vector<TwistBipoint> PrsonaClient::make_votes(
    std::vector<Proof>& validVoteProof,
    const std::vector<Proof>& serverProof,
    const std::vector<TwistBipoint>& oldEncryptedVotes,
    const std::vector<Scalar>& votes,
    const std::vector<bool>& replaces) const
{
    if (!verify_valid_vote_row_proof(serverProof, oldEncryptedVotes))
    {
        std::cerr << "Could not verify proof of valid votes." << std::endl;
        return oldEncryptedVotes;
    }

    std::vector<Scalar> seeds(oldEncryptedVotes.size());
    std::vector<TwistBipoint> newEncryptedVotes(oldEncryptedVotes.size());

    for (size_t i = 0; i < votes.size(); i++)
    {
        if (replaces[i])
        {
            newEncryptedVotes[i] =
                serverPublicKey.twistEncrypt(seeds[i], votes[i]);
        }
        else
        {
            newEncryptedVotes[i] =
                serverPublicKey.rerandomize(seeds[i], oldEncryptedVotes[i]);
        }
    }

    validVoteProof = generate_vote_proof(
        replaces, oldEncryptedVotes, newEncryptedVotes, seeds, votes);
    return newEncryptedVotes;
}

// Get a new fresh generator (happens at initialization and during each epoch)
bool PrsonaClient::receive_fresh_generator(
    const std::vector<Proof>& pi,
    const Twistpoint& freshGenerator)
{
    if (!verify_generator_proof(pi, freshGenerator, numServers))
    {
        std::cerr << "Issue verifying fresh generator proof." << std::endl;
        return false;
    }

    currentFreshGenerator = freshGenerator;
    return true;
}

// Receive a new encrypted score from the servers (each epoch)
bool PrsonaClient::receive_vote_tally(
    const std::vector<Proof>& pi,
    const EGCiphertext& score)
{
    if (!verify_valid_user_tally_proof(pi, score))
    {
        std::cerr << "Could not verify proof of valid tally." << std::endl;
        return false;
    }

    currentEncryptedScore = score;
    currentScore = decrypt_score(score);
    return true;
}

bool PrsonaClient::receive_new_user_data(
    const std::vector<Proof>& mainProof,
    const std::vector<Proof>& serverEncryptedScoreProof,
    const CurveBipoint& serverEncryptedScore,
    const std::vector<Proof>& userEncryptedScoreProof,
    const EGCiphertext& userEncryptedScore,
    const std::vector<Proof>& voteMatrixProof,
    const std::vector<std::vector<TwistBipoint>>& encryptedVoteMatrix,
    const std::vector<Proof>& pseudonymsProof,
    const std::vector<Twistpoint>& currentPseudonyms)
{
    Twistpoint shortTermPublicKey = get_short_term_public_key();

    if (!verify_valid_server_tally_proof(serverEncryptedScoreProof, serverEncryptedScore))
    {
        std::cerr << "Could not verify preliminary proof of server encrypted tally." << std::endl;
        return false;
    }

    if (!verify_valid_user_tally_proof(userEncryptedScoreProof, userEncryptedScore))
    {
        std::cerr << "Could not verify preliminary proof of user encrypted tally." << std::endl;
        return false;
    }

    if (!verify_valid_vote_matrix_proof(voteMatrixProof, encryptedVoteMatrix))
    {
        std::cerr << "Could not verify preliminary proof of encrypted votes." << std::endl;
        return false;
    }

    if (!verify_valid_pseudonyms_proof(pseudonymsProof, currentPseudonyms))
    {
        std::cerr << "Could not verify preliminary proof of pseudonyms." << std::endl;
        return false;
    }

    size_t selfIndex = binary_search(currentPseudonyms, shortTermPublicKey);
    if (currentPseudonyms[selfIndex] != shortTermPublicKey)
    {
        std::cerr << "Was not added to list of pseudonyms." << std::endl;
        return false;
    }

    bool flag = verify_proof_of_added_user(
        mainProof,
        currentFreshGenerator,
        shortTermPublicKey,
        serverPublicKey.get_bipoint_twistgen(),
        serverPublicKey.get_bipoint_twist_subgroup_gen(),
        serverPublicKey.get_bipoint_curvegen(),
        serverPublicKey.get_bipoint_curve_subgroup_gen(),
        selfIndex,
        userEncryptedScore,
        serverEncryptedScore,
        encryptedVoteMatrix);

    if (!flag)
    {
        std::cerr << "There was an issue verifying the proof; this user was not properly added." << std::endl;
        return false;
    }

    currentEncryptedScore = userEncryptedScore;
    currentScore = decrypt_score(userEncryptedScore);

    return true;
}

/*
 * REPUTATION PROOFS
 */

// A pretty straightforward range proof (generation)
std::vector<Proof> PrsonaClient::generate_reputation_proof(
    const Scalar& threshold,
    size_t numClients) const
{
    Proof ownershipProof = generate_ownership_proof();
    
    return PrsonaBase::generate_reputation_proof(
        ownershipProof,
        currentEncryptedScore,
        currentScore,
        threshold,
        inversePrivateKey,
        numClients);
}

bool PrsonaClient::verify_reputation_proof(
    const std::vector<Proof>& pi,
    const Twistpoint& shortTermPublicKey,
    const Scalar& threshold,
    const std::vector<Proof>& encryptedScoreProof,
    const EGCiphertext& encryptedScore) const
{
    if (!verify_valid_user_tally_proof(encryptedScoreProof, encryptedScore))
    {
        std::cerr << "Error getting score from server, aborting." << std::endl;
        return false;
    }

    return PrsonaBase::verify_reputation_proof(
        pi, currentFreshGenerator, shortTermPublicKey, encryptedScore, threshold);
}

Scalar PrsonaClient::get_score() const
{
    return currentScore;
}

/*********************
 * PRIVATE FUNCTIONS *
 *********************/

/*
 * SCORE DECRYPTION
 */

// Basic memoized score decryption
Scalar PrsonaClient::decrypt_score(const EGCiphertext& score)
{
    Twistpoint s, hashedDecrypted;

    // Remove the mask portion of the ciphertext
    s = score.mask * inversePrivateKey;
    hashedDecrypted = score.encryptedMessage - s;
    
    // Check if it's a value we've already seen
    auto lookup = decryption_memoizer.find(hashedDecrypted);
    if (lookup != decryption_memoizer.end())
        return lookup->second;

    // If not, iterate until we find it (adding everything to the memoization)
    max_checked++;
    Twistpoint decryptionCandidate = elGamalBlindGenerator * max_checked;
    while (decryptionCandidate != hashedDecrypted)
    {
        decryption_memoizer[decryptionCandidate] = max_checked;

        decryptionCandidate = decryptionCandidate + elGamalBlindGenerator;
        max_checked++;
    }
    decryption_memoizer[decryptionCandidate] = max_checked;

    // Return the value we found
    return max_checked;
}

/*
 * OWNERSHIP PROOFS
 */

// Prove ownership of the short term public key
Proof PrsonaClient::generate_ownership_proof() const
{
    Twistpoint shortTermPublicKey = currentFreshGenerator * longTermPrivateKey;

    return PrsonaBase::generate_ownership_proof(
        currentFreshGenerator, shortTermPublicKey, longTermPrivateKey);
}

/*
 * VALID VOTE PROOFS
 */

std::vector<Proof> PrsonaClient::generate_vote_proof(
    const std::vector<bool>& replaces,
    const std::vector<TwistBipoint>& oldEncryptedVotes,
    const std::vector<TwistBipoint>& newEncryptedVotes,
    const std::vector<Scalar>& seeds,
    const std::vector<Scalar>& votes) const
{
    Proof pi = generate_ownership_proof();

    return PrsonaBase::generate_vote_proof(
        pi,
        serverPublicKey.get_bipoint_twistgen(),
        serverPublicKey.get_bipoint_twist_subgroup_gen(),
        replaces,
        oldEncryptedVotes,
        newEncryptedVotes,
        seeds,
        votes);
}