#include <iostream>

#include "serverEntity.hpp"

const int MAX_ALLOWED_VOTE = 2;

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

/*
 * CONSTRUCTORS
 */

PrsonaServerEntity::PrsonaServerEntity(size_t numServers)
{
    if (numServers < 1)
    {
        std::cerr << "You have to have at least 1 server. "
            << "I'm making it anyways." << std::endl;
    }

    // Make the first server, which makes the BGN parameters
    PrsonaServer firstServer;
    servers.push_back(firstServer);

    // Make the rest of the servers, which take the BGN parameters
    const BGN& sharedBGN = firstServer.get_bgn_details();
    for (size_t i = 1; i < numServers; i++)
        servers.push_back(PrsonaServer(sharedBGN));

    // After all servers have made their seeds,
    // make sure they have the initial fresh generator
    Curvepoint firstGenerator = get_fresh_generator();
    for (size_t i = 0; i < numServers; i++)
        servers[i].initialize_fresh_generator(firstGenerator);
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

BGNPublicKey PrsonaServerEntity::get_bgn_public_key() const
{
    return servers[0].get_bgn_public_key();
}

Curvepoint PrsonaServerEntity::get_blinding_generator() const
{
    return servers[0].get_blinding_generator();
}

Curvepoint PrsonaServerEntity::get_fresh_generator() const
{
    Curvepoint retval = PrsonaServer::EL_GAMAL_GENERATOR;
    for (size_t j = 0; j < servers.size(); j++)
        retval = servers[j].add_curr_seed_to_generator(retval);

    return retval;
}

/*
 * ENCRYPTED DATA GETTERS
 */

/* Call this in order to get the current encrypted votes cast by a given user
 * (who is identified by their short term public key).
 * In practice, this is intended for clients,
 * who need to know their current votes in order to rerandomize them. */
std::vector<CurveBipoint> PrsonaServerEntity::get_current_votes_by(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    return servers[0].get_current_votes_by(pi, shortTermPublicKey);
}

/* Call this in order to get the current encrypted tally of a given user
 * (who is identified by their short term public key).
 * In practice, this is intended for clients, so that the servers vouch
 * for their ciphertexts being valid as part of their reputation proofs. */
EGCiphertext PrsonaServerEntity::get_current_tally(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    return servers[0].get_current_tally(pi, shortTermPublicKey);
}

/*
 * CLIENT INTERACTIONS
 */

/* Add a new client (who is identified only by their short term public key)
 * One server does the main work, then other servers import their (proven)
 * exported data. */
void PrsonaServerEntity::add_new_client(PrsonaClient& newUser)
{
    Proof proofOfValidSTPK, proofOfCorrectAddition, proofOfValidVotes;
    Curvepoint freshGenerator = get_fresh_generator();

    // Users can't actually announce a short term public key
    // if they don't know the fresh generator.
    newUser.receive_fresh_generator(freshGenerator);
    Curvepoint shortTermPublicKey = newUser.get_short_term_public_key(
                                        proofOfValidSTPK);

    // Do the actual work of adding the client to the first server
    servers[0].add_new_client(
        proofOfValidSTPK, proofOfCorrectAddition, shortTermPublicKey);

    // Then, export the data to the rest of the servers
    std::vector<TwistBipoint> previousVoteTally;
    std::vector<Curvepoint> currentPseudonyms;
    std::vector<EGCiphertext> currentUserEncryptedTallies;
    std::vector<Proof> currentTallyProofs;
    std::vector<std::vector<CurveBipoint>> voteMatrix;
    servers[0].export_updates(
        previousVoteTally,
        currentPseudonyms,
        currentUserEncryptedTallies,
        currentTallyProofs,
        voteMatrix);
    for (size_t j = 1; j < servers.size(); j++)
    {
        servers[j].import_updates(
            proofOfCorrectAddition,
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            currentTallyProofs,
            voteMatrix);
    }

    // Finally, give the user the information it needs
    // about its current tally and votes
    transmit_updates(newUser);
}

// Receive a new vote row from a user (identified by short term public key).
void PrsonaServerEntity::receive_vote(
    const Proof& pi,
    const std::vector<CurveBipoint>& votes,
    const Curvepoint& shortTermPublicKey)
{
    for (size_t j = 0; j < servers.size(); j++)
        servers[j].receive_vote(pi, votes, shortTermPublicKey);
}

// After tallying scores and new vote matrix,
// give those to a user for the new epoch
void PrsonaServerEntity::transmit_updates(PrsonaClient& currUser) const
{
    Proof proofOfValidSTPK, proofOfScore, proofOfCorrectVotes;
    Curvepoint freshGenerator = get_fresh_generator();

    // Get users the next fresh generator so they can correctly
    // ask for their new scores and vote row
    currUser.receive_fresh_generator(freshGenerator);
    Curvepoint shortTermPublicKey = currUser.get_short_term_public_key(
                                        proofOfValidSTPK);

    EGCiphertext score = get_current_tally(proofOfScore, shortTermPublicKey);
    currUser.receive_vote_tally(proofOfScore, score);

    std::vector<CurveBipoint> encryptedVotes = get_current_votes_by(
            proofOfCorrectVotes, shortTermPublicKey);
    currUser.receive_encrypted_votes(proofOfCorrectVotes, encryptedVotes);
}

/*
 * EPOCH
 */

// Do the epoch process
void PrsonaServerEntity::epoch(Proof& pi)
{
    Curvepoint nextGenerator = PrsonaServer::EL_GAMAL_GENERATOR;
    
    std::vector<TwistBipoint> previousVoteTally;
    std::vector<Curvepoint> currentPseudonyms;
    std::vector<EGCiphertext> currentUserEncryptedTallies;
    std::vector<Proof> currentTallyProofs;
    std::vector<std::vector<CurveBipoint>> voteMatrix;

    // go from A_0 to A_0.5
    for (size_t i = 0; i < servers.size(); i++)
    {
        servers[i].build_up_midway_pseudonyms(pi, nextGenerator);
        servers[i].export_updates(
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            currentTallyProofs,
            voteMatrix);
        if (i < servers.size() - 1)
        {
            servers[i + 1].import_updates(
                pi,
                previousVoteTally,
                currentPseudonyms,
                currentUserEncryptedTallies,
                currentTallyProofs,
                voteMatrix);
        }
    }

    /* Imagine that server 0 is encrypting these, then would do a ZKP that it
     * knows a secret mask and encrypted the correct value everyone else already
     * knows. Everyone else then adds a mask and proves they added a secret mask
     * to the committed values. */
    currentUserEncryptedTallies = tally_scores(currentTallyProofs, nextGenerator);
    servers[0].import_updates(
        pi,
        previousVoteTally,
        currentPseudonyms,
        currentUserEncryptedTallies,
        currentTallyProofs,
        voteMatrix);
    
    // go from A_0.5 to A_1
    for (size_t i = 0; i < servers.size(); i++)
    {
        servers[i].break_down_midway_pseudonyms(pi, nextGenerator);
        servers[i].export_updates(
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            currentTallyProofs,
            voteMatrix);
        if (i < servers.size() - 1)
        {
            servers[i + 1].import_updates(
                pi,
                previousVoteTally,
                currentPseudonyms,
                currentUserEncryptedTallies,
                currentTallyProofs,
                voteMatrix);
        }
    }

    // At the end, make sure all servers have same information
    for (size_t i = 0; i < servers.size() - 1; i++)
    {
        servers[i].import_updates(
            pi,
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            currentTallyProofs,
            voteMatrix);
    }
}

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

/*
 * SCORE TALLYING
 */

/* Calculate scores, then scale the values appropriately,
 * so they are in the correct range to be used as vote weights.
 *
 * We're treating it as if we are one server, so that server gets the updated
 * weights to be sent to all other servers for the next epoch. */
std::vector<EGCiphertext> PrsonaServerEntity::tally_scores(
    std::vector<Proof>& tallyProofs, const Curvepoint& nextGenerator)
{
    std::vector<EGCiphertext> retval;
    Proof maxScoreProof;
    std::vector<Scalar> decryptedTalliedScores = servers[0].tally_scores(
                                                    tallyProofs);
    mpz_class maxScorePossibleThisRound =
        servers[0].get_max_possible_score(maxScoreProof).toInt() *
        MAX_ALLOWED_VOTE;
    mpz_class topOfScoreRange = retval.size() * MAX_ALLOWED_VOTE;

    for (size_t i = 0; i < decryptedTalliedScores.size(); i++)
    {
        decryptedTalliedScores[i] =
            Scalar(
                (decryptedTalliedScores[i].toInt() * topOfScoreRange) /
                maxScorePossibleThisRound
            );

        EGCiphertext currCiphertext;
        retval.push_back(currCiphertext);
        Scalar currMask;
        currMask.set_random();

        // Give the server the new weights,
        // to get passed around to the other servers
        servers[0].bgn_system.encrypt(
            servers[0].previousVoteTallies[i], decryptedTalliedScores[i]);
        
        retval[i].mask = servers[0].currentPseudonyms[i] * currMask;
        retval[i].encryptedMessage =
            (nextGenerator * currMask) +
            (PrsonaServer::get_blinding_generator() * decryptedTalliedScores[i]);
    }

    return retval;
}

/*
 * BINARY SEARCH
 */

// Completely normal binary search
size_t PrsonaServerEntity::binary_search(
    const Curvepoint& shortTermPublicKey) const
{
    return servers[0].binary_search(shortTermPublicKey);
}