#include <mutex>
#include <condition_variable>
#include <chrono>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cstdlib>

#include "networkClient.hpp"

using namespace std;

enum EventType {
    CLIENT_MAKE_VOTE = 1,
    CLIENT_MAKE_REP_PROOF
};

struct synchronization_tool exitSync;

// Initialize the classes we use
void initialize_prsona_classes()
{
    Scalar::init();
    PrsonaBase::init();
    PrsonaBase::set_client_malicious();
}

BGNPublicKey get_bgn_public_key_from_file(
    struct synchronization_tool *sync,
    const char *filename)
{
    unique_lock<mutex> lck(sync->mtx);
    ifstream bgnFile(filename);

    BGNPublicKey publicKey;
    bgnFile >> publicKey;

    return publicKey;
}

vector<Proof> get_valid_addition_proof_from_file(
    struct synchronization_tool *sync,
    const char *filename)
{
    unique_lock<mutex> lck(sync->mtx);
    ifstream additionFile(filename);

    vector<Proof> retval;

    BinarySizeT sizeOfVector;
    additionFile >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        Proof currProof;
        additionFile >> currProof;
        retval.push_back(currProof);
    }

    return retval;
}

BGNPublicKey get_bgn_public_key(
    default_random_engine *randomGenerator,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts)
{
    struct synchronization_tool sync;
    char *filename = NULL;

    uniform_int_distribution<size_t> distribution(0, serverIPs.size() - 1);
    size_t whichServer = distribution(*randomGenerator);

    bool flag = false;
    while (!flag)
    {
        struct mg_connection *conn = mg_connect_websocket_client(
            serverIPs[whichServer].c_str(),
            serverPorts[whichServer],
            USE_SSL,
            NULL,
            0,
            PUBLIC_BGN_URI,
            "null",
            file_websocket_data_handler,
            file_websocket_close_handler,
            (void *) &sync);

        if (!conn)
        {
            cerr << "Couldn't obtain BGN details" << endl;
            continue;
        }

        unique_lock<mutex> lck(sync.mtx);
        filename = set_temp_filename(*randomGenerator, conn);
        sync.val = 0;

        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        while (!sync.val)
            sync.cv.wait(lck);

        mg_close_connection(conn);

        flag = true;
    }

    BGNPublicKey retval =
        get_bgn_public_key_from_file(&sync, filename);

    remove(filename);
    delete filename;

    return retval;
}

void verify_valid_addition(
    PrsonaClient *newUser,
    default_random_engine *rng,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const vector<Proof>& proofOfValidAddition,
    const Twistpoint& shortTermPublicKey)
{
    vector<Proof> serverEncryptedScoreProof;
    CurveBipoint serverEncryptedScore =
        get_server_committed_val<CurveBipoint>(
            rng,
            serverIPs,
            serverPorts,
            serverEncryptedScoreProof,
            shortTermPublicKey,
            SERVER_TALLY_URI,
            SERVER_TALLY_COMMIT_URI);

    vector<Proof> userEncryptedScoreProof;
    EGCiphertext userEncryptedScore =
        get_server_committed_val<EGCiphertext>(
            rng,
            serverIPs,
            serverPorts,
            userEncryptedScoreProof,
            shortTermPublicKey,
            USER_TALLY_URI,
            USER_TALLY_COMMIT_URI);

    vector<Proof> voteMatrixProof;
    vector<vector<TwistBipoint>> voteMatrix =
        get_server_committed_val<vector<vector<TwistBipoint>>>(
            rng,
            serverIPs,
            serverPorts,
            voteMatrixProof,
            shortTermPublicKey,
            VOTE_MATRIX_URI,
            VOTE_MATRIX_COMMIT_URI);

    vector<Proof> pseudonymsProof;
    vector<Twistpoint> currentPseudonyms =
        get_server_committed_val<vector<Twistpoint>>(
            rng,
            serverIPs,
            serverPorts,
            pseudonymsProof,
            shortTermPublicKey,
            PSEUDONYMS_URI,
            PSEUDONYMS_COMMIT_URI);

    newUser->receive_new_user_data(
        proofOfValidAddition,
        serverEncryptedScoreProof,
        serverEncryptedScore,
        userEncryptedScoreProof,
        userEncryptedScore,
        voteMatrixProof,
        voteMatrix,
        pseudonymsProof,
        currentPseudonyms);
}

