#ifndef __PRSONA_NETWORK_SERVER_HPP
#define __PRSONA_NETWORK_SERVER_HPP

#include <random>
#include <string>
#include <vector>
#include <atomic>

#include "server.hpp"
#include "networking.hpp"

/* "PUBLIC" FUNCTIONS */

// CREATOR FOR A NEW SERVER
PrsonaServer *create_server(
    std::default_random_engine& rng,
    std::string dealerIP,
    int dealerPort,
    bool bgnDealer,
    size_t numServers);

// CHECK IN FUNCTION USED FOR SYNCHRONIZATION IN SETUP
void check_in_with_dealer(
    std::string dealerIP,
    int dealerPort);

// INITIATER FOR SHARED GLOBAL VALUES
void initiate_generators(
    std::default_random_engine& rng,
    PrsonaServer* prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort);

// FUNCTION TO PERFORM OPERATIONS FOR EXPERIMENT
void make_epoch(
    std::default_random_engine& rng,
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    std::mutex& updateMtx,
    std::atomic<size_t>& epochNum,
    const CivetServer& civetServer,
    std::mutex& buildUpOutputMtx,
    const std::string& buildUpOutputFilename,
    std::mutex& breakDownOutputMtx,
    const std::string& breakDownOutputFilename,
    std::mutex& fullOutputMtx,
    const std::string& fullOutputFilename,
    std::mutex& usageMtx,
    const std::string& usageFilename);

/* "PRIVATE" FUNCTIONS */

// SHARED GLOBAL UPDATE LOCK GETTERS AND RELEASERS
void obtain_update_locks(
    std::unique_lock<std::mutex> &updateLock,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    std::vector<size_t>& bandwidthData);

void release_update_locks(
    std::unique_lock<std::mutex> &updateLock,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    std::vector<size_t>& bandwidthData);

// GETTER FOR DEALER VALUE
BGN get_bgn_private_key(
    std::default_random_engine& rng,
    std::string dealerIP,
    int dealerPort);

// HELPERS TO INITIATE SHARED GLOBAL VALUES
Twistpoint make_generator(
    std::default_random_engine& rng,
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    bool fresh,
    std::vector<Proof>& pi);

void distribute_generator(
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    bool fresh,
    const std::vector<Proof>& pi,
    const Twistpoint& generator);

// HELPERS FOR EPOCH CALCULATIONS
std::vector<Proof> epoch_build_up(
    std::default_random_engine& rng,
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    Twistpoint& nextGenerator,
    const CivetServer& civetServer,
    std::mutex& outputMtx,
    const std::string& outputFilename,
    std::mutex& usageMtx,
    const std::string& usageFilename,
    std::vector<size_t>& bandwidthData);

std::vector<Proof> hbc_epoch_build_up(
    std::default_random_engine& rng,
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    Twistpoint& nextGenerator,
    const CivetServer& civetServer,
    std::mutex& outputMtx,
    const std::string& outputFilename,
    std::mutex& usageMtx,
    const std::string& usageFilename,
    std::vector<size_t>& overallBandwidthData);

void epoch_break_down(
    std::default_random_engine& rng,
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    const std::vector<Proof>& generatorProof,
    const Twistpoint& nextGenerator,
    const CivetServer& civetServer,
    std::mutex& outputMtx,
    const std::string& outputFilename,
    std::mutex& usageMtx,
    const std::string& usageFilename,
    std::vector<size_t>& bandwidthData);

void hbc_epoch_break_down(
    std::default_random_engine& rng,
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    const std::vector<Proof>& generatorProof,
    const Twistpoint& nextGenerator,
    const CivetServer& civetServer,
    std::mutex& outputMtx,
    const std::string& outputFilename,
    std::mutex& usageMtx,
    const std::string& usageFilename,
    std::vector<size_t>& bandwidthData);

// HELPERS FOR EPOCH HELPERS
Twistpoint initiate_epoch_updates(
    std::default_random_engine& rng,
    const std::string& recipient,
    int recipientPort,
    const std::string& data,
    bool isBreakdown,
    std::vector<std::vector<Proof>>& generatorProofHolder,
    std::vector<size_t>& bandwidthData);

struct mg_connection *distribute_epoch_updates(
    const std::string& recipient,
    int recipientPort,
    const std::string& data,
    struct synchronization_tool* sync);

// SCORE TALLYING AND DISTRIBUTION HELPERS
void tally_scores(
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    const Twistpoint& nextGenerator,
    std::vector<EGCiphertext>& userTallyScores,
    std::vector<CurveBipoint>& serverTallyScores,
    std::vector<size_t>& bandwidthData);

