#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
    servers.push_back(PrsonaServer(numServers));

    // Make the rest of the servers, which take the BGN parameters
    const BGN& sharedBGN = servers[0].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;
    Twistpoint 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.
    Twistpoint 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();
}

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

Twistpoint PrsonaServerEntity::get_blinding_generator(
    size_t which) const
{
    std::vector<Proof> pi;
    Twistpoint 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 Twistpoint();
    }

    return retval;
}

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

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

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

Twistpoint PrsonaServerEntity::get_fresh_generator(
    size_t which) const
{
    std::vector<Proof> pi;
    Twistpoint 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 Twistpoint();
    }

    return retval;
}

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

Twistpoint PrsonaServerEntity::get_fresh_generator(
    std::vector<Proof>& pi,
    size_t which) const
{
    Twistpoint 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<TwistBipoint> PrsonaServerEntity::get_current_votes_by(
    Proof& pi,
    const Twistpoint& 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<TwistBipoint> PrsonaServerEntity::get_current_votes_by(
    Proof& pi,
    const Twistpoint& shortTermPublicKey,
    size_t which) const
{
    return servers[which].get_current_votes_by(pi, shortTermPublicKey);
}

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

std::vector<std::vector<TwistBipoint>> 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 Twistpoint& shortTermPublicKey) const
{
    return get_current_user_encrypted_tally(pi, shortTermPublicKey, 0);
}

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

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

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

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

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

/*
 * PROOF COMMITMENT GETTERS
 */

void PrsonaServerEntity::get_other_vote_row_commitments(
    std::vector<Proof>& pi,
    const Twistpoint& request) const
{
    get_other_vote_row_commitments(pi, request, 0);
}

void PrsonaServerEntity::get_other_vote_row_commitments(
    std::vector<Proof>& pi,
    const Twistpoint& request,
    size_t whichNot) const
{
    for (size_t i = 0; i < servers.size(); i++)
    {
        if (i == whichNot)
            continue;

        pi.push_back(servers[i].get_vote_row_commitment(request));
    }
}

void PrsonaServerEntity::get_other_vote_matrix_commitments(
    std::vector<Proof>& pi) const
{
    get_other_vote_matrix_commitments(pi, 0);
}

void PrsonaServerEntity::get_other_vote_matrix_commitments(
    std::vector<Proof>& pi,
    size_t whichNot) const
{
    for (size_t i = 0; i < servers.size(); i++)
    {
        if (i == whichNot)
            continue;

        pi.push_back(servers[i].get_vote_matrix_commitment());
    }
}

void PrsonaServerEntity::get_other_user_tally_commitments(
    std::vector<Proof>& pi,
    const Twistpoint& request) const
{
    get_other_user_tally_commitments(pi, request, 0);
}

void PrsonaServerEntity::get_other_user_tally_commitments(
    std::vector<Proof>& pi,
    const Twistpoint& request,
    size_t whichNot) const
{
    for (size_t i = 0; i < servers.size(); i++)
    {
        if (i == whichNot)
            continue;

        pi.push_back(servers[i].get_user_tally_commitment(request));
    }
}

void PrsonaServerEntity::get_other_server_tally_commitments(
    std::vector<Proof>& pi,
    const Twistpoint& request) const
{
    get_other_server_tally_commitments(pi, request, 0);
}

void PrsonaServerEntity::get_other_server_tally_commitments(
    std::vector<Proof>& pi,
    const Twistpoint& request,
    size_t whichNot) const
{
    for (size_t i = 0; i < servers.size(); i++)
    {
        if (i == whichNot)
            continue;

        pi.push_back(servers[i].get_server_tally_commitment(request));
    }
}

void PrsonaServerEntity::get_other_pseudonyms_commitments(
    std::vector<Proof>& pi) const
{
    get_other_pseudonyms_commitments(pi, 0);
}

void PrsonaServerEntity::get_other_pseudonyms_commitments(
    std::vector<Proof>& pi,
    size_t whichNot) const
{
    for (size_t i = 0; i < servers.size(); i++)
    {
        if (i == whichNot)
            continue;

        pi.push_back(servers[i].get_pseudonyms_commitment());
    }
}

/*
 * 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;
    Twistpoint 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);
    Twistpoint 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<CurveBipoint> previousVoteTallies;
    std::vector<Twistpoint> currentPseudonyms;
    std::vector<EGCiphertext> currentUserEncryptedTallies;
    std::vector<std::vector<TwistBipoint>> voteMatrix;

    previousVoteTallies = servers[which].previousVoteTallies;
    currentPseudonyms = servers[which].currentPseudonyms;
    currentUserEncryptedTallies = servers[which].currentUserEncryptedTallies;
    voteMatrix = servers[which].voteMatrix;

    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_new_user_update(proofOfCorrectAddition, previousVoteTallies, 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<TwistBipoint>& newVotes,
    const Twistpoint& shortTermPublicKey)
{
    return receive_vote(pi, newVotes, shortTermPublicKey, 0);
}

bool PrsonaServerEntity::receive_vote(
    const std::vector<Proof>& pi,
    const std::vector<TwistBipoint>& newVotes,
    const Twistpoint& 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
{
    transmit_new_user_data(pi, newUser, 0);
}

void PrsonaServerEntity::transmit_new_user_data(
    const std::vector<Proof>& pi,
    PrsonaClient& newUser,
    size_t which) const
{
    Twistpoint shortTermPublicKey = newUser.get_short_term_public_key();

    Proof baseProof;
    std::vector<Proof> serverEncryptedScoreProof;
    CurveBipoint serverEncryptedScore = get_current_server_encrypted_tally(baseProof, shortTermPublicKey, which);
    serverEncryptedScoreProof.push_back(baseProof);
    get_other_server_tally_commitments(serverEncryptedScoreProof, shortTermPublicKey, which);

    baseProof.clear();
    std::vector<Proof> userEncryptedScoreProof;
    EGCiphertext userEncryptedScore = get_current_user_encrypted_tally(baseProof, shortTermPublicKey, which);
    userEncryptedScoreProof.push_back(baseProof);
    get_other_user_tally_commitments(userEncryptedScoreProof, shortTermPublicKey, which);

    baseProof.clear();
    std::vector<Proof> voteMatrixProof;
    std::vector<std::vector<TwistBipoint>> voteMatrix = get_all_current_votes(baseProof, which);
    voteMatrixProof.push_back(baseProof);
    get_other_vote_matrix_commitments(voteMatrixProof, which);

    baseProof.clear();
    std::vector<Proof> pseudonymsProof;
    std::vector<Twistpoint> currentPseudonyms = get_current_pseudonyms(baseProof, which);
    pseudonymsProof.push_back(baseProof);
    get_other_pseudonyms_commitments(pseudonymsProof, which);

    newUser.receive_new_user_data(pi, serverEncryptedScoreProof, serverEncryptedScore, userEncryptedScoreProof, userEncryptedScore, voteMatrixProof, voteMatrix, pseudonymsProof, currentPseudonyms);
}

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;
    Twistpoint 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);

    Proof baseProof;
    std::vector<Proof> fullProof;
    Twistpoint shortTermPublicKey = currUser.get_short_term_public_key();
    EGCiphertext encryptedScore = get_current_user_encrypted_tally(baseProof, shortTermPublicKey, which);
    fullProof.push_back(baseProof);
    get_other_user_tally_commitments(fullProof, shortTermPublicKey, which);

    currUser.receive_vote_tally(fullProof, encryptedScore);
}

/*
 * EPOCH
 */

void PrsonaServerEntity::epoch()
{
    epoch(0);
}

// Do the epoch process
void PrsonaServerEntity::epoch(
    size_t which)
{
    Twistpoint nextGenerator = PrsonaServer::EL_GAMAL_GENERATOR;
    
    std::vector<std::vector<std::vector<Proof>>> pi;
    std::vector<std::vector<std::vector<Twistpoint>>> permutationCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymSeedCommits;
    std::vector<std::vector<std::vector<CurveBipoint>>> serverTallyCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> partwayVoteMatrixCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> finalVoteMatrixCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallyMaskCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallyMessageCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallySeedCommits;

    std::vector<std::vector<Proof>> generator_proof_holder(1);
    pi.push_back(generator_proof_holder);

    // go from A_0 to A_0.5
    for (size_t i = 0; i < servers[which].get_num_servers(); i++)
    {   
        size_t realI = (which + i) % servers[which].get_num_servers();

        servers[realI].build_up_midway_pseudonyms(pi, permutationCommits, freshPseudonymCommits, freshPseudonymSeedCommits, serverTallyCommits, partwayVoteMatrixCommits, finalVoteMatrixCommits, nextGenerator);

        for (size_t j = 1; j < servers[which].get_num_servers(); j++)
        {
            std::vector<std::vector<Twistpoint>> currUserTallyMaskCommits;
            std::vector<std::vector<Twistpoint>> currUserTallyMessageCommits;
            std::vector<std::vector<Twistpoint>> currUserTallySeedCommits;

            size_t realJ = (realI + j) % servers[which].get_num_servers();

            servers[realJ].accept_epoch_updates(pi[i + 1], permutationCommits[i], freshPseudonymCommits[i], freshPseudonymSeedCommits[i], serverTallyCommits[i], partwayVoteMatrixCommits[i], finalVoteMatrixCommits[i], currUserTallyMaskCommits, currUserTallyMessageCommits, currUserTallySeedCommits, nextGenerator, false);
        }
    }

    std::vector<Proof> generator_proof = pi[0][0];
    pi.clear();
    permutationCommits.clear();
    freshPseudonymCommits.clear();
    freshPseudonymSeedCommits.clear();
    serverTallyCommits.clear();
    partwayVoteMatrixCommits.clear();
    finalVoteMatrixCommits.clear();

    /* 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. */
    std::vector<EGCiphertext> currentUserEncryptedTallies;
    std::vector<CurveBipoint> currentServerEncryptedTallies;
    tally_scores(nextGenerator, currentUserEncryptedTallies, currentServerEncryptedTallies, which);
    distribute_tallied_scores(currentUserEncryptedTallies, currentServerEncryptedTallies);
    
    // go from A_0.5 to A_1
    for (size_t i = 0; i < servers[which].get_num_servers(); i++)
    {
        size_t realI = (which + i) % servers[which].get_num_servers();
        
        servers[realI].break_down_midway_pseudonyms(generator_proof, pi, permutationCommits, freshPseudonymCommits, freshPseudonymSeedCommits, serverTallyCommits, partwayVoteMatrixCommits, finalVoteMatrixCommits, userTallyMaskCommits, userTallyMessageCommits, userTallySeedCommits, nextGenerator);
        
        for (size_t j = 1; j < servers[which].get_num_servers(); j++)
        {
            size_t realJ = (realI + j) % servers[which].get_num_servers();

            servers[realJ].accept_epoch_updates(pi[i], permutationCommits[i], freshPseudonymCommits[i], freshPseudonymSeedCommits[i], serverTallyCommits[i], partwayVoteMatrixCommits[i], finalVoteMatrixCommits[i], userTallyMaskCommits[i], userTallyMessageCommits[i], userTallySeedCommits[i], nextGenerator, true);
        }
    }
}

void PrsonaServerEntity::print_scores() const
{
    BGN bgnSystem = servers[0].bgnSystem;
    std::vector<CurveBipoint> 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<TwistBipoint>> 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;
    }
}

