#include <iostream>
#include <sstream>
#include <cstring>
#include <cstdio>

#include "networkServer.hpp"

void obtain_update_locks(
    std::unique_lock<std::mutex> &lck,
    const std::vector<std::string>& serverIPs,
    const std::string& selfIP,
    struct synchronization_tool *sync)
{
    size_t i = 0;
    while (i < serverIPs.size())
    {
        if (serverIPs[i] == selfIP)
        {
            lck.lock();
            i++;
            continue;
        }

        struct mg_connection *conn =
            mg_connect_websocket_client(
                serverIPs[i].c_str(),
                PRSONA_PORT,
                USE_SSL,
                NULL,
                0,
                UPDATE_LOCK_URI,
                "null",
                synchro_websocket_data_handler,
                synchro_websocket_close_handler,
                (void *) sync);

        if (!conn)
        {
            std::cerr << "Couldn't get server " << i << "'s lock" << std::endl;
            continue;
        }

        unique_lock<mutex> lck(sync->mtx);
        sync->val = 0;
        sync->val2 = 0;
        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        while (!sync->val2)
            sync->cv.wait(lck);

        mg_close_connection(conn);

        if (sync->val)
            i++;
    }
}

void release_update_locks(
    std::unique_lock<std::mutex> &lck,
    const std::vector<std::string>& serverIPs,
    const std::string& selfIP,
    struct synchronization_tool *sync)
{
    ssize_t i = serverIPs.size() - 1;
    while (i >= 0)
    {
        if (serverIPs[i] == selfIP)
        {
            lck.unlock();
            i--;
            continue;
        }

        struct mg_connection *conn =
            mg_connect_websocket_client(
                serverIPs[i].c_str(),
                PRSONA_PORT,
                USE_SSL,
                NULL,
                0,
                UPDATE_UNLOCK_URI,
                "null",
                synchro_websocket_data_handler,
                synchro_websocket_close_handler,
                (void *) sync);

        if (!conn)
        {
            std::cerr << "Couldn't get server " << i << "'s lock" << std::endl;
            continue;
        }

        unique_lock<mutex> lck(sync->mtx);
        sync->val = 0;
        sync->val2 = 0;
        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        while (!sync->val2)
            sync->cv.wait(lck);

        mg_close_connection(conn);

        if (sync->val)
            i--;
    }
}

std::string make_epoch_initiator_string(
    const std::vector<Proof>& generatorProof,
    const Twistpoint& nextGenerator)
{
    std::stringstream buffer;
    
    BinarySizeT sizeOfVector(generatorProof.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << generatorProof[i];
    
    buffer << nextGenerator;

    return buffer.str();
}

void read_epoch_initiator_string(
    const char *filename,
    std::vector<Proof>& generatorProof,
    Twistpoint& nextGenerator)
{
    ifstream file(filename);
    BinarySizeT sizeOfVector;

    generatorProof.clear();

    file >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        Proof currProof;
        file >> currProof;

        generatorProof.push_back;
    }
    
    file >> nextGenerator;
}

