#include <iostream>

#include "serverEntity.hpp"

/********************
 * 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(numServers);
    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(numServers, sharedBGN));

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

    pi.clear();

    // It's important that no server knows the DLOG between g and h for ElGamal,
    // so have each server collaborate to make h.
    Curvepoint blindGenerator = PrsonaServer::EL_GAMAL_GENERATOR;
    for (size_t i = 0; i < numServers; i++)
    {
        blindGenerator =
            servers[i].add_rand_seed_to_generator(pi, blindGenerator);
    }
    for (size_t i = 0; i < numServers; i++)
        servers[i].set_EG_blind_generator(pi, blindGenerator);
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

BGNPublicKey PrsonaServerEntity::get_bgn_public_key() const
{
    return get_bgn_public_key(0);
}

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

Curvepoint PrsonaServerEntity::get_blinding_generator() const
{
    return get_blinding_generator(0);
}

Curvepoint PrsonaServerEntity::get_blinding_generator(size_t which) const
{
    std::vector<Proof> pi;
    Curvepoint retval = get_blinding_generator(pi, which);

    if (!servers[which].verify_generator_proof(
            pi, retval, servers[which].get_num_servers()))
    {
        std::cerr << "Error making the generator, aborting." << std::endl;
        return Curvepoint();
    }

    return retval;
}

Curvepoint PrsonaServerEntity::get_blinding_generator(
    std::vector<Proof>& pi) const
{
    return get_blinding_generator(pi, 0);
}

Curvepoint PrsonaServerEntity::get_blinding_generator(
    std::vector<Proof>& pi, size_t which) const
{
    return servers[which].get_blinding_generator(pi);
}

Curvepoint PrsonaServerEntity::get_fresh_generator() const
{
    return get_fresh_generator(0);
}

Curvepoint PrsonaServerEntity::get_fresh_generator(size_t which) const
{
    std::vector<Proof> pi;
    Curvepoint retval = get_fresh_generator(pi, which);

    if (!servers[which].verify_generator_proof(
            pi, retval, servers[which].get_num_servers()))
    {
        std::cerr << "Error making the generator, aborting." << std::endl;
        return Curvepoint();
    }

    return retval;
}

Curvepoint PrsonaServerEntity::get_fresh_generator(
    std::vector<Proof>& pi) const
{
    return get_fresh_generator(pi, 0);
}

Curvepoint PrsonaServerEntity::get_fresh_generator(
    std::vector<Proof>& pi, size_t which) const
{
    Curvepoint retval = PrsonaServer::EL_GAMAL_GENERATOR;

    pi.clear();   
    for (size_t j = 0; j < servers[which].get_num_servers(); j++)
    {
        size_t index = (which + j) % servers[which].get_num_servers();
        retval = servers[index].add_curr_seed_to_generator(pi, retval);
    }

    return retval;
}

size_t PrsonaServerEntity::get_num_clients() const
{
    return get_num_clients(0);
}

size_t PrsonaServerEntity::get_num_clients(size_t which) const
{
    return servers[which].get_num_clients();
}

size_t PrsonaServerEntity::get_num_servers() const
{
    return get_num_servers(0);
}

size_t PrsonaServerEntity::get_num_servers(size_t which) const
{
    return servers[which].get_num_servers();
}

/*
 * ENCRYPTED DATA GETTERS
 */

std::vector<CurveBipoint> PrsonaServerEntity::get_current_votes_by(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    return get_current_votes_by(pi, shortTermPublicKey, 0);
}

/* 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, size_t which) const
{
    return servers[which].get_current_votes_by(pi, shortTermPublicKey);
}

std::vector<std::vector<CurveBipoint>> PrsonaServerEntity::get_all_current_votes(
    Proof& pi) const
{
    return get_all_current_votes(pi, 0);
}

std::vector<std::vector<CurveBipoint>> PrsonaServerEntity::get_all_current_votes(
    Proof& pi, size_t which) const
{
    return servers[which].get_all_current_votes(pi);
}

EGCiphertext PrsonaServerEntity::get_current_user_encrypted_tally(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    return get_current_user_encrypted_tally(pi, shortTermPublicKey, 0);
}

EGCiphertext PrsonaServerEntity::get_current_user_encrypted_tally(
    Proof& pi, const Curvepoint& shortTermPublicKey, size_t which) const
{
    return servers[which].get_current_user_encrypted_tally(
        pi, shortTermPublicKey);
}

TwistBipoint PrsonaServerEntity::get_current_server_encrypted_tally(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    return get_current_server_encrypted_tally(pi, shortTermPublicKey, 0);
}

TwistBipoint PrsonaServerEntity::get_current_server_encrypted_tally(
    Proof& pi, const Curvepoint& shortTermPublicKey, size_t which) const
{
    return servers[which].get_current_server_encrypted_tally(
        pi, shortTermPublicKey);
}

std::vector<Curvepoint> PrsonaServerEntity::get_current_pseudonyms(
    Proof& pi) const
{
    return get_current_pseudonyms(pi, 0);
}

std::vector<Curvepoint> PrsonaServerEntity::get_current_pseudonyms(
    Proof& pi, size_t which) const
{
    return servers[which].get_current_pseudonyms(pi);
}

/*
 * CLIENT INTERACTIONS
 */

void PrsonaServerEntity::add_new_client(PrsonaClient& newUser)
{
    add_new_client(newUser, 0);
}

