#include <iostream>

#include "client.hpp"
#include "serverEntity.hpp"

extern const curvepoint_fp_t bn_curvegen;
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 PrsonaClient::EL_GAMAL_GENERATOR = Curvepoint();
Curvepoint PrsonaClient::EL_GAMAL_BLIND_GENERATOR = Curvepoint();
bool PrsonaClient::SERVER_IS_MALICIOUS = false;
bool PrsonaClient::CLIENT_IS_MALICIOUS = false;

// Quick and dirty function to calculate ceil(log base 2) with mpz_class
mpz_class log2(mpz_class x)
{
    mpz_class retval = 0;
    while (x > 0)
    {
        retval++;
        x = x >> 1;
    }

    return retval;
}

mpz_class bit(mpz_class x)
{
    return x > 0 ? 1 : 0;
}

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

/*
 * CONSTRUCTORS
 */

PrsonaClient::PrsonaClient(
    const BGNPublicKey& serverPublicKey,
    const PrsonaServerEntity* servers)
    : serverPublicKey(serverPublicKey),
        servers(servers),
        max_checked(0)
{
    longTermPrivateKey.set_random();
    inversePrivateKey = longTermPrivateKey.curveInverse();

    decryption_memoizer[EL_GAMAL_BLIND_GENERATOR * max_checked] = max_checked;
}

/*
 * SETUP FUNCTIONS
 */

// Must be called once before any usage of this class
void PrsonaClient::init(const Curvepoint& elGamalBlindGenerator)
{
    EL_GAMAL_GENERATOR = Curvepoint(bn_curvegen);
    EL_GAMAL_BLIND_GENERATOR = elGamalBlindGenerator;
}

void PrsonaClient::set_server_malicious()
{
    SERVER_IS_MALICIOUS = true;
}

void PrsonaClient::set_client_malicious()
{
    CLIENT_IS_MALICIOUS = true;
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

Curvepoint 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<CurveBipoint> PrsonaClient::make_votes(
    Proof& pi,
    const std::vector<CurveBipoint>& currentEncryptedVotes,
    const std::vector<Scalar>& votes,
    const std::vector<bool>& replaces)
{
    std::vector<CurveBipoint> retval;

    if (!verify_valid_votes_proof(pi, currentEncryptedVotes))
    {
        std::cerr << "Could not verify proof of valid votes." << std::endl;
        return retval;
    }

    for (size_t i = 0; i < votes.size(); i++)
    {
        CurveBipoint currScore;
        if (replaces[i])
            serverPublicKey.encrypt(currScore, votes[i]);
        else
            currScore = serverPublicKey.rerandomize(currentEncryptedVotes[i]);

        retval.push_back(currScore);
    }

    pi = generate_vote_proof(retval, votes);
    return retval;
}

// Get a new fresh generator (happens at initialization and during each epoch)
void PrsonaClient::receive_fresh_generator(const Curvepoint& freshGenerator)
{
    currentFreshGenerator = freshGenerator;
}

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

    currentEncryptedScore = score;
    decrypt_score(score);
}

/*
 * REPUTATION PROOFS
 */

