#include <iostream>

#include "serverEntity.hpp"

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;

    PrsonaServer firstServer;
    servers.push_back(firstServer);

    const BGN& sharedBGN = firstServer.get_bgn_details();
    Curvepoint sharedEGBlindGenerator = firstServer.get_blinding_generator();

    for (size_t i = 1; i < numServers; i++)
    {
        PrsonaServer currServer(sharedBGN, sharedEGBlindGenerator);
        servers.push_back(currServer);
    }
}

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(Proof& pi) const
{
    Curvepoint retval = PrsonaServer::elGamalGenerator;
    for (size_t j = 0; j < servers.size(); j++)
        servers[j].add_seed_to_generator(pi, retval);

    return retval;
}

void PrsonaServerEntity::add_new_client(PrsonaClient& newUser)
{
    Proof proofOfValidGenerator, proofOfValidSTPK, proofOfCorrectAddition, proofOfDefaultTally, proofOfValidVotes;
    Curvepoint freshGenerator = get_fresh_generator(proofOfValidGenerator);

    newUser.receive_fresh_generator(proofOfValidGenerator, freshGenerator);
    Curvepoint shortTermPublicKey = newUser.get_short_term_public_key(proofOfValidSTPK);

    servers[0].add_new_client(proofOfValidSTPK, proofOfCorrectAddition, shortTermPublicKey);

    std::vector<TwistBipoint> previousVoteTally;
    std::vector<Curvepoint> currentPseudonyms;
    std::vector<std::vector<CurveBipoint>> voteMatrix;
    servers[0].export_updates(previousVoteTally, currentPseudonyms, voteMatrix);

    for (size_t j = 1; j < servers.size(); j++)
        servers[j].import_updates(proofOfCorrectAddition, previousVoteTally, currentPseudonyms, voteMatrix);

    EGCiphertext defaultTally = get_default_tally(proofOfDefaultTally, shortTermPublicKey);
    newUser.receive_vote_tally(proofOfDefaultTally, defaultTally, true);

    std::vector<CurveBipoint> encryptedDefaults = servers[0].get_current_votes_by(proofOfValidVotes, shortTermPublicKey);
    newUser.receive_encrypted_votes(proofOfValidVotes, encryptedDefaults, true);
}

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

void PrsonaServerEntity::epoch()
{
    Proof pi, proofOfCorrectTally;
    Curvepoint nextGenerator = PrsonaServer::elGamalGenerator;
    std::vector<Scalar> decryptedTalliedScores = tally_scores(proofOfCorrectTally);
    std::vector<TwistBipoint> previousVoteTally;
    std::vector<Curvepoint> currentPseudonyms;
    std::vector<std::vector<CurveBipoint>> voteMatrix;

    for (size_t i = 0; i < servers.size() - 1; i++)
    {
        servers[i].epoch_part_one(pi, nextGenerator, decryptedTalliedScores);

        servers[i].export_updates(previousVoteTally, currentPseudonyms, voteMatrix);
        servers[i + 1].import_updates(pi, previousVoteTally, currentPseudonyms, voteMatrix);
    }

    servers[servers.size() - 1].epoch_part_one(pi, nextGenerator, decryptedTalliedScores);

    servers[servers.size() - 1].export_updates(previousVoteTally, currentPseudonyms, voteMatrix);

    encryptedTallies.clear();
    for (size_t i = 0; i < decryptedTalliedScores.size(); i++)
    {
        EGCiphertext currCiphertext;
        encryptedTallies.push_back(currCiphertext);

        Scalar currMask;
        currMask.set_random();

        servers[0].bgn_system.encrypt(previousVoteTally[i], decryptedTalliedScores[i]);
        
        encryptedTallies[i].mask = currentPseudonyms[i] * currMask;
        encryptedTallies[i].encryptedMessage = (nextGenerator * currMask) + (get_blinding_generator() * decryptedTalliedScores[i]);
    }

    pi = generate_epoch_round_one_proof(pi, proofOfCorrectTally);
    servers[0].import_updates(pi, previousVoteTally, currentPseudonyms, voteMatrix);
    
    for (size_t i = 0; i < servers.size() - 1; i++)
    {
        servers[i].epoch_part_two(pi, encryptedTallies);

        servers[i].export_updates(previousVoteTally, currentPseudonyms, voteMatrix);
        servers[i + 1].import_updates(pi, previousVoteTally, currentPseudonyms, voteMatrix);
    }

    servers[servers.size() - 1].epoch_part_two(pi, encryptedTallies);

    servers[servers.size() - 1].export_updates(previousVoteTally, currentPseudonyms, voteMatrix);
    for (size_t i = 0; i < servers.size() - 1; i++)
        servers[i].import_updates(pi, previousVoteTally, currentPseudonyms, voteMatrix);

    for (size_t i = 0; i < encryptedTallies.size(); i++)
        tallyProofs[i] = generate_epoch_proof(pi, encryptedTallies[i]);
}

void PrsonaServerEntity::transmit_score(PrsonaClient& newUser) const
{
    Proof proofOfValidGenerator, proofOfValidSTPK, proofOfScore, proofOfCorrectVotes;
    Curvepoint freshGenerator = get_fresh_generator(proofOfValidGenerator);

    newUser.receive_fresh_generator(proofOfValidGenerator, freshGenerator);
    Curvepoint shortTermPublicKey = newUser.get_short_term_public_key(proofOfValidSTPK);

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

    std::vector<CurveBipoint> encryptedVotes = servers[0].get_current_votes_by(proofOfCorrectVotes, shortTermPublicKey);
    newUser.receive_encrypted_votes(proofOfCorrectVotes, encryptedVotes, false);
}

EGCiphertext PrsonaServerEntity::get_default_tally(Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    Proof unused;
    EGCiphertext retval;

    Scalar lambda;
    lambda.set_random();

    retval.mask = shortTermPublicKey * lambda;
    retval.encryptedMessage = get_fresh_generator(unused) * lambda + get_blinding_generator() * PrsonaServer::defaultTally;

    pi = generate_valid_default_tally_proof(retval, lambda);

    return retval;
}

Proof PrsonaServerEntity::generate_valid_default_tally_proof(const EGCiphertext& encryptedDefaultTally, const Scalar& lambda) const
{
    if (!PrsonaServer::malicious_server)
        return "PROOF";

    return "PROOF";
}

Proof PrsonaServerEntity::generate_epoch_round_one_proof(const Proof& pi1, const Proof& pi2) const
{
    if (!PrsonaServer::malicious_server)
        return "PROOF";

    return "PROOF";
}

Proof PrsonaServerEntity::generate_epoch_proof(const Proof& pi, const EGCiphertext& encryptedTally) const
{
    if (!PrsonaServer::malicious_server)
        return "PROOF";

    return "PROOF";
}

EGCiphertext PrsonaServerEntity::get_current_tally(Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    size_t requester = binary_search(shortTermPublicKey);

    pi = tallyProofs[requester];

    return encryptedTallies[requester];
}

std::vector<Scalar> PrsonaServerEntity::tally_scores(Proof& pi)
{
    return servers[0].tally_scores(pi);
}

size_t PrsonaServerEntity::binary_search(const Curvepoint& shortTermPublicKey) const
{
    return servers[0].binary_search(shortTermPublicKey);
}