#include <iostream>

#include "server.hpp"

extern const curvepoint_fp_t bn_curvegen;
const Curvepoint PrsonaServer::elGamalGenerator(bn_curvegen);
extern const scalar_t bn_n;
const Scalar PrsonaServer::scalarN(bn_n);
const Scalar PrsonaServer::defaultTally(0);
const Scalar PrsonaServer::defaultVote(0);

bool PrsonaServer::malicious_server = false;
bool PrsonaServer::malicious_client = false;

PrsonaServer::PrsonaServer()
{
    Scalar lambda;
    lambda.set_random();

    elGamalBlindGenerator = elGamalGenerator * lambda;

    currentSeed.set_random();
}

PrsonaServer::PrsonaServer(const BGN& other_bgn, const Curvepoint& other_blind_gen)
: bgn_system(other_bgn), elGamalBlindGenerator(other_blind_gen)
{
    currentSeed.set_random();
}

void PrsonaServer::set_malicious_server()
{
    malicious_server = true;
}

void PrsonaServer::set_malicious_client()
{
    malicious_client = true;
}

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

Curvepoint PrsonaServer::get_blinding_generator() const
{
    return elGamalBlindGenerator;
}

Curvepoint PrsonaServer::add_seed_to_generator(Proof& pi, const Curvepoint& currGenerator) const
{
    pi = generate_valid_fresh_generator_proof(pi);

    return currGenerator * currentSeed;
}

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

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

    TwistBipoint newTalliedVote;
    bgn_system.encrypt(newTalliedVote, Scalar(defaultTally));
    previousVoteTally.push_back(newTalliedVote);

    CurveBipoint encryptedDefaultVote;
    bgn_system.encrypt(encryptedDefaultVote, Scalar(defaultVote));
    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);
    }

    encryptedDefaultVote = bgn_system.rerandomize(encryptedDefaultVote);
    newRow.push_back(encryptedDefaultVote);
    voteMatrix.push_back(newRow);

    order_data();

    proofOfValidAddition = generate_proof_of_added_user(shortTermPublicKey);
}

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

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

std::vector<Scalar> PrsonaServer::tally_scores(Proof& pi)
{
    std::vector<Quadripoint> BGNEncryptedTallies;
    std::vector<Scalar> decryptedTallies;
    for (size_t i = 0; i < voteMatrix.size(); i++)
    {
        std::vector<Quadripoint> weightedVotes;

        for (size_t j = 0; j < previousVoteTally.size(); j++)
        {
            Quadripoint curr = bgn_system.homomorphic_multiplication(voteMatrix[j][i], previousVoteTally[j]);
            weightedVotes.push_back(curr);
        }

        Quadripoint currEncryptedTally = weightedVotes[0];
        for (size_t j = 1; j < weightedVotes.size(); j++)
            currEncryptedTally = bgn_system.homomorphic_addition(currEncryptedTally, weightedVotes[j]);

        currEncryptedTally = bgn_system.rerandomize(currEncryptedTally);
        decryptedTallies.push_back(bgn_system.decrypt(currEncryptedTally));
    }

    pi = generate_proof_of_correct_tally(BGNEncryptedTallies, decryptedTallies);

    return decryptedTallies;
}

void PrsonaServer::export_updates(std::vector<TwistBipoint>& otherPreviousVoteTally,
    std::vector<Curvepoint>& otherCurrentPseudonyms, std::vector<std::vector<CurveBipoint>>& otherVoteMatrix) const
{
    otherPreviousVoteTally = previousVoteTally;
    otherCurrentPseudonyms = currentPseudonyms;
    otherVoteMatrix = voteMatrix;
}

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

    previousVoteTally = otherPreviousVoteTally;
    currentPseudonyms = otherCurrentPseudonyms;
    voteMatrix = otherVoteMatrix;
}

void PrsonaServer::epoch_part_one(Proof& pi, Curvepoint& nextGenerator, std::vector<Scalar>& decryptedTallies)
{
    nextSeed.set_random();
    nextGenerator = nextGenerator * nextSeed;

    for (size_t i = 0; i < currentPseudonyms.size(); i++)
        currentPseudonyms[i] = currentPseudonyms[i] * nextSeed;
        
    std::vector<size_t> shuffleOrder = order_data();
    rerandomize_vote_matrix();

    std::vector<Scalar> reorderedTallies;
    for (size_t i = 0; i < shuffleOrder.size(); i++)
        reorderedTallies.push_back(decryptedTallies[shuffleOrder[i]]);
    decryptedTallies = reorderedTallies;

    pi = generate_proof_of_shuffle(shuffleOrder);
}

void PrsonaServer::epoch_part_two(Proof& pi, std::vector<EGCiphertext>& encryptedTallies)
{
    // TOFIX
    Scalar inverseSeed = scalarN - currentSeed;

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

    std::vector<size_t> shuffleOrder = order_data();
    rerandomize_vote_matrix();

    std::vector<EGCiphertext> reorderedTallies;
    for (size_t i = 0; i < shuffleOrder.size(); i++)
        reorderedTallies.push_back(encryptedTallies[shuffleOrder[i]]);
    encryptedTallies = reorderedTallies;

    pi = generate_proof_of_shuffle(shuffleOrder);

    currentSeed = nextSeed;
}

std::vector<size_t> PrsonaServer::order_data()
{
    std::vector<size_t> retval;
    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());

    std::vector<Curvepoint> newPseudonyms;
    std::vector<TwistBipoint> newVoteTallies;
    std::vector<std::vector<CurveBipoint>> newVoteMatrix;

    for (size_t i = 0; i < sortTracker.size(); i++)
    {
        newPseudonyms.push_back(sortTracker[i].pseudonym);
        newVoteTallies.push_back(previousVoteTally[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);
    }

    currentPseudonyms = newPseudonyms;
    previousVoteTally = newVoteTallies;
    voteMatrix = newVoteMatrix;

    return retval;
}

void PrsonaServer::rerandomize_vote_matrix()
{
    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]);
        }
    }
}

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

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

    return pi == "PROOF";
}

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

    return pi == "PROOF";
}

bool PrsonaServer::verify_update_proof(const Proof& pi, const std::vector<TwistBipoint>& otherPreviousVoteTally, const std::vector<Curvepoint>& otherCurrentPseudonyms, const std::vector<std::vector<CurveBipoint>>& otherVoteMatrix) const
{
    if (!malicious_server)
        return pi == "PROOF";

    return pi == "PROOF";
}

Proof PrsonaServer::generate_valid_fresh_generator_proof(const Proof& oldProof) const
{
    if (!malicious_server)
        return "PROOF";

    return "PROOF";
}

Proof PrsonaServer::generate_votes_valid_proof(const std::vector<CurveBipoint>& votes, const Curvepoint& voter) const
{
    if (!malicious_server)
        return "PROOF";

    return "PROOF";
}

Proof PrsonaServer::generate_proof_of_added_user(const Curvepoint& shortTermPublicKey) const
{
    if (!malicious_server)
        return "PROOF";

    return "PROOF";
}

Proof PrsonaServer::generate_score_proof(const EGCiphertext& score) const
{
    if (!malicious_server)
        return "PROOF";

    return "PROOF";
}

Proof PrsonaServer::generate_proof_of_correct_tally(const std::vector<Quadripoint>& BGNEncryptedTallies, const std::vector<Scalar>& decryptedTallies) const
{
    if (!malicious_server)
        return "PROOF";

    return "PROOF";
}

Proof PrsonaServer::generate_proof_of_shuffle(const std::vector<size_t>& shuffle_order) const
{
    if (!malicious_server)
        return "PROOF";

    return "PROOF";
}