#include <iostream>

#include "server.hpp"

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

/*
 * CONSTRUCTORS
 */

// Used to generate the first server; instantiates BGN for the first time
PrsonaServer::PrsonaServer(size_t numServers)
: numServers(numServers)
{
    currentSeed.set_random();
}

// Used for all other servers, so they have the same BGN parameters
PrsonaServer::PrsonaServer(size_t numServers, const BGN& otherBgn)
: numServers(numServers), bgnSystem(otherBgn)
{
    currentSeed.set_random();
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

BGNPublicKey PrsonaServer::get_bgn_public_key() const
{
    return bgnSystem.get_public_key();
}

size_t PrsonaServer::get_num_clients() const
{
    return currentPseudonyms.size();
}

size_t PrsonaServer::get_num_servers() const
{
    return numServers;
}

/*
 * FRESH GENERATOR CALCULATION
 */

// To calculate the current epoch's generator, start from the base generator,
// then have every server call this function on it iteratively (in any order).
Curvepoint PrsonaServer::add_curr_seed_to_generator(
    std::vector<Proof>& pi,
    const Curvepoint& currGenerator) const
{
    pi.push_back(add_to_generator_proof(currGenerator, currentSeed));

    return currGenerator * currentSeed;
}

// To calculate the next epoch's generator, start from the base generator,
// then have every server call this function on it iteratively (in any order).
Curvepoint PrsonaServer::add_next_seed_to_generator(
    std::vector<Proof>& pi,
    const Curvepoint& currGenerator) const
{
    pi.push_back(add_to_generator_proof(currGenerator, nextSeed));

    return currGenerator * nextSeed;
}

/*
 * 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> PrsonaServer::get_current_votes_by(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    std::vector<CurveBipoint> retval;
    size_t voteSubmitter = binary_search(shortTermPublicKey);
    retval = voteMatrix[voteSubmitter];

    pi = generate_valid_vote_row_proof();
    return retval;
}

std::vector<std::vector<CurveBipoint>> PrsonaServer::get_all_current_votes(
    Proof& pi) const
{
    pi = generate_valid_vote_matrix_proof();
    return voteMatrix;
}

/* 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 PrsonaServer::get_current_user_encrypted_tally(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    EGCiphertext retval;
    size_t tallyOwner = binary_search(shortTermPublicKey);
    retval = currentUserEncryptedTallies[tallyOwner];
    
    pi = generate_valid_user_tally_proof();
    return retval;
}

TwistBipoint PrsonaServer::get_current_server_encrypted_tally(
    Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    TwistBipoint retval;
    size_t tallyOwner = binary_search(shortTermPublicKey);
    retval = previousVoteTallies[tallyOwner];
    
    pi = generate_valid_server_tally_proof();
    return retval;
}

std::vector<Curvepoint> PrsonaServer::get_current_pseudonyms(Proof& pi) const
{    
    pi = generate_valid_pseudonyms_proof();
    return currentPseudonyms;
}

/*
 * CLIENT INTERACTIONS
 */

/* Add a new client (who is identified only by their short term public key)
 * One server will do this, then ask all other servers to import their
 * (proven) exported data. */
void PrsonaServer::add_new_client(
    std::vector<Proof>& proofOfValidAddition,
    const Proof& proofOfValidKey,
    const Curvepoint& shortTermPublicKey)
{
    if (!verify_ownership_proof(
            proofOfValidKey, currentFreshGenerator, shortTermPublicKey))
    {
        std::cerr << "Could not verify proof of valid key." << std::endl;
        return;
    }

    currentPseudonyms.push_back(shortTermPublicKey);

    // The first epoch's score for a new user will be low,
    // but will typically converge on an average score quickly
    Scalar tallySeed;
    TwistBipoint encryptedDefaultTally =
        bgnSystem.get_public_key().twistEncrypt(tallySeed, DEFAULT_TALLY);
    previousVoteTallies.push_back(encryptedDefaultTally);

    Scalar seedForUserTally;
    seedForUserTally.set_random();
    EGCiphertext newUserEncryptedTally;
    newUserEncryptedTally.mask = shortTermPublicKey * seedForUserTally;
    newUserEncryptedTally.encryptedMessage =
        currentFreshGenerator * seedForUserTally +
        elGamalBlindGenerator * DEFAULT_TALLY;
    currentUserEncryptedTallies.push_back(newUserEncryptedTally);

    // Users are defaulted to casting a neutral vote for others.
    CurveBipoint encryptedDefaultVote, encryptedSelfVote;
    Scalar currDefaultSeed, currSelfSeed;
    encryptedDefaultVote = 
        bgnSystem.get_public_key().curveEncrypt(currDefaultSeed, DEFAULT_VOTE);
    encryptedSelfVote =
        bgnSystem.get_public_key().curveEncrypt(currSelfSeed, Scalar(MAX_ALLOWED_VOTE));

    std::vector<CurveBipoint> newRow;
    std::vector<Scalar> userVoteSeeds;
    std::vector<Scalar> otherVoteSeeds;
    for (size_t i = 0; i < voteMatrix.size(); i++)
    {
        Scalar addedSeed;

        encryptedDefaultVote = bgnSystem.get_public_key().rerandomize(addedSeed, encryptedDefaultVote);
        currDefaultSeed = currDefaultSeed.curveAdd(addedSeed);

        otherVoteSeeds.push_back(Scalar());
        otherVoteSeeds[i] = currDefaultSeed;
        
        voteMatrix[i].push_back(encryptedDefaultVote);

        encryptedDefaultVote = bgnSystem.get_public_key().rerandomize(addedSeed, encryptedDefaultVote);
        currDefaultSeed = currDefaultSeed.curveAdd(addedSeed);

        userVoteSeeds.push_back(Scalar());
        userVoteSeeds[i] = currDefaultSeed;

        newRow.push_back(encryptedDefaultVote);
    }
    
    // Because we are adding the new user to the end (and then sorting it),
    // this last element (bottom right corner) is always the self vote.
    userVoteSeeds.push_back(Scalar());
    userVoteSeeds[newRow.size()] = currSelfSeed;
    
    otherVoteSeeds.push_back(Scalar());
    otherVoteSeeds[newRow.size()] = currSelfSeed;

    newRow.push_back(encryptedSelfVote);
    voteMatrix.push_back(newRow);

    Proof unused;
    std::vector<size_t> sortOrder = order_data(unused);
    std::vector<Scalar> newUserVoteSeeds;
    std::vector<Scalar> newOtherVoteSeeds;
    for (size_t i = 0; i < sortOrder.size(); i++)
    {
        newUserVoteSeeds.push_back(userVoteSeeds[sortOrder[i]]);
        newOtherVoteSeeds.push_back(otherVoteSeeds[sortOrder[i]]);
    }

    proofOfValidAddition = generate_proof_of_added_user(
        tallySeed,
        seedForUserTally,
        newUserVoteSeeds,
        newOtherVoteSeeds);
}

// Receive a new vote row from a user (identified by short term public key).
bool PrsonaServer::receive_vote(
    const std::vector<Proof>& pi,
    const std::vector<CurveBipoint>& newVotes,
    const Curvepoint& shortTermPublicKey)
{
    size_t voteSubmitter = binary_search(shortTermPublicKey);
    std::vector<CurveBipoint> oldVotes = voteMatrix[voteSubmitter];

    if (!verify_vote_proof(pi, oldVotes, newVotes, shortTermPublicKey))
        return false;

    voteMatrix[voteSubmitter] = newVotes;
    return true;
}

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

/*
 * CONSTRUCTOR HELPERS
 */

const BGN& PrsonaServer::get_bgn_details() const
{
    return bgnSystem;
}

bool PrsonaServer::initialize_fresh_generator(
    const std::vector<Proof>& pi,
    const Curvepoint& firstGenerator)
{
    if (!verify_generator_proof(pi, firstGenerator, numServers))
    {
        std::cerr << "Could not verify generator proof, aborting." << std::endl;
        return false;
    }

    currentFreshGenerator = firstGenerator;
    return true;
}

// To calculate the blind generator for ElGamal, start from the base generator,
// then have every server call this function on it iteratively (in any order).
Curvepoint PrsonaServer::add_rand_seed_to_generator(
    std::vector<Proof>& pi,
    const Curvepoint& currGenerator) const
{
    Scalar lambda;
    lambda.set_random();

    pi.push_back(add_to_generator_proof(currGenerator, lambda));

    return currGenerator * lambda;
}

bool PrsonaServer::set_EG_blind_generator(
    const std::vector<Proof>& pi,
    const Curvepoint& currGenerator)
{
    return PrsonaBase::set_EG_blind_generator(pi, currGenerator, numServers);
}

/*
 * SCORE TALLYING
 */

/* Calculate scores homomorphically (as an individual server would normally)
 * and then decrypt them (which the servers would normally work together to do).
 *
 * Note that since these calculations are just for us, we don't need to do any
 * expensive rerandomizations of intermediate values. */
std::vector<Scalar> PrsonaServer::tally_scores()
{
    std::vector<Quadripoint> BGNEncryptedTallies;
    std::vector<Scalar> decryptedTallies;
    for (size_t i = 0; i < voteMatrix.size(); i++)
    {
        std::vector<Quadripoint> weightedVotes;

        // ZIP
        for (size_t j = 0; j < previousVoteTallies.size(); j++)
        {
            Quadripoint curr =
                bgnSystem.homomorphic_multiplication_no_rerandomize(
                    voteMatrix[j][i], previousVoteTallies[j]);

            weightedVotes.push_back(curr);
        }

        // FOLDL
        Quadripoint currEncryptedTally = weightedVotes[0];
        for (size_t j = 1; j < weightedVotes.size(); j++)
        {
            currEncryptedTally =
                bgnSystem.homomorphic_addition_no_rerandomize(
                    currEncryptedTally, weightedVotes[j]);
        }

        // DECRYPT
        decryptedTallies.push_back(bgnSystem.decrypt(currEncryptedTally));
    }

    return decryptedTallies;
}

/* Calculate what the maximum possible score this round was (that is,
 * given the current user weights, what was the highest possible score?).
 *
 * As with individual scores, this also does the decryption that servers
 * would ordinarily work together to form. */
Scalar PrsonaServer::get_max_possible_score(Proof& pi)
{
    // FOLDL
    TwistBipoint currEncryptedVal = previousVoteTallies[0];
    for (size_t i = 1; i < previousVoteTallies.size(); i++)
    {
        currEncryptedVal =
            bgnSystem.homomorphic_addition_no_rerandomize(
                currEncryptedVal, previousVoteTallies[i]);
    }

    // DECRYPT
    Scalar retval = bgnSystem.decrypt(currEncryptedVal);

    pi = generate_proof_of_correct_sum();
    return retval;
}

/*
 * EPOCH ROUNDS
 */

// The first round, going from A_0 to A_0.5
void PrsonaServer::build_up_midway_pseudonyms(
    Proof& pi, Curvepoint& nextGenerator)
{
    nextSeed.set_random();
    nextGenerator = nextGenerator * nextSeed;

    currentUserEncryptedTallies.clear();

    for (size_t i = 0; i < currentPseudonyms.size(); i++)
        currentPseudonyms[i] = currentPseudonyms[i] * nextSeed;

    rerandomize_data();
    order_data(pi);
}

// In between these rounds, scores are tallied, decrypted,
// and encrypted to fresh user pseudonyms (possible through weird math)

// The second round, going from A_0.5 to A_1
void PrsonaServer::break_down_midway_pseudonyms(
    Proof& pi, const Curvepoint& nextGenerator)
{
    Scalar inverseSeed = currentSeed.curveInverse();

    for (size_t i = 0; i < currentPseudonyms.size(); i++)
    {
        currentPseudonyms[i] = currentPseudonyms[i] * inverseSeed;
        currentUserEncryptedTallies[i].mask =
            currentUserEncryptedTallies[i].mask * inverseSeed;
    }

    currentSeed = nextSeed;
    currentFreshGenerator = nextGenerator;

    rerandomize_data();
    order_data(pi);
}

/*
 * DATA MAINTENANCE
 */

void PrsonaServer::import_updates(
    const Proof& pi,
    const std::vector<TwistBipoint>& otherPreviousVoteTallies,
    const std::vector<Curvepoint>& otherCurrentPseudonyms,
    const std::vector<EGCiphertext>& otherCurrentUserEncryptedTallies,
    const std::vector<std::vector<CurveBipoint>>& otherVoteMatrix)
{
    if (!verify_update_proof(pi))
    {
        std::cerr << "Could not verify valid update." << std::endl;
        return;
    }

    previousVoteTallies = otherPreviousVoteTallies;
    currentPseudonyms = otherCurrentPseudonyms;
    currentUserEncryptedTallies = otherCurrentUserEncryptedTallies;
    voteMatrix = otherVoteMatrix;
}

void PrsonaServer::export_updates(
    std::vector<TwistBipoint>& otherPreviousVoteTallies,
    std::vector<Curvepoint>& otherCurrentPseudonyms,
    std::vector<EGCiphertext>& otherCurrentUserEncryptedTallies,
    std::vector<std::vector<CurveBipoint>>& otherVoteMatrix) const
{
    otherPreviousVoteTallies = previousVoteTallies;
    otherCurrentPseudonyms = currentPseudonyms;
    otherCurrentUserEncryptedTallies = currentUserEncryptedTallies;
    otherVoteMatrix = voteMatrix;
}

/*
 * DATA SAFEKEEPING
 */

// Everything needs to be rerandomized during epoch rounds
// NOTE: this may need to add something about proofs later?
void PrsonaServer::rerandomize_data()
{
    for (size_t i = 0; i < voteMatrix.size(); i++)
    {
        for (size_t j = 0; j < voteMatrix[0].size(); j++)
            voteMatrix[i][j] = bgnSystem.rerandomize(voteMatrix[i][j]);

        bgnSystem.rerandomize(previousVoteTallies[i]);
        if (!currentUserEncryptedTallies.empty())
        {
            Scalar rerandomizer;
            rerandomizer.set_random();

            currentUserEncryptedTallies[i].mask =
                currentUserEncryptedTallies[i].mask +
                currentPseudonyms[i] * rerandomizer;
            currentUserEncryptedTallies[i].encryptedMessage =
                currentUserEncryptedTallies[i].encryptedMessage +
                currentFreshGenerator * rerandomizer;
        }
    }
}

/* This is what powers the "shuffle"; really, as pseudonyms get updated,
 * the pseudonyms are no longer in the order prescribed by operator<().
 * So, we put them (and everything else) back into that order,
 * effectively shuffling them (and making lookups easier later on). */
std::vector<size_t> PrsonaServer::order_data(Proof& pi)
{
    std::vector<size_t> retval;

    // SortingType's index member allows us to replicate the "sort" across
    std::vector<SortingType> sortTracker;
    for (size_t i = 0; i < currentPseudonyms.size(); i++)
    {
        SortingType curr;
        
        curr.pseudonym = currentPseudonyms[i];
        curr.index = i;

        sortTracker.push_back(curr);
    }
    std::sort(sortTracker.begin(), sortTracker.end());

    // Order all other data in the same way, for consistency
    std::vector<Curvepoint> newPseudonyms;
    std::vector<TwistBipoint> newVoteTallies;
    std::vector<EGCiphertext> newUserEncryptedTallies;
    std::vector<std::vector<CurveBipoint>> newVoteMatrix;

    for (size_t i = 0; i < sortTracker.size(); i++)
    {
        newPseudonyms.push_back(sortTracker[i].pseudonym);
        newVoteTallies.push_back(previousVoteTallies[sortTracker[i].index]);
        
        if (!currentUserEncryptedTallies.empty())
        {
            newUserEncryptedTallies.push_back(
                currentUserEncryptedTallies[sortTracker[i].index]);
        }

        std::vector<CurveBipoint> currNewRow;
        for (size_t j = 0; j < currentPseudonyms.size(); j++)
        {
            currNewRow.push_back(
                voteMatrix[sortTracker[i].index][sortTracker[j].index]);
        }
        newVoteMatrix.push_back(currNewRow);

        retval.push_back(sortTracker[i].index);
    }

    previousVoteTallies = newVoteTallies;
    currentPseudonyms = newPseudonyms;
    currentUserEncryptedTallies = newUserEncryptedTallies;
    voteMatrix = newVoteMatrix;

    pi = generate_proof_of_shuffle();
    return retval;
}

/*
 * BINARY SEARCH
 */

// Completely normal binary search
size_t PrsonaServer::binary_search(const Curvepoint& index) const
{
    return PrsonaBase::binary_search(currentPseudonyms, index);
}

/*
 * VALID VOTE PROOFS
 */

bool PrsonaServer::verify_vote_proof(
    const std::vector<Proof>& pi,
    const std::vector<CurveBipoint>& oldVotes,
    const std::vector<CurveBipoint>& newVotes,
    const Curvepoint& shortTermPublicKey) const
{
    const BGNPublicKey& pubKey = bgnSystem.get_public_key();
    return PrsonaBase::verify_vote_proof(
        pubKey.get_bipoint_curvegen(),
        pubKey.get_bipoint_curve_subgroup_gen(),
        pi,
        oldVotes,
        newVotes,
        currentFreshGenerator,
        shortTermPublicKey);
}