#include <iostream>

#include "client.hpp"
#include "serverEntity.hpp"

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

/*
 * CONSTRUCTORS
 */

PrsonaClient::PrsonaClient(
    const std::vector<Proof>& generatorProof,
    const Curvepoint& elGamalBlindGenerator,
    const BGNPublicKey& serverPublicKey,
    const PrsonaServerEntity* servers)
    : serverPublicKey(serverPublicKey),
        servers(servers),
        max_checked(0)
{
    set_EG_blind_generator(
        generatorProof, elGamalBlindGenerator, servers->get_num_servers());

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

    decryption_memoizer[elGamalBlindGenerator * max_checked] = max_checked;
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

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

Curvepoint 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<CurveBipoint> PrsonaClient::make_votes(
    std::vector<Proof>& validVoteProof,
    const Proof& serverProof,
    const std::vector<CurveBipoint>& oldEncryptedVotes,
    const std::vector<Scalar>& votes,
    const std::vector<bool>& replaces) const
{
    std::vector<Scalar> seeds(oldEncryptedVotes.size());
    std::vector<CurveBipoint> newEncryptedVotes(oldEncryptedVotes.size());

    if (!verify_valid_vote_row_proof(serverProof))
    {
        std::cerr << "Could not verify proof of valid votes." << std::endl;
        return newEncryptedVotes;
    }

    for (size_t i = 0; i < votes.size(); i++)
    {
        if (replaces[i])
        {
            newEncryptedVotes[i] =
                serverPublicKey.curveEncrypt(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 Curvepoint& freshGenerator)
{
    if (!verify_generator_proof(pi, freshGenerator, servers->get_num_servers()))
        return false;

    currentFreshGenerator = freshGenerator;
    return true;
}

// Receive a new encrypted score from the servers (each epoch)
bool PrsonaClient::receive_vote_tally()
{
    Proof pi;
    Curvepoint shortTermPublicKey = get_short_term_public_key();
    EGCiphertext score =
        servers->get_current_user_encrypted_tally(pi, shortTermPublicKey);

    if (!verify_valid_user_tally_proof(pi))
    {
        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)
{
    Proof currProof;
    Curvepoint shortTermPublicKey = get_short_term_public_key();
    
    EGCiphertext userEncryptedScore =
        servers->get_current_user_encrypted_tally(currProof, shortTermPublicKey);
    if (!verify_valid_user_tally_proof(currProof))
    {
        std::cerr << "Could not verify preliminary proof of user encrypted tally." << std::endl;
        return false;
    }

    TwistBipoint serverEncryptedScore =
        servers->get_current_server_encrypted_tally(currProof, shortTermPublicKey);
    if (!verify_valid_server_tally_proof(currProof))
    {
        std::cerr << "Could not verify preliminary proof of server encrypted tally." << std::endl;
        return false;
    }

    std::vector<std::vector<CurveBipoint>> encryptedVoteMatrix = 
        servers->get_all_current_votes(currProof);
    if (!verify_valid_vote_matrix_proof(currProof))
    {
        std::cerr << "Could not verify preliminary proof of encrypted votes." << std::endl;
        return false;
    }

    std::vector<Curvepoint> currentPseudonyms = 
        servers->get_current_pseudonyms(currProof);
    if (!verify_valid_pseudonyms_proof(currProof))
    {
        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,
        elGamalBlindGenerator,
        serverPublicKey.get_bipoint_curvegen(),
        serverPublicKey.get_bipoint_curve_subgroup_gen(),
        serverPublicKey.get_bipoint_twistgen(),
        serverPublicKey.get_bipoint_twist_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) const
{
    Proof ownershipProof = generate_ownership_proof();
    
    return PrsonaBase::generate_reputation_proof(
        ownershipProof,
        currentEncryptedScore,
        currentScore,
        threshold,
        inversePrivateKey,
        servers->get_num_clients());
}

bool PrsonaClient::verify_reputation_proof(
    const std::vector<Proof>& pi,
    const Curvepoint& shortTermPublicKey,
    const Scalar& threshold) const
{
    Proof serverProof;
    EGCiphertext encryptedScore =
        servers->get_current_user_encrypted_tally(serverProof, shortTermPublicKey);

    if (!verify_valid_user_tally_proof(serverProof))
    {
        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)
{
    Curvepoint 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++;
    Curvepoint 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
{
    Curvepoint 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<CurveBipoint>& oldEncryptedVotes,
    const std::vector<CurveBipoint>& 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_curvegen(),
        serverPublicKey.get_bipoint_curve_subgroup_gen(),
        replaces,
        oldEncryptedVotes,
        newEncryptedVotes,
        seeds,
        votes);
}