std::string make_epoch_update_string(
    const std::vector<std::vector<Proof>>& pi,
    const std::vector<std::vector<Twistpoint>>& permutationCommits,
    const std::vector<std::vector<Twistpoint>>& freshPseudonymCommits,
    const std::vector<std::vector<Twistpoint>>& freshPseudonymSeedCommits,
    const std::vector<std::vector<CurveBipoint>>& serverTallyCommits,
    const std::vector<std::vector<std::vector<TwistBipoint>>>& partwayVoteMatrixCommits,
    const std::vector<std::vector<std::vector<TwistBipoint>>>& finalVoteMatrixCommits,
    const std::vector<std::vector<Twistpoint>>& userTallyMaskCommits,
    const std::vector<std::vector<Twistpoint>>& userTallyMessageCommits,
    const std::vector<std::vector<Twistpoint>>& userTallySeedCommits,
    const Twistpoint& nextGenerator,
    bool doUserTallies)
{
    std::stringstream buffer;
    
    BinarySizeT sizeOfVectorI, sizeOfVectorJ;

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

        for (size_t j = 0; j < sizeOfVectorJ.val(); j++)
            buffer << pi[i][j];
    }

    sizeOfVectorI.set(permutationCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            buffer << permutationCommits[i][j];

    sizeOfVectorI.set(freshPseudonymCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            buffer << freshPseudonymCommits[i][j];

    sizeOfVectorI.set(freshPseudonymSeedCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            buffer << freshPseudonymSeedCommits[i][j];

    sizeOfVectorI.set(serverTallyCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            buffer << serverTallyCommits[i][j];

    sizeOfVectorI.set(partwayVoteMatrixCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            for (size_t k = 0; k < sizeOfVectorI.val(); k++)
                buffer << partwayVoteMatrixCommits[i][j][k];

    sizeOfVectorI.set(finalVoteMatrixCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            for (size_t k = 0; k < sizeOfVectorI.val(); k++)
                buffer << finalVoteMatrixCommits[i][j][k];

    sizeOfVectorI.set(userTallyMaskCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            buffer << userTallyMaskCommits[i][j];

    sizeOfVectorI.set(userTallyMessageCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            buffer << userTallyMessageCommits[i][j];

    sizeOfVectorI.set(userTallySeedCommits.size());
    buffer << sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
            buffer << userTallySeedCommits[i][j];

    buffer << nextGenerator;

    BinaryBool flag(doUserTallies);
    buffer << flag;

    return buffer.str();
}

bool read_epoch_update_string(
    const char *filename,
    std::vector<std::vector<Proof>>& pi,
    std::vector<std::vector<Twistpoint>>& permutationCommits,
    std::vector<std::vector<Twistpoint>>& freshPseudonymCommits,
    std::vector<std::vector<Twistpoint>>& freshPseudonymSeedCommits,
    std::vector<std::vector<CurveBipoint>>& serverTallyCommits,
    std::vector<std::vector<std::vector<TwistBipoint>>>& partwayVoteMatrixCommits,
    std::vector<std::vector<std::vector<TwistBipoint>>>& finalVoteMatrixCommits,
    std::vector<std::vector<Twistpoint>>& userTallyMaskCommits,
    std::vector<std::vector<Twistpoint>>& userTallyMessageCommits,
    std::vector<std::vector<Twistpoint>>& userTallySeedCommits,
    Twistpoint& nextGenerator)
{
    ifstream file(filename);
    BinarySizeT sizeOfVectorI, sizeOfVectorJ;
    
    pi.clear();
    permutationCommits.clear();
    freshPseudonymCommits.clear();
    freshPseudonymSeedCommits.clear();
    serverTallyCommits.clear();
    partwayVoteMatrixCommits.clear();
    finalVoteMatrixCommits.clear();
    userTallyMaskCommits.clear();
    userTallyMessageCommits.clear();
    userTallySeedCommits.clear();


    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<Proof> currRow;

        file >> sizeOfVectorJ;
        for (size_t j = 0; j < sizeOfVectorJ.val(); j++)
        {
            Proof currProof;
            file >> currProof;

            currRow.push_back(currProof);
        }

        pi.push_back(currRow);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<Twistpoint> currRow;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            Twistpoint currCommit;
            file >> currCommit;

            currRow.push_back(currCommit);
        }
        permutationCommits.push_back(currRow);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<Twistpoint> currRow;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            Twistpoint currCommit;
            file >> currCommit;

            currRow.push_back(currCommit);
        }
        freshPseudonymCommits.push_back(currRow);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<Twistpoint> currRow;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            Twistpoint currCommit;
            file >> currCommit;

            currRow.push_back(currCommit);
        }
        freshPseudonymSeedCommits.push_back(currRow);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<CurveBipoint> currRow;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            CurveBipoint currCommit;
            file >> currCommit;

            currRow.push_back(currCommit);
        }
        serverTallyCommits.push_back(currRow);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<std::vector<TwistBipoint>> currMatrix;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            std::vector<TwistBipoint> currRow;
            for (size_t k = 0; k < sizeOfVectorI.val(); k++)
            {
                TwistBipoint currCommit;
                file >> currCommit;

                currRow.push_back(currCommit);    
            }
            currMatrix.push_back(currRow);
        }
        partwayVoteMatrixCommits.push_back(currMatrix);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<std::vector<TwistBipoint>> currMatrix;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            std::vector<TwistBipoint> currRow;
            for (size_t k = 0; k < sizeOfVectorI.val(); k++)
            {
                TwistBipoint currCommit;
                file >> currCommit;

                currRow.push_back(currCommit);    
            }
            currMatrix.push_back(currRow);
        }
        finalVoteMatrixCommits.push_back(currMatrix);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<Twistpoint> currRow;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            Twistpoint currCommit;
            file >> currCommit;

            currRow.push_back(currCommit);
        }
        userTallyMaskCommits.push_back(currRow);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<Twistpoint> currRow;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            Twistpoint currCommit;
            file >> currCommit;

            currRow.push_back(currCommit);
        }
        userTallyMessageCommits.push_back(currRow);
    }

    file >> sizeOfVectorI;
    for (size_t i = 0; i < sizeOfVectorI.val(); i++)
    {
        std::vector<Twistpoint> currRow;
        for (size_t j = 0; j < sizeOfVectorI.val(); j++)
        {
            Twistpoint currCommit;
            file >> currCommit;

            currRow.push_back(currCommit);
        }
        userTallySeedCommits.push_back(currRow);
    }

    file >> nextGenerator;
    
    BinaryBool doUserTallies;
    file >> doUserTallies;

    return doUserTallies.val();
}

