#include <iostream>

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

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

/*
 * CONSTRUCTORS
 */

PrsonaClient::PrsonaClient(
    const BGNPublicKey& serverPublicKey,
    const std::vector<Proof>& generatorProof,
    const Curvepoint& elGamalBlindGenerator,
    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(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_votes_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.encrypt(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)
void PrsonaClient::receive_vote_tally(
    const Proof& pi, const EGCiphertext& score)
{
    if (!verify_valid_tally_proof(pi))
    {
        std::cerr << "Could not verify proof of valid tally." << std::endl;
        return;
    }

    currentEncryptedScore = score;
    decrypt_score(score);
}

/*
 * 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_tally(serverProof, shortTermPublicKey);

    if (!verify_valid_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
void 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())
    {
        currentScore = lookup->second;
        return;
    }

    // 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;

    // Set the value we found
    currentScore = 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);
}