void register_new_client(
    default_random_engine *generator,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    PrsonaClient *newUser,
    const Proof& proofOfValidSTPK,
    const Twistpoint& shortTermPublicKey)
{
    struct synchronization_tool sync;
    char *filename = NULL;
    stringstream buffer;
    string data;

    buffer << proofOfValidSTPK;
    buffer << shortTermPublicKey;

    data = buffer.str();

    uniform_int_distribution<size_t> distribution(0, serverIPs.size() - 1);
    size_t whichServer = distribution(*generator);

    bool flag = false;
    while (!flag)
    {
        struct mg_connection *conn = mg_connect_websocket_client(
            serverIPs[whichServer].c_str(),
            serverPorts[whichServer],
            USE_SSL,
            NULL,
            0,
            PUBLIC_BGN_URI,
            "null",
            file_websocket_data_handler,
            file_websocket_close_handler,
            (void *) &sync);

        if (!conn)
        {
            cerr << "Couldn't register new client" << endl;
            continue;
        }

        unique_lock<mutex> lck(sync.mtx);
        filename = set_temp_filename(*generator, conn);
        sync.val = 0;
        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_BINARY,
            data.c_str(),
            data.length());

        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        while (!sync.val)
            sync.cv.wait(lck);

        mg_close_connection(conn);

        flag = true;
    }

    vector<Proof> proofOfValidAddition = 
        get_valid_addition_proof_from_file(&sync, filename);

    remove(filename);
    delete filename;

    verify_valid_addition(newUser, generator, serverIPs, serverPorts, proofOfValidAddition, shortTermPublicKey);
}

PrsonaClient *create_client(
    default_random_engine *generator,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    size_t numServers)
{
    BGNPublicKey publicKey =
        get_bgn_public_key(
            generator,
            serverIPs,
            serverPorts);

    vector<Proof> generatorProof;
    Twistpoint blindGenerator =
        get_generator(
            generator,
            serverIPs,
            serverPorts,
            generatorProof,
            false);

    PrsonaClient *retval =
        new PrsonaClient(generatorProof, blindGenerator, publicKey, numServers);

    Proof proofOfValidSTPK;
    Twistpoint shortTermPublicKey = retval->get_short_term_public_key(
                                        proofOfValidSTPK);

    register_new_client(
        generator,
        serverIPs,
        serverPorts,
        retval,
        proofOfValidSTPK,
        shortTermPublicKey);

    return retval;
}