void distribute_epoch_updates(
    const std::string& recipient,
    const std::string& data,
    const struct synchronization_tool* sync)
{
    bool flag = false;
    while (!flag)
    {
        struct mg_connection *conn =
            mg_connect_websocket_client(
                serverIPs[i].c_str(),
                PRSONA_PORT,
                USE_SSL,
                NULL,
                0,
                ACCEPT_EPOCH_UPDATES_URI,
                "null",
                synchro_websocket_data_handler,
                empty_websocket_close_handler,
                (void *) sync);

        if (!conn)
        {
            std::cerr << "Trouble giving epoch updates to server at " << recipient << std::endl;
            continue;
        }
        
        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_BINARY,
            data.c_str(),
            data.length());

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

        mg_close_connection(conn);

        flag = true;
    }   
}

/*
 * CONSTRUCTORS
 */

PrsonaServerWebSocketHandler::PrsonaServerWebSocketHandler(
    const PrsonaServer *prsonaServer,
    const std::mutex *updateMtx,
    const std::atomic<size_t> *epochNum,
    const std::vector<std::string> &serverIPs,
    const std::string &selfIP)
: prsonaServer(prsonaServer), updateMtx(updateMtx), epochNum(epochNum),
    serverIPs(serverIPs), selfIP(selfIP)
{ /* */ }

virtual bool PrsonaServerWebSocketHandler::handleConnection(
    CivetServer *server,
    const struct mg_connection *conn)
{
    const struct mg_request_info *info = mg_get_request_info(conn);
    
    bool flag = info->query_string && info->query_string[0] >= PRSONA_ADD_CLIENT && info->query_string[0] <= PRSONA_RECEIVE_PARTIAL_DECRYPTION;
    flag = flag || (info->query_string && info->query_string[0] == PRSONA_GET_FRESH_GENERATOR);
    flag = flag || (info->query_string && info->query_string[0] == PRSONA_GET_EG_BLIND_GENERATOR);

    return flag;
}

virtual void PrsonaServerWebSocketHandler::handleReadyState(
    CivetServer *server,
    struct mg_connection *conn)
{
    const struct mg_request_info *info = mg_get_request_info(conn);

    switch (info->query_string[0])
    {
        case PRSONA_ADD_CLIENT:
        case PRSONA_RECEIVE_VOTE:
        case PRSONA_GET_VOTES_BY:
        case PRSONA_GET_USER_TALLY:
        case PRSONA_GET_SERVER_TALLY:
        case PRSONA_GET_VOTE_ROW_COMMITMENT:
        case PRSONA_GET_USER_TALLY_COMMITMENT:
        case PRSONA_GET_SERVER_TALLY_COMMITMENT:
        case PRSONA_ADD_CURR_SEED_TO_GENERATOR:
        case PRSONA_SET_FRESH_GENERATOR:
        case PRSONA_ADD_RAND_SEED_TO_GENERATOR:
        case PRSONA_SET_EG_BLIND_GENERATOR:
        case PRSONA_EPOCH_UPDATE:
        case PRSONA_NEW_USER_UPDATE:
        case PRSONA_RECEIVE_PARTIAL_DECRYPTION:
            set_temp_filename(conn);
            break;

        default:
            mg_set_user_connection_data(conn, NULL);
            break;
    }
}

virtual bool PrsonaServerWebSocketHandler::handleData(
    CivetServer *server,
    struct mg_connection *conn,
    int bits,
    char *data,
    size_t data_len)
{
    char *filename = (char *) mg_get_user_connection_data(conn);

    if ((bits & 0xf) == MG_WEBSOCKET_OPCODE_DATACOMPLETE)
    {
        generate_response(conn, filename);
        return false;
    }

    if ((bits & 0xf) != MG_WEBSOCKET_OPCODE_BINARY && (bits & 0xf) != MG_WEBSOCKET_OPCODE_CONTINUATION)
    {
        std::cerr << "Unknown opcode: failing." << std::endl;
        return false;
    }

    FILE *currFile = fopen(filename, "ab");
    fwrite(data, sizeof(char), data_len, currFile);
    fclose(currFile);

    return true;
}