void distribute_tallied_scores(
    PrsonaServer *prsonaServer,
    const std::vector<std::string>& serverIPs,
    const std::vector<int>& serverPorts,
    const std::string& selfIP,
    int selfPort,
    const Twistpoint& nextGenerator,
    const std::vector<EGCiphertext>& userTallyScores,
    const std::vector<CurveBipoint>& serverTallyScores,
    std::vector<size_t>& bandwidthData);

// FILE I/O HELPERS
BGN get_bgn_private_key_from_file(
    const char *filename);

Twistpoint get_generator_from_file(
    const char *filename,
    Proof& pi);

Twistpoint get_generator_from_file(
    const char *filename,
    std::vector<Proof>& pi);

// EPOCH DATA SERIALIZERS/UN-SERIALIZERS
std::string make_epoch_initiator_string(
    const std::vector<Proof>& generatorProof,
    const Twistpoint& nextGenerator);

ssize_t read_epoch_initiator_string(
    const char *filename,
    std::vector<Proof>& generatorProof,
    Twistpoint& 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::string make_hbc_epoch_update_string(
    const std::vector<Proof>& generatorProof,
    const std::vector<Twistpoint>& newFreshPseudonyms,
    const std::vector<CurveBipoint>& newServerTallies,
    const std::vector<std::vector<TwistBipoint>>& newVoteMatrix,
    const std::vector<EGCiphertext>& newUserTallies,
    const Twistpoint& nextGenerator,
    bool doUserTallies);

ssize_t 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,
    bool& doUserTallies);

ssize_t read_hbc_epoch_update_string(
    const char *filename,
    std::vector<Proof>& generatorProof,
    std::vector<Twistpoint>& newFreshPseudonyms,
    std::vector<CurveBipoint>& newServerTallies,
    std::vector<std::vector<TwistBipoint>>& newVoteMatrix,
    std::vector<Twistpoint>& newUserTallies,
    Twistpoint& nextGenerator,
    bool& doUserTallies);

/* OTHER SERVER-RELEVANT HANDLERS */

// Used to tell orchestrator when the system is ready to do an epoch change
class EpochReadyHandler : public CivetHandler
{
    public:
        EpochReadyHandler(
            struct synchronization_tool *exitSync,
            struct synchronization_tool *readySync,
            std::mutex& updateMtx,
            size_t numServers);

        bool handleGet(
            CivetServer *server,
            struct mg_connection *conn);

    private:
        struct synchronization_tool *exitSync, *readySync;
        std::mutex& updateMtx;
        const size_t numServers;

};

// Used to indicate which epoch the system is currently in
class EpochNumHandler : public CivetHandler
{
    public:
        EpochNumHandler(
            std::atomic<size_t>& epochNum);
        
        bool handleGet(
            CivetServer *server,
            struct mg_connection *conn);

    private:
        std::atomic<size_t>& epochNum;
};

// Used to take and release update locks
class UpdateLockWebSocketHandler : public CivetWebSocketHandler
{
    public:
        UpdateLockWebSocketHandler(
            std::mutex& updateMtx,
            std::unique_lock<std::mutex> **lockHolder,
            bool isLocking);

        ~UpdateLockWebSocketHandler();

        bool handleConnection(
            CivetServer *server,
            const struct mg_connection *conn);

        void handleReadyState(
            CivetServer *server,
            struct mg_connection *conn);

        bool handleData(
            CivetServer *server,
            struct mg_connection *conn,
            int bits,
            char *data,
            size_t data_len);

        void handleClose(
            CivetServer *server,
            const struct mg_connection *conn);

    private:
        std::mutex& updateMtx;
        std::unique_lock<std::mutex> **lockHolder;
        const bool isLocking;
};

/* SERVER-SPECIFIC HANDLER */

class PrsonaServerWebSocketHandler : public CivetWebSocketHandler  {
    public:
        // CONSTRUCTORS
        PrsonaServerWebSocketHandler(
            std::default_random_engine& rng,
            PrsonaServer *prsonaServer,
            const std::vector<std::string>& serverIPs,
            const std::vector<int>& serverPorts,
            const std::string& selfIP,
            int selfPort,
            std::mutex& updateMtx,
            std::atomic<size_t>& epochNum,
            std::mutex& buildUpOutputMtx,
            const std::string& buildUpOutputFilename,
            std::mutex& breakDownOutputMtx,
            const std::string& breakDownOutputFilename,
            std::mutex& updateOutputMtx,
            const std::string& updateOutputFilename,
            std::mutex& voteOutputMtx,
            const std::string& voteOutputFilename,
            std::mutex& usageMtx,
            const std::string& usageFilename);

        // REQUIRED BY INHERITED CLASS
        bool handleConnection(
            CivetServer *server,
            const struct mg_connection *conn);

        void handleReadyState(
            CivetServer *server,
            struct mg_connection *conn);

        bool handleData(
            CivetServer *server,
            struct mg_connection *conn,
            int bits,
            char *data,
            size_t data_len);