string make_vote_string(
    const vector<Proof>& pi,
    const vector<TwistBipoint>& newVotes,
    const Twistpoint& shortTermPublicKey)
{
    stringstream buffer;
    BinarySizeT sizeOfVector;

    sizeOfVector.set(pi.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << pi[i];

    sizeOfVector.set(newVotes.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << newVotes[i];

    buffer << shortTermPublicKey;

    BinaryBool shouldDeal(true);
    buffer << shouldDeal;

    return buffer.str();
}

void send_item(
    const string& target, int targetPort, const string& data, const char* whichUri)
{
    struct synchronization_tool sync;
    bool flag = false;
    while (!flag)
    {
        struct mg_connection *conn = mg_connect_websocket_client(
            target.c_str(),
            targetPort,
            USE_SSL,
            NULL,
            0,
            whichUri,
            "null",
            synchro_websocket_data_handler,
            empty_websocket_close_handler,
            (void *) &sync);

        if (!conn)
        {
            cerr << "Couldn't obtain BGN details" << endl;
            continue;
        }

        unique_lock<mutex> lck(sync.mtx);
        sync.val = 0;

        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_BINARY,
            data.c_str(),
            data.length());

        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        while (!sync.val)
            sync.cv.wait(lck);

        mg_close_connection(conn);

        flag = true;
    }
}

void make_vote(
    default_random_engine& generator,
    const string& target,
    int targetPort,
    PrsonaClient* prsonaClient,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    size_t numClients)
{
    uniform_int_distribution<int> voteDistribution(
        0, PrsonaBase::get_max_allowed_vote());
    uniform_int_distribution<int> numVoteDistribution(
        0, numClients);
    size_t numVotes = numVoteDistribution(generator);

    // Make the correct number of new votes, but shuffle where they go
    vector<Scalar> votes;
    vector<bool> replaces;
    for (size_t j = 0; j < numClients; j++)
    {
        votes.push_back(Scalar(voteDistribution(generator)));
        replaces.push_back(j < numVotes);
    }
    shuffle(replaces.begin(), replaces.end(), generator);

    vector<Proof> generatorProof;
    Twistpoint freshGenerator =
        get_generator(
            &generator,
            serverIPs,
            serverPorts,
            generatorProof,
            true);

    prsonaClient->receive_fresh_generator(
        generatorProof, freshGenerator);
    
    Twistpoint shortTermPublicKey =
        prsonaClient->get_short_term_public_key();

    vector<Proof> fullProof;
    vector<TwistBipoint> encryptedVotes =
        get_server_committed_val<vector<TwistBipoint>>(
            &generator,
            serverIPs,
            serverPorts,
            fullProof,
            shortTermPublicKey,
            VOTES_BY_URI,
            VOTES_BY_COMMIT_URI);
    
    vector<Proof> voteProof;
    encryptedVotes =
        prsonaClient->make_votes(
            voteProof,
            fullProof,
            encryptedVotes,
            votes,
            replaces);

    string data =
        make_vote_string(voteProof, encryptedVotes, shortTermPublicKey);

    send_item(target, targetPort, data, GIVE_NEW_VOTE_URI);
}

string make_vote_string(
    const vector<Proof>& pi,
    const Twistpoint& shortTermPublicKey,
    const Scalar& threshold)
{
    stringstream buffer;
    BinarySizeT sizeOfVector;

    sizeOfVector.set(pi.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << pi[i];

    buffer << shortTermPublicKey;
    buffer << threshold;

    return buffer.str();
}

string make_rep_proof_string(
    const vector<Proof>& pi,
    const Twistpoint& shortTermPublicKey,
    const Scalar& threshold)
{
    stringstream buffer;
    BinarySizeT sizeOfVector;

    sizeOfVector.set(pi.size());
    
    buffer >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << pi[i];

    buffer << shortTermPublicKey;
    buffer << threshold;

    return buffer.str();
}

void make_reputation_proof(
    default_random_engine& generator,
    const string& target,
    int targetPort,
    PrsonaClient* prsonaClient,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    size_t numClients)
{
    vector<Proof> generatorProof;
    Twistpoint freshGenerator =
        get_generator(
            &generator,
            serverIPs,
            serverPorts,
            generatorProof,
            true);

    prsonaClient->receive_fresh_generator(
        generatorProof, freshGenerator);

    Twistpoint shortTermPublicKey =
        prsonaClient->get_short_term_public_key();

    vector<Proof> encryptedScoreProof;
    EGCiphertext encryptedScore =
        get_server_committed_val<EGCiphertext>(
            &generator,
            serverIPs,
            serverPorts,
            encryptedScoreProof,
            shortTermPublicKey,
            USER_TALLY_URI,
            USER_TALLY_COMMIT_URI);

    prsonaClient->receive_vote_tally(
        encryptedScoreProof,
        encryptedScore);

    Scalar threshold(0);

    vector<Proof> repProof =
        prsonaClient->generate_reputation_proof(
            threshold, numClients);

    string data =
        make_rep_proof_string(repProof, shortTermPublicKey, threshold);

    send_item(target, targetPort, data, REP_PROOF_URI);
}

int main(int argc, char *argv[])
{
    initialize_prsona_classes();

#if USE_SSL
    mg_init_library(0);
#else
    mg_init_library(MG_FEATURES_SSL);
#endif

    string id = "";
    if (argc > 1)
        id = argv[1];

    string seedStr;
    if (id.empty())
        seedStr = "default-client";
    else
    {
        seedStr = id;
        seedStr += "-client";
    }
    
    vector<string> serverIPs, clientIPs;
    vector<int> serverPorts, clientPorts;
    string selfIP, selfPortStr;
    int selfPort = 0;

    char buffer[46], *helper;
    ifstream serverConfig("cfg/serverIPs.cfg");
    while (!serverConfig.eof())
    {
        serverConfig.getline(buffer, 46);
        if (strlen(buffer) > 0)
        {
            helper = buffer;
            if (strchr(helper, ':'))
            {
                helper = strtok(helper, ":");
                serverIPs.push_back(string(helper));
                helper = strtok(NULL, ":");
                serverPorts.push_back(atoi(helper));
            }
            else
            {
                serverIPs.push_back(string(helper));
                serverPorts.push_back(atoi(PRSONA_PORT_STR));
            }
        }
    }

    ifstream clientConfig("cfg/clientIPs.cfg");
    while (!clientConfig.eof())
    {
        clientConfig.getline(buffer, 46);
        if (strlen(buffer) > 0)
        {
            helper = buffer;
            if (strchr(helper, ':'))
            {
                helper = strtok(helper, ":");
                clientIPs.push_back(string(helper));
                helper = strtok(NULL, ":");
                clientPorts.push_back(atoi(helper));
            }
            else
            {
                clientIPs.push_back(string(helper));
                clientPorts.push_back(atoi(PRSONA_PORT_STR));
            }
        }
    }

    string selfConfigFilename = "cfg/selfIP";
    if (!id.empty())
    {
        selfConfigFilename += "-";
        selfConfigFilename += id;
    }
    selfConfigFilename += ".cfg";

    ifstream selfConfig(selfConfigFilename);
    while (!selfConfig.eof())
    {
        selfConfig.getline(buffer, 46);
        if (strlen(buffer) > 0)
        {
            helper = buffer;
            if (strchr(helper, ':'))
            {
                helper = strtok(helper, ":");
                selfIP = helper;
                helper = strtok(NULL, ":");
                selfPortStr = helper;
                selfPort = atoi(helper);
            }
            else
            {
                selfIP = helper;
                selfPortStr = PRSONA_PORT_STR;
                selfPort = atoi(PRSONA_PORT_STR);
            }
        }
    }

    // Defaults
    size_t numServers = serverIPs.size();
    size_t numClients = clientIPs.size();
    bool maliciousServers = true;
    uniform_int_distribution<size_t> distribution(0, numServers - 1);

    const char *options[] = {"listening_ports", selfPortStr.c_str(), 0};

    if (argc > 2)
    {
        bool setting = argv[2][0] == 't' || argv[2][0] == 'T';
        maliciousServers = setting;
    }

    seed_seq seed(seedStr.begin(), seedStr.end());
    default_random_engine generator(seed);

    cout << "[" << seedStr << "] Establishing PRSONA client with the following parameters: " << endl;
    cout << "[" << seedStr << "] " << numServers << " PRSONA servers" << endl;
    cout << "[" << seedStr << "] " << numClients << " PRSONA clients" << endl;
    cout << "[" << seedStr << "] Servers are set to " << (maliciousServers ? "MALICIOUS" : "HBC") << " security" << endl;
    cout << "[" << seedStr << "] This client is at IP address: " << selfIP << ":" << selfPort << endl;
    cout << endl;

    // Set malicious flags where necessary
    if (maliciousServers)
        PrsonaBase::set_server_malicious();

    // Entities we operate with
    cout << "[" << seedStr << "] Creating PRSONA client." << endl;
    PrsonaClient *prsonaClient =
        create_client(&generator, serverIPs, serverPorts, numServers);
    CivetServer server(options);

    cout << "[" << seedStr << "] Setting up handlers for client." << endl;

    PrsonaClientWebSocketHandler wsHandler(prsonaClient, serverIPs, serverPorts, &generator);
    server.addWebSocketHandler("/ws", wsHandler);

    unique_lock<mutex> exitLock(exitSync.mtx);
    exitSync.val = 0;
    exitSync.val2 = 0;
    RemoteControlHandler exitHandler(&exitSync, "Client coming down!");
    server.addHandler(EXIT_URI, exitHandler);

    AltRemoteControlHandler triggerVoteHandler(CLIENT_MAKE_VOTE, &exitSync, "Client will make new votes!");
    server.addHandler(TRIGGER_VOTE_URI, triggerVoteHandler);

    AltRemoteControlHandler triggerRepHandler(CLIENT_MAKE_REP_PROOF, &exitSync, "Client will make new votes!");
    server.addHandler(TRIGGER_REP_URI, triggerRepHandler);

    cout << "[" << seedStr << "] Entering main ready loop." << endl;

    while (!exitSync.val)
    {
        while (!exitSync.val && !exitSync.val2)
            exitSync.cv.wait(exitLock);

        size_t whichServer = distribution(generator);
        string fullQuery, target;
        size_t colonLocation;
        int targetPort;

        switch (exitSync.val2)
        {
            case CLIENT_MAKE_VOTE:
                make_vote(
                    generator,
                    serverIPs[whichServer],
                    serverPorts[whichServer],
                    prsonaClient,
                    serverIPs,
                    serverPorts,
                    numClients);
                break;

            case CLIENT_MAKE_REP_PROOF:
                fullQuery = triggerRepHandler.getQuery();
                colonLocation = fullQuery.find(":");

                target = fullQuery.substr(0, colonLocation);
                targetPort = stoi(fullQuery.substr(colonLocation + 1));
                make_reputation_proof(
                    generator,
                    target,
                    targetPort,
                    prsonaClient,
                    serverIPs,
                    serverPorts,
                    numClients);
                break;

            default:
                break;
        }

        exitSync.val2 = 0;
    }

    mg_exit_library();

    delete prsonaClient;

    return 0;
}