#include <iostream>

#include "server.hpp"

extern const curvepoint_fp_t bn_curvegen;
extern const scalar_t bn_n;
const int MAX_ALLOWED_VOTE = 2;

/* These lines need to be here so these static variables are defined,
 * but in C++ putting code here doesn't actually execute
 * (or at least, with g++, whenever it would execute is not at a useful time)
 * so we have an init() function to actually put the correct values in them. */
Curvepoint PrsonaServer::EL_GAMAL_GENERATOR = Curvepoint();
Curvepoint PrsonaServer::EL_GAMAL_BLIND_GENERATOR = Curvepoint();
Scalar PrsonaServer::SCALAR_N = Scalar();
Scalar PrsonaServer::DEFAULT_TALLY = Scalar();
Scalar PrsonaServer::DEFAULT_VOTE = Scalar();
bool PrsonaServer::SERVER_IS_MALICIOUS = false;
bool PrsonaServer::CLIENT_IS_MALICIOUS = false;

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

/*
 * CONSTRUCTORS
 */

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

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

/*
 * SETUP FUNCTIONS
 */

// Must be called once before any usage of this class
Curvepoint PrsonaServer::init()
{
    Scalar lambda;
    lambda.set_random();

    EL_GAMAL_GENERATOR = Curvepoint(bn_curvegen);
    EL_GAMAL_BLIND_GENERATOR = EL_GAMAL_GENERATOR * lambda;
    SCALAR_N = Scalar(bn_n);
    DEFAULT_TALLY = Scalar(1);
    DEFAULT_VOTE = Scalar(1);

    return EL_GAMAL_BLIND_GENERATOR;
}

// Call this (once) if using malicious-security servers
void PrsonaServer::set_server_malicious()
{
    SERVER_IS_MALICIOUS = true;
}