// A pretty straightforward range proof (generation)
std::vector<Proof> PrsonaClient::generate_reputation_proof(
    const Scalar& threshold) const
{
    std::vector<Proof> retval;

    // Don't even try if the user asks to make an illegitimate proof
    if (threshold > currentScore)
        return retval;

    // Base case
    if (!CLIENT_IS_MALICIOUS)
    {
        Proof currProof;
        currProof.basic = "PROOF";

        retval.push_back(currProof);
        return retval;
    }

    // We really have two consecutive proofs in a junction.
    // The first is to prove that we are the stpk we claim we are
    retval.push_back(generate_ownership_proof());

    // The value we're actually using in our proof
    mpz_class proofVal = currentScore.curveSub(threshold).toInt();
    // Top of the range in our proof determined by what scores are even possible
    mpz_class proofBits =
        log2(
            servers->get_num_clients() * MAX_ALLOWED_VOTE -
            threshold.toInt());
    
    // Don't risk a situation that would divulge our private key
    if (proofBits <= 1)
        proofBits = 2;

    // This seems weird, but remember our base is A_t^r, not g^t
    std::vector<Scalar> masksPerBit;
    masksPerBit.push_back(inversePrivateKey);
    for (size_t i = 1; i < proofBits; i++)
    {
        Scalar currMask;
        currMask.set_random();

        masksPerBit.push_back(currMask);
        masksPerBit[0] =
            masksPerBit[0].curveSub(currMask.curveMult(Scalar(1 << i)));
    }

    // Taken from Fig. 1 in https://eprint.iacr.org/2014/764.pdf
    for (size_t i = 0; i < proofBits; i++)
    {
        Proof currProof;
        Curvepoint g, h, c, c_a, c_b;
        g = currentEncryptedScore.mask;
        h = EL_GAMAL_BLIND_GENERATOR;
    
        mpz_class currBit = bit(proofVal & (1 << i));
        Scalar a, s, t, m, r;
        a.set_random();
        s.set_random();
        t.set_random();
        m = Scalar(currBit);
        r = masksPerBit[i];
        
        c = g * r + h * m;
        currProof.partialUniversals.push_back(c);

        c_a = g * s + h * a;

        Scalar am = a.curveMult(m);
        c_b = g * t + h * am;

        std::stringstream oracleInput;
        oracleInput << g << h << c << c_a << c_b;

        Scalar x = oracle(oracleInput.str());
        currProof.challengeParts.push_back(x);

        Scalar f, z_a, z_b;
        Scalar mx = m.curveMult(x);
        f = mx.curveAdd(a);

        Scalar rx = r.curveMult(x);
        z_a = rx.curveAdd(s);

        Scalar x_f = x.curveSub(f);
        Scalar r_x_f = r.curveMult(x_f);
        z_b = r_x_f.curveAdd(t);

        currProof.responseParts.push_back(f);
        currProof.responseParts.push_back(z_a);
        currProof.responseParts.push_back(z_b);

        retval.push_back(currProof);
    }

    return retval;
}

// A pretty straightforward range proof (verification)
bool PrsonaClient::verify_reputation_proof(
    const std::vector<Proof>& pi,
    const Curvepoint& shortTermPublicKey,
    const Scalar& threshold) const
{
    // Reject outright if there's no proof to check
    if (pi.empty())
    {
        std::cerr << "Proof was empty, aborting." << std::endl;
        return false;
    }

    // Base case
    if (!CLIENT_IS_MALICIOUS)
        return pi[0].basic == "PROOF";

    // User should be able to prove they are who they say they are
    if (!verify_ownership_proof(pi[0], shortTermPublicKey))
    {
        std::cerr << "Schnorr proof failed, aborting." << std::endl;
        return false;
    }

    // Get the encrypted score in question from the servers
    Proof serverProof;
    EGCiphertext encryptedScore =
        servers->get_current_tally(serverProof, shortTermPublicKey);

    // Rough for the prover but if the server messes up,
    // no way to prove the thing anyways
    if (!verify_valid_tally_proof(serverProof, encryptedScore))
    {
        std::cerr << "Server error prevented proof from working, aborting." << std::endl;
        return false;
    }

    // X is the thing we're going to be checking in on throughout
    // to try to get our score commitment back in the end.
    Curvepoint X;
    for (size_t i = 1; i < pi.size(); i++)
    {
        Curvepoint c, g, h;
        c = pi[i].partialUniversals[0];
        g = encryptedScore.mask;
        h = EL_GAMAL_BLIND_GENERATOR;

        X = X + c * Scalar(1 << (i - 1));

        Scalar x, f, z_a, z_b;
        x = pi[i].challengeParts[0];
        f = pi[i].responseParts[0];
        z_a = pi[i].responseParts[1];
        z_b = pi[i].responseParts[2];

        // Taken from Fig. 1 in https://eprint.iacr.org/2014/764.pdf
        Curvepoint c_a, c_b;
        c_a = g * z_a + h * f - c * x;
        Scalar x_f = x.curveSub(f);
        c_b = g * z_b - c * x_f;

        std::stringstream oracleInput;
        oracleInput << g << h << c << c_a << c_b;

        if (oracle(oracleInput.str()) != pi[i].challengeParts[0])
        {
            std::cerr << "0 or 1 proof failed at index " << i << " of " << pi.size() - 1 << ", aborting." << std::endl;
            return false;
        }
    }

    Scalar negThreshold;
    negThreshold = Scalar(0).curveSub(threshold);

    Curvepoint scoreCommitment =
        encryptedScore.encryptedMessage +
        EL_GAMAL_BLIND_GENERATOR * negThreshold;
    
    return X == scoreCommitment;
}

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

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