/* 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, size_t which)
{
    Proof proofOfValidSTPK, proofOfValidVotes;
    std::vector<Proof> proofOfValidGenerator, proofOfCorrectAddition;
    Curvepoint freshGenerator =
        get_fresh_generator(proofOfValidGenerator, which);

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

    // Do the actual work of adding the client to the first server
    servers[which].add_new_client(
        proofOfCorrectAddition, proofOfValidSTPK, 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<std::vector<CurveBipoint>> voteMatrix;
    servers[which].export_updates(
        previousVoteTally,
        currentPseudonyms,
        currentUserEncryptedTallies,
        voteMatrix);
    
    // FIX ME!!!!!!!
    Proof unused("PROOF");

    for (size_t j = 1; j < servers[which].get_num_servers(); j++)
    {
        size_t index = (which + j) % servers[which].get_num_servers();
        servers[index].import_updates(
            unused,
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            voteMatrix);
    }

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

// Receive a new vote row from a user (identified by short term public key).
bool PrsonaServerEntity::receive_vote(
    const std::vector<Proof>& pi,
    const std::vector<CurveBipoint>& newVotes,
    const Curvepoint& shortTermPublicKey)
{
    return receive_vote(pi, newVotes, shortTermPublicKey, 0);
}

bool PrsonaServerEntity::receive_vote(
    const std::vector<Proof>& pi,
    const std::vector<CurveBipoint>& newVotes,
    const Curvepoint& shortTermPublicKey,
    size_t which)
{
    bool retval = true;

    for (size_t i = 0; i < servers[which].get_num_servers(); i++)
    {
        size_t index = (i + which) % servers[which].get_num_servers();

        retval = retval &&
            servers[index].receive_vote(pi, newVotes, shortTermPublicKey);
    }

    return retval;
}

void PrsonaServerEntity::transmit_new_user_data(
    const std::vector<Proof>& pi, PrsonaClient& newUser) const
{
    newUser.receive_new_user_data(pi);
}

void PrsonaServerEntity::transmit_updates(PrsonaClient& currUser) const
{
    transmit_updates(currUser, 0);
}

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

    // Get users the next fresh generator so they can correctly
    // ask for their new scores and vote row
    currUser.receive_fresh_generator(proofOfValidGenerator, freshGenerator);

    currUser.receive_vote_tally();
}

/*
 * EPOCH
 */

void PrsonaServerEntity::epoch(Proof& pi)
{
    epoch(pi, 0);
}

// Do the epoch process
void PrsonaServerEntity::epoch(Proof& pi, size_t which)
{
    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[which].get_num_servers(); i++)
    {
        size_t index = (which + i) % servers[which].get_num_servers();
        size_t nextIndex = (which + i + 1) % servers[which].get_num_servers();

        servers[index].build_up_midway_pseudonyms(pi, nextGenerator);
        servers[index].export_updates(
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            voteMatrix);
        servers[nextIndex].import_updates(
            pi,
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            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(nextGenerator, which);
    
    // go from A_0.5 to A_1
    for (size_t i = 0; i < servers[which].get_num_servers(); i++)
    {
        size_t index = (which + i) % servers[which].get_num_servers();
        size_t nextIndex = (which + i + 1) % servers[which].get_num_servers();
        
        servers[index].break_down_midway_pseudonyms(pi, nextGenerator);
        servers[index].export_updates(
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            voteMatrix);
        servers[nextIndex].import_updates(
            pi,
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            voteMatrix);
    }

    // At the end, make sure all servers have same information
    for (size_t i = 1; i < servers[which].get_num_servers() - 1; i++)
    {
        size_t index = (which + i) % servers[which].get_num_servers();

        servers[index].import_updates(
            pi,
            previousVoteTally,
            currentPseudonyms,
            currentUserEncryptedTallies,
            voteMatrix);
    }
}

void PrsonaServerEntity::print_scores() const
{
    BGN bgnSystem = servers[0].bgnSystem;
    std::vector<TwistBipoint> scores = servers[0].previousVoteTallies;

    std::cout << "[";
    for (size_t i = 0; i < scores.size(); i++)
    {
            std::cout << bgnSystem.decrypt(scores[i])
                << (i == scores.size() - 1 ? "]" : " ");
    }
    std::cout << std::endl;
}

void PrsonaServerEntity::print_votes() const
{
    BGN bgnSystem = servers[0].bgnSystem;
    std::vector<std::vector<CurveBipoint>> voteMatrix = servers[0].voteMatrix;

    for (size_t i = 0; i < voteMatrix.size(); i++)
    {
        std::cout << (i == 0 ? "[[" : " [");
        for (size_t j = 0; j < voteMatrix[i].size(); j++)
        {

            std::cout << bgnSystem.decrypt(voteMatrix[i][j])
                << (j == voteMatrix[i].size() - 1 ? "]" : " ");
        }
        std::cout << (i == voteMatrix.size() - 1 ? "]" : " ") << std::endl;
    }
}

/*********************
 * 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(
    const Curvepoint& nextGenerator,
    size_t which)
{
    std::vector<EGCiphertext> retval;
    Proof maxScoreProof;
    std::vector<Scalar> decryptedTalliedScores = servers[which].tally_scores();
    mpz_class maxScorePossibleThisRound =
        servers[which].get_max_possible_score(maxScoreProof).toInt() *
        PrsonaBase::get_max_allowed_vote();
    mpz_class topOfScoreRange =
        decryptedTalliedScores.size() * PrsonaBase::get_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[which].bgnSystem.encrypt(
            servers[which].previousVoteTallies[i], decryptedTalliedScores[i]);
        
        retval[i].mask = servers[which].currentPseudonyms[i] * currMask;
        retval[i].encryptedMessage =
            (nextGenerator * currMask) +
            (servers[which].get_blinding_generator() * decryptedTalliedScores[i]);
    }

    servers[which].currentUserEncryptedTallies = retval;
    return retval;
}

/*
 * BINARY SEARCH
 */

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