void PrsonaServerEntity::print_current_votes_by(
    const Twistpoint& index) const
{
    BGN bgnSystem = servers[0].bgnSystem;
    size_t realIndex = servers[0].binary_search(index);
    std::vector<TwistBipoint> scores = servers[0].voteMatrix[realIndex];

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

/*********************
 * 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. */
void PrsonaServerEntity::tally_scores(
    const Twistpoint& nextGenerator,
    std::vector<EGCiphertext>& userTallyScores,
    std::vector<CurveBipoint>& serverTallyScores,
    size_t which)
{
    std::vector<EGCiphertext> retval;
    std::vector<Scalar> decryptedTalliedScores = servers[which].tally_scores();
    mpz_class maxScorePossibleThisRound = servers[which].get_max_possible_score().toInt() * PrsonaBase::get_max_allowed_vote();
    mpz_class topOfScoreRange = decryptedTalliedScores.size() * PrsonaBase::get_max_allowed_vote();

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

        EGCiphertext currCiphertext;
        userTallyScores.push_back(currCiphertext);
        CurveBipoint currServerScore;
        serverTallyScores.push_back(currServerScore);
        Scalar currMask;
        currMask.set_random();

        // Give the server the new weights,
        // to get passed around to the other servers
        servers[which].bgnSystem.encrypt(serverTallyScores[i], decryptedTalliedScores[i]);

        userTallyScores[i].mask = servers[which].currentPseudonyms[i] * currMask;
        userTallyScores[i].encryptedMessage = (nextGenerator * currMask) + (servers[which].get_blinding_generator() * decryptedTalliedScores[i]);
    }
}

void PrsonaServerEntity::distribute_tallied_scores(
    const std::vector<EGCiphertext>& userScores,
    const std::vector<CurveBipoint>& serverScores)
{
    for (size_t i = 0; i < servers[0].get_num_servers(); i++)
        servers[i].receive_tallied_scores(userScores, serverScores);
}

/*
 * BINARY SEARCH
 */

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