#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;
    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].currentUserEncryptedTallies = userScores;
        servers[i].previousVoteTallies = 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);
}