// Call this (once) if using malicious-security clients
void PrsonaServer::set_client_malicious()
{
    CLIENT_IS_MALICIOUS = true;
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

Curvepoint PrsonaServer::get_blinding_generator()
{
    return EL_GAMAL_BLIND_GENERATOR;
}

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

/*
 * 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(
    const Curvepoint& currGenerator) const
{
    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(
    const Curvepoint& currGenerator) const
{
    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_votes_valid_proof(retval, shortTermPublicKey);
    return retval;
}

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

/*
 * 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(
    const Proof& proofOfValidKey,
    Proof& proofOfValidAddition,
    const Curvepoint& shortTermPublicKey)
{
    if (!verify_valid_key_proof(proofOfValidKey, 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
    TwistBipoint encryptedDefaultTally;
    bgn_system.encrypt(encryptedDefaultTally, DEFAULT_TALLY);
    previousVoteTallies.push_back(encryptedDefaultTally);

    Scalar mask;
    mask.set_random();
    EGCiphertext newUserEncryptedTally;
    newUserEncryptedTally.mask = shortTermPublicKey * mask;
    newUserEncryptedTally.encryptedMessage =
        currentFreshGenerator * mask +
        get_blinding_generator() * DEFAULT_TALLY;
    currentUserEncryptedTallies.push_back(newUserEncryptedTally);
    currentTallyProofs.push_back(
        generate_valid_default_tally_proof(newUserEncryptedTally, mask));

    // Users are defaulted to casting a neutral vote for others.
    CurveBipoint encryptedDefaultVote, encryptedSelfVote;
    bgn_system.encrypt(encryptedDefaultVote, DEFAULT_VOTE);
    bgn_system.encrypt(encryptedSelfVote, Scalar(MAX_ALLOWED_VOTE));
    std::vector<CurveBipoint> newRow;
    for (size_t i = 0; i < voteMatrix.size(); i++)
    {
        encryptedDefaultVote = bgn_system.rerandomize(encryptedDefaultVote);
        voteMatrix[i].push_back(encryptedDefaultVote);

        encryptedDefaultVote = bgn_system.rerandomize(encryptedDefaultVote);
        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.
    newRow.push_back(encryptedSelfVote);
    voteMatrix.push_back(newRow);

    order_data(proofOfValidAddition);

    proofOfValidAddition = generate_proof_of_added_user(shortTermPublicKey);
}

// Receive a new vote row from a user (identified by short term public key).
void PrsonaServer::receive_vote(
    const Proof& pi,
    const std::vector<CurveBipoint>& votes,
    const Curvepoint& shortTermPublicKey)
{
    if (!verify_vote_proof(pi, votes, shortTermPublicKey))
    {
        std::cerr << "Could not verify votes." << std::endl;
        return;
    }

    size_t voteSubmitter = binary_search(shortTermPublicKey);
    voteMatrix[voteSubmitter] = votes;
}

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

/*
 * CONSTRUCTOR HELPERS
 */

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

void PrsonaServer::initialize_fresh_generator(const Curvepoint& firstGenerator)
{
    currentFreshGenerator = firstGenerator;
}

/*
 * 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<Proof>& tallyProofs)
{
    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 =
                bgn_system.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 =
                bgn_system.homomorphic_addition_no_rerandomize(
                    currEncryptedTally, weightedVotes[j]);
        }

        // DECRYPT
        decryptedTallies.push_back(bgn_system.decrypt(currEncryptedTally));
        tallyProofs.push_back(
            generate_proof_of_correct_tally(
                currEncryptedTally, decryptedTallies[i]));
    }

    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 =
            bgn_system.homomorphic_addition_no_rerandomize(
                currEncryptedVal, previousVoteTallies[i]);
    }

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

    pi = generate_proof_of_correct_sum(currEncryptedVal, retval);
    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();
    currentTallyProofs.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<Proof>& otherCurrentTallyProofs,
    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;
    currentTallyProofs = otherCurrentTallyProofs;
    voteMatrix = otherVoteMatrix;
}

void PrsonaServer::export_updates(
    std::vector<TwistBipoint>& otherPreviousVoteTallies,
    std::vector<Curvepoint>& otherCurrentPseudonyms,
    std::vector<EGCiphertext>& otherCurrentUserEncryptedTallies,
    std::vector<Proof>& otherCurrentTallyProofs,
    std::vector<std::vector<CurveBipoint>>& otherVoteMatrix) const
{
    otherPreviousVoteTallies = previousVoteTallies;
    otherCurrentPseudonyms = currentPseudonyms;
    otherCurrentUserEncryptedTallies = currentUserEncryptedTallies;
    otherCurrentTallyProofs = currentTallyProofs;
    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] = bgn_system.rerandomize(voteMatrix[i][j]);

        bgn_system.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<Proof> newTallyProofs;
    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]);
        }
        if (!currentTallyProofs.empty())
        {
            newTallyProofs.push_back(
                currentTallyProofs[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;
    currentTallyProofs = newTallyProofs;
    voteMatrix = newVoteMatrix;

    pi = generate_proof_of_shuffle(retval);
    return retval;
}

/*
 * BINARY SEARCH
 */

/* Completely normal binary search
 * There might be a standard function for this in <algorithms>?
 * But it returns an iterator, not a size_t, so less useful. */
size_t PrsonaServer::binary_search(const Curvepoint& index) const
{
    size_t lo, hi;
    lo = 0;
    hi = currentPseudonyms.size() - 1;

    while (lo < hi)
    {
        size_t mid = (lo + hi) / 2;
        if (currentPseudonyms[mid] < index)
            lo = mid + 1;
        else if (index == currentPseudonyms[mid])
            return mid;
        else hi = mid - 1;
    }

    return lo;
}

/*
 * PROOF VERIFICATION
 */

bool PrsonaServer::verify_valid_key_proof(
    const Proof& pi,
    const Curvepoint& shortTermPublicKey) const
{
    if (!CLIENT_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

bool PrsonaServer::verify_vote_proof(
    const Proof& pi,
    const std::vector<CurveBipoint>& votes,
    const Curvepoint& shortTermPublicKey) const
{
    if (!CLIENT_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

bool PrsonaServer::verify_update_proof(
    const Proof& pi) const
{
    if (!SERVER_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

/*
 * PROOF GENERATION
 */

Proof PrsonaServer::generate_valid_default_tally_proof(
    const EGCiphertext& newUserEncryptedTally, const Scalar& mask) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}

Proof PrsonaServer::generate_valid_fresh_generator_proof(
    const Proof& oldProof) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}

Proof PrsonaServer::generate_votes_valid_proof(
    const std::vector<CurveBipoint>& votes, const Curvepoint& voter) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}

Proof PrsonaServer::generate_proof_of_added_user(
    const Curvepoint& shortTermPublicKey) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}

Proof PrsonaServer::generate_score_proof(const EGCiphertext& score) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}

Proof PrsonaServer::generate_proof_of_correct_tally(
    const Quadripoint& BGNEncryptedTally,
    const Scalar& decryptedTally) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}

Proof PrsonaServer::generate_proof_of_correct_sum(
    const TwistBipoint& BGNEncryptedSum, const Scalar& decryptedSum) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}

Proof PrsonaServer::generate_proof_of_shuffle(
    const std::vector<size_t>& shuffle_order) const
{
    Proof retval;

    if (!SERVER_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    retval.basic = "PROOF";
    return retval;
}