/*
 * SCORE DECRYPTION
 */

// Basic memoized score decryption
void PrsonaClient::decrypt_score(const EGCiphertext& score)
{
    Curvepoint 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())
    {
        currentScore = lookup->second;
        return;
    }

    // If not, iterate until we find it (adding everything to the memoization)
    max_checked++;
    Curvepoint decryptionCandidate = EL_GAMAL_BLIND_GENERATOR * max_checked;
    while (decryptionCandidate != hashedDecrypted)
    {
        decryption_memoizer[decryptionCandidate] = max_checked;

        decryptionCandidate = decryptionCandidate + EL_GAMAL_BLIND_GENERATOR;
        max_checked++;
    }
    decryption_memoizer[decryptionCandidate] = max_checked;

    // Set the value we found
    currentScore = max_checked;
}

/*
 * OWNERSHIP PROOFS
 */

// Very basic Schnorr proof (generation)
Proof PrsonaClient::generate_ownership_proof() const
{
    Proof retval;
    if (!CLIENT_IS_MALICIOUS)
    {
        retval.basic = "PROOF";
        return retval;
    }

    std::stringstream oracleInput;
    
    Scalar r;
    r.set_random();

    Curvepoint shortTermPublicKey = currentFreshGenerator * longTermPrivateKey;
    Curvepoint u = currentFreshGenerator * r;
    oracleInput << currentFreshGenerator << shortTermPublicKey << u;

    Scalar c = oracle(oracleInput.str());
    Scalar z = r.curveAdd(c.curveMult(longTermPrivateKey));

    retval.basic = "PROOF";
    retval.challengeParts.push_back(c);
    retval.responseParts.push_back(z);

    return retval;
}

// Very basic Schnorr proof (verification)
bool PrsonaClient::verify_ownership_proof(
    const Proof& pi, const Curvepoint& shortTermPublicKey) const
{
    if (!CLIENT_IS_MALICIOUS)
        return pi.basic == "PROOF";

    Scalar c = pi.challengeParts[0];
    Scalar z = pi.responseParts[0];

    Curvepoint u = currentFreshGenerator * z - shortTermPublicKey * c;

    std::stringstream oracleInput;
    oracleInput << currentFreshGenerator << shortTermPublicKey << u;
    
    return c == oracle(oracleInput.str());
}

/*
 * PROOF VERIFICATION
 */

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

    return pi.basic == "PROOF";
}

bool PrsonaClient::verify_generator_proof(
    const Proof& pi, const Curvepoint& generator) const
{
    if (!SERVER_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

bool PrsonaClient::verify_default_tally_proof(
    const Proof& pi, const EGCiphertext& score) const
{
    if (!SERVER_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

bool PrsonaClient::verify_valid_tally_proof(
    const Proof& pi, const EGCiphertext& score) const
{
    if (!SERVER_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

bool PrsonaClient::verify_default_votes_proof(
    const Proof& pi, const std::vector<CurveBipoint>& votes) const
{
    if (!SERVER_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

bool PrsonaClient::verify_valid_votes_proof(
    const Proof& pi, const std::vector<CurveBipoint>& votes) const
{
    if (!SERVER_IS_MALICIOUS)
        return pi.basic == "PROOF";

    return pi.basic == "PROOF";
}

/*
 * PROOF GENERATION
 */

Proof PrsonaClient::generate_vote_proof(
    const std::vector<CurveBipoint>& encryptedVotes,
    const std::vector<Scalar>& vote) const
{
    Proof retval;

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

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