void PrsonaServerWebSocketHandler::generate_response(
    struct mg_connection *conn,
    char *filename)
{
    const struct mg_request_info *info = mg_get_request_info(conn);

    switch (info->query_string[0])
    {
        case PRSONA_ADD_CLIENT:
            add_new_client(conn, filename);
            break;

        case PRSONA_RECEIVE_VOTE:
            receive_vote(filename);
            break;

        case PRSONA_GET_BGN_PUBKEY:
            get_bgn_public_key(conn);
            break;

        case PRSONA_GET_NUM_CLIENTS:
            get_num_clients(conn);
            break;

        case PRSONA_GET_NUM_SERVERS:
            get_num_servers(conn);
            break;

        case PRSONA_GET_VOTES_BY:
            get_current_votes_by(conn, filename);
            break;

        case PRSONA_GET_ALL_VOTES:
            get_all_current_votes(conn);
            break;

        case PRSONA_GET_USER_TALLY:
            get_current_user_encrypted_tally(conn, filename);
            break;

        case PRSONA_GET_SERVER_TALLY:
            get_current_server_encrypted_tally(conn, filename);
            break;

        case PRSONA_GET_PSEUDONYMS:
            get_current_pseudonyms(conn);
            break;

        case PRSONA_GET_VOTE_ROW_COMMITMENT:
            get_vote_row_commitment(conn, filename);
            break;

        case PRSONA_GET_VOTE_MATRIX_COMMITMENT:
            get_vote_matrix_commitment(conn);
            break;

        case PRSONA_GET_USER_TALLY_COMMITMENT:
            get_user_tally_commitment(conn, filename);
            break;

        case PRSONA_GET_SERVER_TALLY_COMMITMENT:
            get_server_tally_commitment(conn, filename);
            break;

        case PRSONA_GET_PSEUDONYMS_COMMITMENT:
            get_pseudonyms_commitment(conn);
            break;

        case PRSONA_GET_BGN_DETAILS:
            get_bgn_details(conn);
            break;

        case PRSONA_ADD_CURR_SEED_TO_GENERATOR:
            add_seed_to_generator(conn, filename, true);
            break;

        case PRSONA_SET_FRESH_GENERATOR:
            set_generator(filename, true);
            break;

        case PRSONA_ADD_RAND_SEED_TO_GENERATOR:
            add_seed_to_generator(conn, filename, false);
            break;

        case PRSONA_SET_EG_BLIND_GENERATOR:
            set_generator(filename, false);
            break;

        case PRSONA_EPOCH_BUILD_UP:
            build_up_midway_pseudonyms(conn);
            break;

        case PRSONA_EPOCH_BREAK_DOWN:
            break_down_midway_pseudonyms(conn);
            break;

        case PRSONA_EPOCH_UPDATE:
            accept_epoch_updates(filename);
            break;

        case PRSONA_NEW_USER_UPDATE:
            import_new_user_update(filename);
            break;

        case PRSONA_GET_PARTIAL_DECRYPTION:
            get_partial_decryption(conn, filename);
            break;

        case PRSONA_RECEIVE_PARTIAL_DECRYPTION:
            receive_tallied_scores(conn, filename);
            break;

        case PRSONA_GET_FRESH_GENERATOR:
            get_generator(conn, true);
            break;

        case PRSONA_GET_EG_BLIND_GENERATOR:
            get_generator(conn, false);
            break;

        default:
            break;
    }
}

virtual void PrsonaServerWebSocketHandler::handleClose(
    CivetServer *server,
    const struct mg_connection *conn)
{
    char *filename = (char *) mg_get_user_connection_data(conn);
    if (!filename)
        return;

    remove(filename);
    delete filename;
}

/*
 * BASIC PUBLIC SYSTEM INFO GETTERS
 */

