#include <iostream>
#include <mutex>
#include <thread>

#include "client.hpp"

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

/*
 * CONSTRUCTORS
 */

PrsonaClient::PrsonaClient(
    const std::vector<Proof>& generatorProof,
    const Twistpoint& elGamalBlindGenerator,
    const BGNPublicKey& serverPublicKey,
    size_t numServers)
: serverPublicKey(serverPublicKey),
        numServers(numServers),
        max_checked(0)
{
    set_EG_blind_generator(generatorProof, elGamalBlindGenerator, numServers);

    longTermPrivateKey.set_random();
    inversePrivateKey = longTermPrivateKey.curveMultInverse();

    decryption_memoizer[elGamalBlindGenerator * max_checked] = max_checked;
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

Twistpoint PrsonaClient::get_short_term_public_key() const
{
    return currentFreshGenerator * longTermPrivateKey;
}

Twistpoint PrsonaClient::get_short_term_public_key(
    Proof &pi) const
{
    pi = generate_ownership_proof();
    return currentFreshGenerator * longTermPrivateKey;
}

/*
 * SERVER INTERACTIONS
 */

/* Generate a new vote vector to give to the servers
 * @replaces controls which votes are actually being updated and which are not
 *
 * You may really want to make currentEncryptedVotes a member variable, 
 * but it doesn't behave correctly when adding new clients after this one. */
std::vector<TwistBipoint> PrsonaClient::make_votes(
    std::vector<Proof>& validVoteProof,
    const std::vector<Proof>& serverProof,
    const std::vector<TwistBipoint>& oldEncryptedVotes,
    const std::vector<Scalar>& votes,
    const std::vector<bool>& replaces) const
{
    if (!verify_valid_vote_row_proof(serverProof, oldEncryptedVotes))
    {
        std::cerr << "Could not verify proof of valid votes." << std::endl;
        return oldEncryptedVotes;
    }

    std::vector<Scalar> seeds(oldEncryptedVotes.size());
    std::vector<TwistBipoint> newEncryptedVotes(oldEncryptedVotes.size());

    for (size_t i = 0; i < votes.size(); i++)
    {
        if (replaces[i])
            newEncryptedVotes[i] = serverPublicKey.twistEncrypt(seeds[i], votes[i]);
        else
            newEncryptedVotes[i] = serverPublicKey.rerandomize(seeds[i], oldEncryptedVotes[i]);
    }

    validVoteProof = generate_vote_proof(replaces, oldEncryptedVotes, newEncryptedVotes, seeds, votes);
    return newEncryptedVotes;
}

// Get a new fresh generator (happens at initialization and during each epoch)
bool PrsonaClient::receive_fresh_generator(
    const std::vector<Proof>& pi,
    const Twistpoint& freshGenerator)
{
    if (!verify_generator_proof(pi, freshGenerator, numServers))
    {
        std::cerr << "Issue verifying fresh generator proof." << std::endl;
        return false;
    }

    currentFreshGenerator = freshGenerator;
    return true;
}

// Receive a new encrypted score from the servers (each epoch)
bool PrsonaClient::receive_vote_tally(
    const std::vector<Proof>& pi,
    const EGCiphertext& score)
{
    if (!verify_valid_user_tally_proof(pi, score))
    {
        std::cerr << "Could not verify proof of valid tally." << std::endl;
        return false;
    }

    currentEncryptedScore = score;
    currentScore = decrypt_score(score);
    return true;
}

bool PrsonaClient::receive_new_user_data(
    const std::vector<Proof>& mainProof,
    const std::vector<Proof>& serverEncryptedScoreProof,
    const CurveBipoint& serverEncryptedScore,
    const std::vector<Proof>& userEncryptedScoreProof,
    const EGCiphertext& userEncryptedScore,
    const std::vector<Proof>& voteMatrixProof,
    const std::vector<std::vector<TwistBipoint>>& encryptedVoteMatrix,
    const std::vector<Proof>& pseudonymsProof,
    const std::vector<Twistpoint>& currentPseudonyms)
{
    Twistpoint shortTermPublicKey = get_short_term_public_key();

    if (!verify_valid_server_tally_proof(serverEncryptedScoreProof, serverEncryptedScore))
    {
        std::cerr << "Could not verify preliminary proof of server encrypted tally." << std::endl;
        return false;
    }

    if (!verify_valid_user_tally_proof(userEncryptedScoreProof, userEncryptedScore))
    {
        std::cerr << "Could not verify preliminary proof of user encrypted tally." << std::endl;
        return false;
    }

    if (!verify_valid_vote_matrix_proof(voteMatrixProof, encryptedVoteMatrix))
    {
        std::cerr << "Could not verify preliminary proof of encrypted votes." << std::endl;
        return false;
    }

    if (!verify_valid_pseudonyms_proof(pseudonymsProof, currentPseudonyms))
    {
        std::cerr << "Could not verify preliminary proof of pseudonyms." << std::endl;
        return false;
    }

    size_t selfIndex = binary_search(currentPseudonyms, shortTermPublicKey);
    if (currentPseudonyms[selfIndex] != shortTermPublicKey)
    {
        std::cerr << "Was not added to list of pseudonyms." << std::endl;
        return false;
    }

    bool flag = verify_proof_of_added_user(mainProof, currentFreshGenerator, shortTermPublicKey, serverPublicKey.get_bipoint_twistgen(), serverPublicKey.get_bipoint_twist_subgroup_gen(), serverPublicKey.get_bipoint_curvegen(), serverPublicKey.get_bipoint_curve_subgroup_gen(), selfIndex, userEncryptedScore, serverEncryptedScore, encryptedVoteMatrix);

    if (!flag)
    {
        std::cerr << "There was an issue verifying the proof; this user was not properly added." << std::endl;
        return false;
    }

    currentEncryptedScore = userEncryptedScore;
    currentScore = decrypt_score(userEncryptedScore);

    return true;
}

/*
 * REPUTATION PROOFS
 */

// A pretty straightforward range proof (generation)
std::vector<Proof> PrsonaClient::generate_reputation_proof(
    const Scalar& threshold,
    size_t numClients) const
{
    Proof ownershipProof = generate_ownership_proof();
    
    return PrsonaBase::generate_reputation_proof(ownershipProof, currentEncryptedScore, currentScore, threshold, inversePrivateKey, numClients);
}

bool PrsonaClient::verify_reputation_proof(
    const std::vector<Proof>& pi,
    const Twistpoint& shortTermPublicKey,
    const Scalar& threshold,
    const std::vector<Proof>& encryptedScoreProof,
    const EGCiphertext& encryptedScore) const
{
    if (!verify_valid_user_tally_proof(encryptedScoreProof, encryptedScore))
    {
        std::cerr << "Error getting score from server, aborting." << std::endl;
        return false;
    }

    return PrsonaBase::verify_reputation_proof(pi, currentFreshGenerator, shortTermPublicKey, encryptedScore, threshold);
}

Scalar PrsonaClient::get_score() const
{
    return currentScore;
}

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

/*
 * SCORE DECRYPTION
 */

void alarm_bell(void *mutex, int *value)
{
    std::this_thread::sleep_for(std::chrono::seconds(1));

    std::mutex *realMtx = (std::mutex *) mutex;
    std::unique_lock<std::mutex> lck(*realMtx);
    *value = 1;
}

// Basic memoized score decryption
Scalar PrsonaClient::decrypt_score(
    const EGCiphertext& score)
{
    Twistpoint s, hashedDecrypted;

    // Remove the mask portion of the ciphertext
    s = score.mask * inversePrivateKey;
    hashedDecrypted = score.encryptedMessage - s;
    
    // Check if it's a value we've already seen
    auto lookup = decryption_memoizer.find(hashedDecrypted);
    if (lookup != decryption_memoizer.end())
        return lookup->second;

    int alarm = 0;
    bool flag = false;
    std::mutex alarmMtx;
    std::unique_lock<std::mutex> lck(alarmMtx, std::defer_lock);

    std::thread alarmThread(alarm_bell, &alarmMtx, &alarm);

    // If not, iterate until we find it (adding everything to the memoization)
    max_checked++;
    Twistpoint decryptionCandidate = elGamalBlindGenerator * max_checked;

    lck.lock();
    while (!alarm && decryptionCandidate != hashedDecrypted)
    {
        lck.unlock();

        decryption_memoizer[decryptionCandidate] = max_checked;

        decryptionCandidate = decryptionCandidate + elGamalBlindGenerator;
        max_checked++;

        lck.lock();
    }
    if (alarm)
        flag = true;
    lck.unlock();

    decryption_memoizer[decryptionCandidate] = max_checked;

    alarmThread.join();

    if (flag)
        return Scalar(0);

    // Return the value we found
    return max_checked;
}

/*
 * OWNERSHIP PROOFS
 */

// Prove ownership of the short term public key
Proof PrsonaClient::generate_ownership_proof() const
{
    Twistpoint shortTermPublicKey = currentFreshGenerator * longTermPrivateKey;

    return PrsonaBase::generate_ownership_proof(currentFreshGenerator, shortTermPublicKey, longTermPrivateKey);
}

/*
 * VALID VOTE PROOFS
 */

std::vector<Proof> PrsonaClient::generate_vote_proof(
    const std::vector<bool>& replaces,
    const std::vector<TwistBipoint>& oldEncryptedVotes,
    const std::vector<TwistBipoint>& newEncryptedVotes,
    const std::vector<Scalar>& seeds,
    const std::vector<Scalar>& votes) const
{
    Proof pi = generate_ownership_proof();

    return PrsonaBase::generate_vote_proof(pi, serverPublicKey.get_bipoint_twistgen(), serverPublicKey.get_bipoint_twist_subgroup_gen(), replaces, oldEncryptedVotes, newEncryptedVotes, seeds, votes);
}