        void handleClose(
            CivetServer *server,
            const struct mg_connection *conn);

    private:
        std::default_random_engine& rng;
        
        PrsonaServer *prsonaServer;
        
        const std::vector<std::string> serverIPs;
        const std::vector<int> serverPorts;
        
        const std::string selfIP;
        const int selfPort;
        
        std::mutex& updateMtx;
        
        std::atomic<size_t>& epochNum;
        
        std::mutex& buildUpOutputMtx;
        const std::string buildUpOutputFilename;
        std::mutex& breakDownOutputMtx;
        const std::string breakDownOutputFilename;
        std::mutex& updateOutputMtx;
        const std::string updateOutputFilename;
        std::mutex& voteOutputMtx;
        const std::string voteOutputFilename;
        std::mutex& usageMtx;
        const std::string usageFilename;

        struct synchronization_tool updateSynch, distributeSynch;

        // RESPONSE ROUTER FUNCTION
        void generate_response(
            CivetServer *server,
            struct mg_connection *conn,
            const char *filename);

        // BASIC PUBLIC SYSTEM INFO GETTERS
        void get_bgn_public_key(
            struct mg_connection *conn
        ) const;

        void get_num_clients(
            struct mg_connection *conn
        ) const;

        void get_num_servers(
            struct mg_connection *conn
        ) const;

        void get_generator(
            struct mg_connection *conn,
            bool fresh
        );

        // ENCRYPTED DATA GETTERS
        void get_current_votes_by(
            struct mg_connection *conn,
            const char *filename
        ) const;

        void get_all_current_votes(
            struct mg_connection *conn
        ) const;

        void get_current_user_encrypted_tally(
            struct mg_connection *conn,
            const char *filename
        ) const;

        void get_current_server_encrypted_tally(
            struct mg_connection *conn,
            const char *filename
        ) const;

        void get_current_pseudonyms(
            struct mg_connection *conn
        ) const;

        // PROOF COMMITMENT GETTERS
        void get_vote_row_commitment(
            struct mg_connection *conn,
            const char *filename
        ) const;

        void get_vote_matrix_commitment(
            struct mg_connection *conn
        ) const;

        void get_user_tally_commitment(
            struct mg_connection *conn,
            const char *filename
        ) const;

        void get_server_tally_commitment(
            struct mg_connection *conn,
            const char *filename
        ) const;

        void get_pseudonyms_commitment(
            struct mg_connection *conn
        ) const;

        // CLIENT INTERACTIONS
        void add_new_client(
            struct mg_connection *conn,
            const char *filename
        );

        void receive_vote(
            CivetServer *civetServer,
            struct mg_connection *conn,
            const char *filename
        );

        // DISTRIBUTION HELPERS FOR CLIENT INTERACTIONS
        void distribute_new_user_updates(
            const std::vector<Proof>& proofOfValidAddition,
            const std::vector<CurveBipoint>& previousVoteTallies,
            const std::vector<Twistpoint>& currentPseudonyms,
            const std::vector<EGCiphertext>& currentUserEncryptedTallies,
            const std::vector<std::vector<TwistBipoint>>& voteMatrix
        ) const;

        void distribute_new_vote(
            std::vector<Proof> pi,
            std::vector<TwistBipoint> newVotes,
            Twistpoint shortTermPublicKey,
            std::vector<size_t>& bandwidthData
        ) const;

        void import_new_user_update(
            struct mg_connection *conn,
            const char *filename
        );

        // CONSTRUCTOR HELPERS
        void get_bgn_details(
            struct mg_connection *conn
        ) const;

        void add_seed_to_generator(
            struct mg_connection *conn,
            const char *filename,
            bool fresh
        ) const;

        void set_generator(
            const char *filename,
            bool fresh
        );

        // EPOCH ROUNDS
        void build_up_midway_pseudonyms(
            CivetServer *civetServer,
            struct mg_connection *conn,
            const char *filename
        );

        void hbc_build_up_midway_pseudonyms(
            CivetServer *civetServer,
            struct mg_connection *conn,
            const char *filename
        );

        void break_down_midway_pseudonyms(
            CivetServer *civetServer,
            struct mg_connection *conn,
            const char *filename
        );

        void hbc_break_down_midway_pseudonyms(
            CivetServer *civetServer,
            struct mg_connection *conn,
            const char *filename
        );

        void accept_epoch_updates(
            CivetServer *civetServer,
            struct mg_connection *conn,
            const char *filename
        );

        void hbc_accept_epoch_updates(
            CivetServer *civetServer,
            struct mg_connection *conn,
            const char *filename
        );

        // DISTRIBUTED BGN
        void get_partial_decryption(
            struct mg_connection *conn
        ) const;

        void receive_tallied_scores(
            struct mg_connection *conn,
            const char *filename
        );
}; 

#endif