void PrsonaServerWebSocketHandler::get_bgn_public_key(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    BGNPublicKey pubKey = prsonaServer->get_bgn_public_key();
    buffer << pubKey;
    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_num_clients(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    BinarySizeT numClients(prsonaServer->get_num_clients());
    buffer << numClients;
    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_num_servers(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    BinarySizeT numServers(prsonaServer->get_num_servers());
    buffer << numServers;
    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

/*
 * ENCRYPTED DATA GETTERS
 */

void PrsonaServerWebSocketHandler::get_current_votes_by(
    struct mg_connection *conn, const char *filename) const
{
    ifstream file(filename);

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    std::stringstream buffer;
    std::string data;

    Proof pi;
    std::vector<TwistBipoint> votes =
        prsonaServer->get_current_votes_by(pi, shortTermPublicKey);
    
    buffer << pi;
    
    BinarySizeT sizeOfVector(votes.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << votes[i];

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_all_current_votes(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    Proof pi;
    std::vector<std::vector<TwistBipoint>> votes =
        prsonaServer->get_all_current_votes(pi);
    
    buffer << pi;
    
    BinarySizeT sizeOfVector(votes.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        for (size_t j = 0; j < sizeOfVector.val(); j++)
            buffer << votes[i][j];

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_current_user_encrypted_tally(
    struct mg_connection *conn, const char *filename) const
{
    ifstream file(filename);

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    std::stringstream buffer;
    std::string data;

    Proof pi;
    EGCiphertext tally =
        prsonaServer->get_current_user_encrypted_tally(pi, shortTermPublicKey);
    
    buffer << pi;
    buffer << tally;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_current_server_encrypted_tally(
    struct mg_connection *conn, const char *filename) const
{
    ifstream file(filename);

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    std::stringstream buffer;
    std::string data;

    Proof pi;
    CurveBipoint tally =
        prsonaServer->get_current_server_encrypted_tally(pi, shortTermPublicKey);
    
    buffer << pi;
    buffer << tally;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_current_pseudonyms(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    Proof pi;
    std::vector<Twistpoint> pseudonyms =
        prsonaServer->get_current_pseudonyms(pi);
    
    buffer << pi;

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

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

/*
 * PROOF COMMITMENT GETTERS
 */

void PrsonaServerWebSocketHandler::get_vote_row_commitment(
    struct mg_connection *conn, const char *filename) const
{
    ifstream file(filename);

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    std::stringstream buffer;
    std::string data;

    Proof pi =
        prsonaServer->get_vote_row_commitment(shortTermPublicKey);
    
    buffer << pi;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_vote_matrix_commitment(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    Proof pi =
        prsonaServer->get_vote_matrix_commitment();
    
    buffer << pi;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_user_tally_commitment(
    struct mg_connection *conn, const char *filename) const
{
    ifstream file(filename);

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    std::stringstream buffer;
    std::string data;

    Proof pi =
        prsonaServer->get_user_tally_commitment(shortTermPublicKey);
    
    buffer << pi;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_server_tally_commitment(
    struct mg_connection *conn, const char *filename) const
{
    ifstream file(filename);

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    std::stringstream buffer;
    std::string data;

    Proof pi =
        prsonaServer->get_server_tally_commitment(shortTermPublicKey);
    
    buffer << pi;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::get_pseudonyms_commitment(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    Proof pi =
        prsonaServer->get_pseudonyms_commitment();
    
    buffer << pi;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::distribute_new_user_updates(
    std::vector<Proof> proofOfValidAddition,
    std::vector<CurveBipoint> previousVoteTallies,
    std::vector<Twistpoint> currentPseudonyms,
    std::vector<EGCiphertext> currentUserEncryptedTallies,
    std::vector<std::vector<TwistBipoint>> voteMatrix) const
{
    std::stringstream buffer;
    std::string data;
    BinarySizeT sizeOfVector;

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

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

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

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

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

    data = buffer.str();

    size_t i = 0;
    while (i < serverIPs.size())
    {
        if (serverIPs[i] == selfIP)
            continue;

        struct mg_connection *conn =
            mg_connect_websocket_client(
                serverIPs[i].c_str(),
                PRSONA_PORT,
                USE_SSL,
                NULL,
                0,
                GIVE_NEW_USER_URI,
                "null",
                synchro_websocket_data_handler,
                synchro_websocket_close_handler,
                (void *) distributeSync);

        if (!conn)
        {
            std::cerr << "Couldn't give server " << i << " new user" << std::endl;
            continue;
        }

        unique_lock<mutex> syncLock(distributeSync->mtx);
        distributeSync->val = 0;
        distributeSync->val2 = 0;
        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_BINARY,
            data.c_str(),
            data.length());

        while (!distributeSync->val2)
            distributeSync->cv.wait(syncLock);

        mg_close_connection(conn);

        if (distributeSync->val)
            i++;
    }
}

void PrsonaServerWebSocketHandler::distribute_new_vote(
    std::vector<Proof> pi;
    std::vector<TwistBipoint> newVotes,
    Twistpoint shortTermPublicKey) const
{
    std::stringstream buffer;
    std::string data;
    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 flag(false);
    buffer << flag;

    data = buffer.str();

    size_t i = 0;
    while (i < serverIPs.size())
    {
        if (serverIPs[i] == selfIP)
            continue;

        struct mg_connection *conn =
            mg_connect_websocket_client(
                serverIPs[i].c_str(),
                PRSONA_PORT,
                USE_SSL,
                NULL,
                0,
                GIVE_NEW_VOTE_URI,
                "null",
                synchro_websocket_data_handler,
                synchro_websocket_close_handler,
                (void *) distributeSync);

        if (!conn)
        {
            std::cerr << "Couldn't give server " << i << " new user" << std::endl;
            continue;
        }

        unique_lock<mutex> syncLock(distributeSync->mtx);
        distributeSync->val = 0;
        distributeSync->val2 = 0;
        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_BINARY,
            data.c_str(),
            data.length());

        while (!distributeSync->val2)
            distributeSync->cv.wait(syncLock);

        mg_close_connection(conn);

        if (distributeSync->val)
            i++;
    }
}

/*
 * CLIENT INTERACTIONS
 */

void PrsonaServerWebSocketHandler::add_new_client(
    struct mg_connection *conn, const char *filename)
{
    ifstream file(filename);

    Proof proofOfValidKey;
    file >> proofOfValidKey;

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    std::unique_lock<std::mutex> lck(*updateMtx, std::defer_lock);
    obtain_update_locks(
        lck,
        serverIPs,
        selfIP,
        &updateSync);

    std::vector<Proof> proofOfValidAddition;
    prsonaServer->add_new_client(
        proofOfValidAddition,
        proofOfValidKey,
        shortTermPublicKey);

    std::vector<CurveBipoint> previousVoteTallies;
    std::vector<Twistpoint> currentPseudonyms;
    std::vector<EGCiphertext> currentUserEncryptedTallies;
    std::vector<std::vector<TwistBipoint>> voteMatrix;
    prsonaServer->export_new_user_update(
        previousVoteTallies,
        currentPseudonyms,
        currentUserEncryptedTallies,
        voteMatrix);

    distribute_new_user_updates(
        proofOfValidAddition,
        previousVoteTallies,
        currentPseudonyms,
        currentUserEncryptedTallies,
        voteMatrix);

    release_update_locks(
        lck,
        serverIPs,
        selfIP,
        &updateSync);

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

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::receive_vote(
    const char *filename)
{
    ifstream file(filename);

    BinarySizeT sizeOfVector;
    file >> sizeOfVector;

    std::vector<Proof> pi;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        Proof currProof;
        file >> currProof;

        pi.push_back(currProof);
    }

    file >> sizeOfVector;

    std::vector<TwistBipoint> newVotes;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        TwistBipoint currVote;
        file >> currVote;

        newVotes.push_back(currVote);
    }

    Twistpoint shortTermPublicKey;
    file >> shortTermPublicKey;

    BinaryBool shouldDeal;
    file >> shouldDeal;

    std::unique_lock<std::mutex> lck(*updateMtx, std::defer_lock);
    if (shouldDeal.val())
    {
        obtain_update_locks(
            lck,
            serverIPs,
            selfIP,
            &distributeSync);
    }

    prsonaServer->receive_vote(
        pi,
        newVotes,
        shortTermPublicKey);

    if (shouldDeal.val())
    {
        distribute_new_vote(pi, newVotes, shortTermPublicKey);
        release_update_locks(
            lck,
            serverIPs,
            selfIP,
            &distributeSync);
    }

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
}

/*
 * CONSTRUCTOR HELPERS
 */

void PrsonaServerWebSocketHandler::get_bgn_details(
    struct mg_connection *conn) const
{
    std::stringstream buffer;
    std::string data;

    const BGN& sharedBGN = prsonaServer->get_bgn_details();
    buffer << sharedBGN;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::add_seed_to_generator(
    struct mg_connection *conn, const char *filename, bool fresh) const
{
    ifstream file(filename);

    Twistpoint currGenerator;
    file >> currGenerator;

    std::stringstream buffer;
    std::string data;

    std::vector<Proof> pi;
    if (fresh)
    {
        currGenerator =
            prsonaServer->add_curr_seed_to_generator(pi, currGenerator);
    }
    else
    {
        currGenerator =
            prsonaServer->add_rand_seed_to_generator(pi, currGenerator);
    }
    
    buffer << pi[0];
    buffer << currGenerator;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
}

void PrsonaServerWebSocketHandler::set_generator(
    const char *filename, bool fresh)
{
    ifstream file(filename);

    BinarySizeT sizeOfVector;
    file >> sizeOfVector;

    std::vector<Proof> pi;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        Proof currProof;
        file >> currProof;

        pi.push_back(currProof);
    }

    Twistpoint newGenerator;
    file >> newGenerator;

    if (fresh)
        prsonaServer->initialize_fresh_generator(pi, newGenerator);
    else
        prsonaServer->set_EG_blind_generator(pi, newGenerator);
}

/*
 * EPOCH ROUNDS
 */

void PrsonaServerWebSocketHandler::build_up_midway_pseudonyms(
    struct mg_connection *conn, const char *filename)
{
    std::vector<std::vector<Proof>> generatorProofHolder;
    std::vector<Proof> generatorProof;
    Twistpoint nextGenerator;

    read_epoch_initiator_string(
        filename,
        generatorProof,
        nextGenerator);

    generatorProofHolder.push_back(generatorProof);

    std::vector<std::vector<std::vector<Proof>>> pi;
    std::vector<std::vector<std::vector<Twistpoint>>> permutationCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymSeedCommits;
    std::vector<std::vector<std::vector<CurveBipoint>>> serverTallyCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> partwayVoteMatrixCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> finalVoteMatrixCommits;

    pi.push_back(generatorProofHolder);

    prsonaServer->build_up_midway_pseudonyms(
        pi,
        permutationCommits,
        freshPseudonymCommits,
        freshPseudonymSeedCommits,
        serverTallyCommits,
        partwayVoteMatrixCommits,
        finalVoteMatrixCommits,
        nextGenerator);

    std::vector<std::vector<Twistpoint>> userTallyMaskCommits, userTallyMessageCommits, userTallySeedCommits;

    string data =
        make_epoch_update_string(
            pi[1],
            permutationCommits[0],
            freshPseudonymCommits[0],
            freshPseudonymSeedCommits[0],
            serverTallyCommits[0],
            partwayVoteMatrixCommits[0],
            finalVoteMatrixCommits[0],
            userTallyMaskCommits,
            userTallyMessageCommits,
            userTallySeedCommits,
            nextGenerator,
            false);

    struct synchronization_tool epochSync;
    epochSync->val = 1;
    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        if (serverIPs[i] == selfIP)
            continue;

        distribute_epoch_updates(
            serverIPs[i],
            data,
            &epochSync);
    }

    unique_lock<mutex> lck(epochSync);
    while (epochSync.val < serverIPs.size())
        epochSync.cv.wait(lck);
    
    data = make_epoch_initiator_string(
        pi[0][0],
        nextGenerator);

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
}

void PrsonaServerWebSocketHandler::break_down_midway_pseudonyms(
    struct mg_connection *conn, const char *filename)
{
    std::vector<Proof> generatorProof;
    Twistpoint nextGenerator;

    read_epoch_initiator_string(
        filename,
        generatorProof,
        nextGenerator);

    std::vector<std::vector<std::vector<Proof>>> pi;
    std::vector<std::vector<std::vector<Twistpoint>>> permutationCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymSeedCommits;
    std::vector<std::vector<std::vector<CurveBipoint>>> serverTallyCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> partwayVoteMatrixCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> finalVoteMatrixCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallyMaskCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallyMessageCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallySeedCommits;

    prsonaServer->break_down_midway_pseudonyms(
        generatorProof,
        pi,
        permutationCommits,
        freshPseudonymCommits,
        freshPseudonymSeedCommits,
        serverTallyCommits,
        partwayVoteMatrixCommits,
        finalVoteMatrixCommits,
        userTallyMaskCommits,
        userTallyMessageCommits,
        userTallySeedCommits,
        nextGenerator);

    string data =
        make_epoch_update_string(
            pi[0],
            permutationCommits[0],
            freshPseudonymCommits[0],
            freshPseudonymSeedCommits[0],
            serverTallyCommits[0],
            partwayVoteMatrixCommits[0],
            finalVoteMatrixCommits[0],
            userTallyMaskCommits[0],
            userTallyMessageCommits[0],
            userTallySeedCommits[0],
            nextGenerator,
            true);

    struct synchronization_tool epochSync;
    epochSync.val = 1;
    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        if (serverIPs[i] == selfIP)
            continue;

        distribute_epoch_updates(
            serverIPs[i],
            data,
            &epochSync);
    }

    unique_lock<mutex> lck(epochSync.mtx);
    while (epochSync.val < serverIPs.size())
        epochSync.cv.wait(lck);

    epochNum->fetch_add(1);

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
}

void PrsonaServerWebSocketHandler::accept_epoch_updates(
    struct mg_connection *conn, const char *filename)
{
    std::vector<std::vector<Proof>> pi;
    std::vector<std::vector<Twistpoint>> permutationCommits;
    std::vector<std::vector<Twistpoint>> freshPseudonymCommits;
    std::vector<std::vector<Twistpoint>> freshPseudonymSeedCommits;
    std::vector<std::vector<CurveBipoint>> serverTallyCommits;
    std::vector<std::vector<std::vector<TwistBipoint>>> partwayVoteMatrixCommits;
    std::vector<std::vector<std::vector<TwistBipoint>>> finalVoteMatrixCommits;
    std::vector<std::vector<Twistpoint>> userTallyMaskCommits;
    std::vector<std::vector<Twistpoint>> userTallyMessageCommits;
    std::vector<std::vector<Twistpoint>> userTallySeedCommits;
    Twistpoint nextGenerator;
    bool doUserTallies =
        read_epoch_update_string(
            filename,
            pi,
            permutationCommits,
            freshPseudonymCommits,
            freshPseudonymSeedCommits,
            serverTallyCommits,
            partwayVoteMatrixCommits,
            finalVoteMatrixCommits,
            userTallyMaskCommits,
            userTallyMessageCommits,
            userTallySeedCommits,
            nextGenerator);

    prsonaServer->accept_epoch_updates(
        pi,
        permutationCommits,
        freshPseudonymCommits,
        freshPseudonymSeedCommits,
        serverTallyCommits,
        partwayVoteMatrixCommits,
        finalVoteMatrixCommits,
        userTallyMaskCommits,
        userTallyMessageCommits,
        userTallySeedCommits,
        nextGenerator,
        doUserTallies);

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
}

/*
 * DATA MAINTENANCE
 */

void PrsonaServerWebSocketHandler::import_new_user_update(
    const char *filename)
{
    std::vector<Proof> proofOfValidAddition;
    std::vector<CurveBipoint> previousVoteTallies;
    std::vector<Twistpoint> currentPseudonyms;
    std::vector<EGCiphertext> currentUserEncryptedTallies;
    std::vector<std::vector<TwistBipoint>> voteMatrix;

    ifstream file(filename);

    BinarySizeT sizeOfVector;
    
    file >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        Proof currProof;
        file >> currProof;

        proofOfValidAddition.push_back(currProof);
    }

    file >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        CurveBipoint currTally;
        file >> currTally;

        previousVoteTallies.push_back(currTally);
    }

    file >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        Twistpoint currNym;
        file >> currNym;

        currentPseudonyms.push_back(currNym);
    }

    file >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        EGCiphertext currTally;
        file >> currTally;

        currentUserEncryptedTallies.push_back(currTally);
    }

    file >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        std::vector<TwistBipoint> currRow;
        for (size_t j = 0; j < sizeOfVector.val(); j++)
        {
            TwistBipoint currVote;
            file >> currVote;

            currRow.push_back(currVote);
        }
        voteMatrix.push_back(currRow);
    }

    prsonaServer->import_new_user_update(
        proofOfValidAddition,
        previousVoteTallies,
        currentPseudonyms,
        currentUserEncryptedTallies,
        voteMatrix);
}

void PrsonaServerWebSocketHandler::get_partial_decryption(
    struct mg_connection *conn)
{
    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
}

void PrsonaServerWebSocketHandler::receive_tallied_scores(
    struct mg_connection *conn, const char *filename)
{
    ifstream file(filename);

    std::vector<EGCiphertext> userScores;
    std::vector<CurveBipoint> serverScores;

    BinarySizeT sizeOfVector;
    file >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        EGCiphertext currScore;
        file >> currScore;

        userScores.push_back(currScore);
    }

    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        CurveBipoint currScore;
        file >> currScore;

        serverScores.push_back(currScore);
    }

    prsonaServer->receive_tallied_scores(userScores, serverScores);

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
}

void PrsonaServerWebSocketHandler::get_generator(
    struct mg_connection *conn, bool fresh)
{
    Twistpoint generator;
    std::vector<Proof> pi;
    if (fresh)
        generator = prsonaServer->get_fresh_generator(pi);
    else
        generator = prsonaServer->get_blinding_generator(pi);

    std::stringstream buffer;
    std::string data;

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

    buffer << generator;

    data = buffer.str();

    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_BINARY, data.c_str(), data.length());
    mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
}