#include <iostream>
#include <functional>
#include "../App/appconfig.hpp"

// The next line suppresses a deprecation warning within boost
#define BOOST_BIND_GLOBAL_PLACEHOLDERS
#include "boost/property_tree/ptree.hpp"
#include "boost/property_tree/json_parser.hpp"
#include <boost/asio.hpp>
#include <thread>
#include "gcm.h"
#include "sgx_tcrypto.h"
#include "clients.hpp"

#define CEILDIV(x,y) (((x)+(y)-1)/(y))

// Split a hostport string like "127.0.0.1:12000" at the rightmost colon
// into a host part "127.0.0.1" and a port part "12000".
static bool split_host_port(std::string &host, std::string &port,
    const std::string &hostport)
{
    size_t colon = hostport.find_last_of(':');
    if (colon == std::string::npos) {
        std::cerr << "Cannot parse \"" << hostport << "\" as host:port\n";
        return false;
    }
    host = hostport.substr(0, colon);
    port = hostport.substr(colon+1);
    return true;
}

// Convert a single hex character into its value from 0 to 15. Return
// true on success, false if it wasn't a hex character.
static inline bool hextoval(unsigned char &val, char hex)
{
    if (hex >= '0' && hex <= '9') {
        val = ((unsigned char)hex)-'0';
    } else if (hex >= 'a' && hex <= 'f') {
        val = ((unsigned char)hex)-'a'+10;
    } else if (hex >= 'A' && hex <= 'F') {
        val = ((unsigned char)hex)-'A'+10;
    } else {
        return false;
    }
    return true;
}

// Convert a 2*len hex character string into a len-byte buffer. Return
// true on success, false on failure.
static bool hextobuf(unsigned char *buf, const char *str, size_t len)
{
    if (strlen(str) != 2*len) {
        std::cerr << "Hex string was not the expected size\n";
        return false;
    }
    for (size_t i=0;i<len;++i) {
        unsigned char hi, lo;
        if (!hextoval(hi, str[2*i]) || !hextoval(lo, str[2*i+1])) {
            std::cerr << "Cannot parse string as hex\n";
            return false;
        }
        buf[i] = (unsigned char)((hi << 4) + lo);
    }
    return true;
}

void displayMessage(unsigned char *msg, uint16_t msg_size)
{
    clientid_t sid, rid;
    unsigned char *ptr = msg;
    sid = *((clientid_t*) ptr);
    ptr+=sizeof(sid);
    rid = *((clientid_t*) ptr);
    ptr+=sizeof(sid);
    printf("Sender ID: %d, Receiver ID: %d, Token: N/A\n", sid, rid );
    printf("Message: ");
    for(int j = 0; j<msg_size - sizeof(sid)*2; j++) {
        printf("%x", (*ptr));
        ptr++;
    }
    printf("\n");
}

void displayPtMessageBundle(unsigned char *bundle, uint16_t priv_out,
    uint16_t msg_size)
{
    unsigned char *ptr = bundle;

    for(int i=0; i<priv_out; i++) {
        displayMessage(ptr, msg_size);
        printf("\n");
        ptr+=msg_size;
    }
    printf("\n");
}

void displayEncMessageBundle(unsigned char *bundle, uint16_t priv_out,
    uint16_t msg_size)
{
    unsigned char *ptr = bundle;
    uint64_t header = *((uint64_t*) ptr);
    ptr+=sizeof(uint64_t);

    printf("IV: ");
    for(int i=0; i<SGX_AESGCM_IV_SIZE; i++) {
        printf("%x", ptr[i]);
    }
    printf("\n");
    ptr+= SGX_AESGCM_IV_SIZE;

    for(int i=0; i<priv_out; i++) {
        displayMessage(ptr, msg_size);
        ptr+=msg_size;
    }

    printf("MAC: ");
    for(int i=0; i<SGX_AESGCM_MAC_SIZE; i++) {
        printf("%x", ptr[i]);
    }
    printf("\n");
}


static inline uint32_t encMsgBundleSize(uint16_t priv_out, uint16_t msg_size)
{
    return(SGX_AESGCM_IV_SIZE + (priv_out * msg_size) + SGX_AESGCM_MAC_SIZE);
}

static inline uint32_t ptMsgBundleSize(uint16_t priv_out, uint16_t msg_size)
{
    return((priv_out * msg_size));
}

bool config_parse(Config &config, const std::string configstr,
    std::vector<NodeConfig> &ingestion_nodes,
    std::vector<NodeConfig> &storage_nodes,
    std::vector<uint16_t> &storage_map)
{
    bool found_params = false;
    bool ret = true;

    std::istringstream configstream(configstr);
    boost::property_tree::ptree conftree;

    read_json(configstream, conftree);
    uint16_t node_num = 0;

    for (auto & entry : conftree) {
        if (!entry.first.compare("params")) {
            for (auto & pentry : entry.second) {
                if (!pentry.first.compare("msg_size")) {
                    config.msg_size = pentry.second.get_value<uint16_t>();
                } else if (!pentry.first.compare("user_count")) {
                    config.user_count = pentry.second.get_value<uint32_t>();
                } else if (!pentry.first.compare("priv_out")) {
                    config.m_priv_out = pentry.second.get_value<uint8_t>();
                } else if (!pentry.first.compare("priv_in")) {
                    config.m_priv_in = pentry.second.get_value<uint8_t>();
                } else if (!pentry.first.compare("pub_out")) {
                    config.m_pub_out = pentry.second.get_value<uint8_t>();
                } else if (!pentry.first.compare("pub_in")) {
                    config.m_pub_in = pentry.second.get_value<uint8_t>();
                // A hardcoded shared secret to derive various
                // keys for client -> server communications and tokens
                } else if (!pentry.first.compare("master_secret")) {
                    std::string hex_key = pentry.second.data();
                    memcpy(config.master_secret, hex_key.c_str(), SGX_AESGCM_KEY_SIZE);

                } else {
                    std::cerr << "Unknown field in params: " <<
                        pentry.first << "\n";
                    ret = false;
                }
            }
            found_params = true;
        } else if (!entry.first.compare("nodes")) {

            for (auto & node : entry.second) {
                NodeConfig nc;
                // All nodes need to be assigned their role in manifest.yaml
                nc.roles = 0;
                for (auto & nentry : node.second) {
                    if (!nentry.first.compare("name")) {
                        nc.name = nentry.second.get_value<std::string>();
                    } else if (!nentry.first.compare("pubkey")) {
                        ret &= hextobuf((unsigned char *)&nc.pubkey,
                            nentry.second.get_value<std::string>().c_str(),
                            sizeof(nc.pubkey));
                    } else if (!nentry.first.compare("weight")) {
                        nc.weight = nentry.second.get_value<std::uint8_t>();
                    } else if (!nentry.first.compare("listen")) {
                        ret &= split_host_port(nc.listenhost, nc.listenport,
                            nentry.second.get_value<std::string>());
                    } else if (!nentry.first.compare("clisten")) {
                        ret &= split_host_port(nc.clistenhost, nc.clistenport,
                            nentry.second.get_value<std::string>());
                    } else if (!nentry.first.compare("roles")) {
                        nc.roles = nentry.second.get_value<std::uint8_t>();
                    } else {
                        std::cerr << "Unknown field in host config: " <<
                            nentry.first << "\n";
                        ret = false;
                    }
                }
                if(nc.roles & ROLE_INGESTION) {
                    ingestion_nodes.push_back(std::move(nc));
                }
                if(nc.roles & ROLE_STORAGE) {
                    storage_nodes.push_back(std::move(nc));
                    storage_map.push_back(node_num);
                }
                node_num++;
            }
        } else {
            std::cerr << "Unknown key in config: " <<
                entry.first << "\n";
            ret = false;
        }
    }

    if (!found_params) {
        std::cerr << "Could not find params in config\n";
        ret = false;
    }

    return ret;
}

static void usage(const char *argv0)
{
    fprintf(stderr, "%s [-t nthreads] < config.json\n",
        argv0);
    exit(1);
}

/*

    Generate EMK (Encryption master Secret Key) and TMK (Token master Secret Key)

*/
int generateMasterKeys(sgx_aes_gcm_128bit_key_t master_secret,
   aes_key &EMK, aes_key &TMK )
{
    unsigned char zeroes[SGX_AESGCM_KEY_SIZE];
    unsigned char iv[SGX_AESGCM_IV_SIZE];
    unsigned char mac[SGX_AESGCM_MAC_SIZE];
    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memset(zeroes, 0, SGX_AESGCM_KEY_SIZE);
    memcpy(iv, "Encryption", sizeof("Encryption"));

    if (sizeof(zeroes) != gcm_encrypt(zeroes, SGX_AESGCM_KEY_SIZE, NULL, 0,
            master_secret, iv, SGX_AESGCM_IV_SIZE, EMK, mac)) {
        printf("Client: generateMasterKeys FAIL\n");
        return -1;
    }

    printf("\n\nEncryption Master Key: ");
    for(int i=0;i<SGX_AESGCM_KEY_SIZE;i++) {
        printf("%x", EMK[i]);
    }
    printf("\n");

    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memcpy(iv, "Token", sizeof("Token"));
    if (sizeof(zeroes) != gcm_encrypt(zeroes, SGX_AESGCM_KEY_SIZE, NULL, 0,
            master_secret, iv, SGX_AESGCM_IV_SIZE, TMK, mac)) {
        printf("generateMasterKeys failed\n");
        return -1;
    }

    printf("Token Master Key: ");
    for(int i=0;i<SGX_AESGCM_KEY_SIZE;i++) {
        printf("%x", TMK[i]);
    }
    printf("\n\n");

    return 1;
}

/*
    Takes the client_number, the master aes_key for generating client keys
    for encrypted communication with ingestion (client_ing_key) and
    storage servers (client_stg_key)
*/
int generateClientKeys(clientid_t client_number, aes_key &EMK,
    aes_key &client_ing_key, aes_key &client_stg_key)
{
    unsigned char zeroes[SGX_AESGCM_KEY_SIZE];
    unsigned char iv[SGX_AESGCM_IV_SIZE];
    unsigned char tag[SGX_AESGCM_MAC_SIZE];
    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memset(zeroes, 0, SGX_AESGCM_KEY_SIZE);
    memset(tag, 0, SGX_AESGCM_KEY_SIZE);
    memcpy(iv, &client_number, sizeof(client_number));

    /*
    printf("Client Key: (before Gen) ");
    for(int i=0;i<SGX_AESGCM_KEY_SIZE;i++) {
        printf("%x", client_ing_key[i]);
    }
    printf("\n");
    */

    if (sizeof(zeroes) != gcm_encrypt(zeroes, SGX_AESGCM_KEY_SIZE, NULL, 0, EMK,
            iv, SGX_AESGCM_IV_SIZE, client_ing_key, tag)) {
        printf("generateClientKeys failed\n");
        return -1;
    }

    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memcpy(iv, &client_number, sizeof(client_number));
    memcpy(iv +sizeof(client_number), "STG", sizeof("STG"));
    if (sizeof(zeroes) != gcm_encrypt(zeroes, SGX_AESGCM_KEY_SIZE, NULL, 0, EMK,
            iv, SGX_AESGCM_IV_SIZE, client_stg_key, tag)) {
        printf("generateClientKeys failed\n");
        return -1;
    }

    /*
    if(client_number % 10 == 0) {
        printf("Client %d Storage Key: (after Gen) ", client_number);
        for(int i=0;i<SGX_AESGCM_KEY_SIZE;i++) {
            printf("%x", client_stg_key[i]);
        }
        printf("\n");
    }
    */

    return 1;
}


void Client::initClient(clientid_t cid, aes_key ikey, aes_key skey,
    uint16_t num_storage_nodes, std::vector<uint16_t> &storage_map)
{
    sim_id = cid;
    uint16_t stg_no = cid % num_storage_nodes;
    uint16_t stg_id = storage_map[stg_no];
    id = stg_id << DEST_UID_BITS;
    id += (cid/num_storage_nodes);

    memcpy(ing_key, ikey, SGX_AESGCM_KEY_SIZE);
    memcpy(stg_key, skey, SGX_AESGCM_KEY_SIZE);
}

void Client::initializeSocket(boost::asio::io_context &ioc,
    NodeConfig &ing_server)
{

    boost::system::error_code err;
    boost::asio::ip::tcp::resolver resolver(ioc);

    while(1) {
#ifdef VERBOSE_NET
        std::cerr << "Connecting to " << ing_server.name << "...\n";
        std::cout << ing_server.clistenhost << ":" << ing_server.clistenport;
#endif
        // ingestion_sock needs io_context
        ingestion_sock = new boost::asio::ip::tcp::socket(ioc);
        boost::asio::connect(*ingestion_sock,
            resolver.resolve(ing_server.clistenhost,
                ing_server.clistenport), err);
        if (!err) break;
        std::cerr << "Connection to " << ing_server.name <<
            " refused, will , epoch_noretry.\n";
        sleep(1);
    }
}

/*

    Populates the buffer pt_msgbundle with a valid message pt_msgbundle.
    Assumes that it is supplied with a pt_msgbundle buffer of the correct length

    Correct length for pt_msgbundle  =  8 + (priv_out)*(msg_size) + 16 bytes

*/
void Client::generateMessageBundle(uint8_t priv_out, uint32_t msg_size,
    unsigned char *pt_msgbundle)
{
    unsigned char *ptr = pt_msgbundle;

    // Setup message pt_msgbundle
    for(uint32_t i = 0; i < priv_out; i++) {
        memcpy(ptr, &id, sizeof(id));
        ptr+=(sizeof(id));

        memcpy(ptr, &id, sizeof(id));
        ptr+=(sizeof(id));

        uint32_t remaining_message_size = msg_size - (sizeof(id)*2);
        memset(ptr, 0, remaining_message_size);
        ptr+=(remaining_message_size);
    }
}


bool Client::encryptMessageBundle(uint32_t enc_bundle_size, unsigned char *pt_msgbundle,
    unsigned char *enc_msgbundle)
{
    // Encrypt the pt_msgbundle
    unsigned char *pt_msgbundle_start = pt_msgbundle;
    unsigned char *enc_msgbundle_start = enc_msgbundle + SGX_AESGCM_IV_SIZE;
    unsigned char *enc_tag = enc_msgbundle + enc_bundle_size - SGX_AESGCM_MAC_SIZE;
    size_t bytes_to_encrypt = enc_bundle_size - SGX_AESGCM_MAC_SIZE - SGX_AESGCM_IV_SIZE;
    if (bytes_to_encrypt != gcm_encrypt(pt_msgbundle_start, bytes_to_encrypt,
        NULL, 0, ing_key, ing_iv, SGX_AESGCM_IV_SIZE, enc_msgbundle_start, enc_tag)) {
            printf("Client: encryptMessageBundle FAIL\n");
            return 0;
    }

    // Copy IV into the bundle
    memcpy(enc_msgbundle, ing_iv, SGX_AESGCM_IV_SIZE);

    // Update IV
    uint64_t *iv_ctr = (uint64_t*) ing_iv;
    (*iv_ctr)+=1;
    return 1;
}

/*

      Assumes pt_msgbundle is a buffer of size messageBundleSize(priv_out, msg_size)
*/

void Client::sendMessageBundle(uint16_t priv_out, uint16_t msg_size,
    unsigned char *pt_msgbundle, unsigned char *enc_msgbundle)
{
    uint32_t enc_bundle_size = encMsgBundleSize(priv_out, msg_size);

    generateMessageBundle(priv_out, msg_size, pt_msgbundle);

    encryptMessageBundle(enc_bundle_size, pt_msgbundle, enc_msgbundle);

#ifdef VERBOSE_CLIENT
    displayPtMessageBundle(pt_msgbundle, priv_out, msg_size);
#endif

    boost::asio::write(*ingestion_sock,
        boost::asio::buffer(enc_msgbundle, enc_bundle_size));
}

int Client::sendAuthMessage(unsigned long epoch_no)
{
    uint32_t auth_size = sizeof(clientid_t) + sizeof(unsigned long) + SGX_AESGCM_KEY_SIZE;
    unsigned char *auth_message = (unsigned char*) malloc(auth_size);
    unsigned char *am_ptr = auth_message;

    memcpy(am_ptr, &sim_id, sizeof(sim_id));
    am_ptr+=sizeof(sim_id);

    memcpy(am_ptr, &epoch_no, sizeof(unsigned long));
    am_ptr+=sizeof(unsigned long);

    unsigned char zeroes[SGX_AESGCM_KEY_SIZE] = {0};
    unsigned char tag[SGX_AESGCM_MAC_SIZE] = {0};
    unsigned char epoch_iv[SGX_AESGCM_IV_SIZE] = {0};
    memcpy(epoch_iv, &epoch_no, sizeof(epoch_no));

    if (sizeof(zeroes) != gcm_encrypt(zeroes, SGX_AESGCM_KEY_SIZE, NULL, 0, ing_key,
            epoch_iv, SGX_AESGCM_IV_SIZE, am_ptr, tag)) {
        printf("generateClientKeys failed\n");
        return -1;
    }

#ifdef VERBOSE_CLIENT
    printf("Client %d auth_message: \n", id);
    for(int i=0; i<auth_size; i++) {
        printf("%x", auth_message[i]);
    }
    printf("\n");
#endif

    boost::asio::write(*ingestion_sock,
        boost::asio::buffer(auth_message, auth_size));

    return 1;
}

void generateClients(boost::asio::io_context &io_context,
    uint32_t cstart, uint32_t cstop, Client* &clients,
    aes_key &EMK, Config &config, std::vector<NodeConfig> &ingestion_nodes,
    std::vector<NodeConfig> &storage_nodes, std::vector<uint16_t> &storage_map,
    uint32_t num_clients_total, uint32_t clients_per_ing,
    uint32_t ing_with_additional)
{
    aes_key client_ing_key;
    aes_key client_stg_key;
    uint16_t num_stg_nodes = storage_nodes.size();
    uint16_t *stg_map = new uint16_t[num_stg_nodes];

    for(uint32_t i=cstart; i<cstop; i++) {
        uint16_t ing_node_this_client = i/clients_per_ing;
        if(ing_node_this_client > ing_with_additional && ing_with_additional!=0) {
            uint16_t leftover = num_clients_total - (ing_with_additional * clients_per_ing);
            ing_node_this_client = ing_with_additional + (leftover / (clients_per_ing-1));
        }

        int ret = generateClientKeys(i, EMK, client_ing_key, client_stg_key);

        clients[i].initClient(i, client_ing_key, client_stg_key, num_stg_nodes, storage_map);

        clients[i].initializeSocket(io_context, ingestion_nodes[ing_node_this_client]);

        /*
        // Test that the keys generated match those generated within
        // enclave config
        unsigned char *ckey;
        ckey = clients[i].getKey();
        printf("Client %d, id = %d, key: ", i, clients[i].getid());
        for(int j=0;j<SGX_AESGCM_KEY_SIZE;j++) {
            printf("%x", ckey[j]);
        }
        printf("\n\n");
        */
    }

    struct timespec ep;
    clock_gettime(CLOCK_REALTIME_COARSE, &ep);
    unsigned long ep_time = ep.tv_sec * 1000000 + ep.tv_nsec/1000;
    unsigned long epoch_no = CEILDIV(ep_time, EPOCH_INTERVAL);

    for(uint32_t i=cstart; i<cstop; i++) {
        clients[i].sendAuthMessage(epoch_no);
    }
}


void sendMessageBundles(uint32_t cstart, uint32_t cstop, Client* &clients,
    Config &config)
{
    uint16_t priv_out = config.m_priv_out;
    uint16_t msg_size = config.msg_size;
    uint32_t pt_bundle_size = ptMsgBundleSize(priv_out, msg_size);
    uint32_t enc_bundle_size = encMsgBundleSize(priv_out, msg_size);
    unsigned char *pt_msgbundle = (unsigned char*) malloc (pt_bundle_size);
    unsigned char *enc_msgbundle = (unsigned char*) malloc (enc_bundle_size);

    for(uint32_t i=cstart; i<cstop; i++) {
        clients[i].sendMessageBundle(priv_out, msg_size, pt_msgbundle, enc_msgbundle);
    }

    free(pt_msgbundle);
    free(enc_msgbundle);
}
/*
    Spin config.user_client actual clients. Each client:
    1) Retrieve messages and tokens from their storage server
    2) Send all their messages to the ingestion server
    3) Wait for a predetermined EPOCH_DURATION time
    4) Repeat from 1)
*/


int main(int argc, char **argv)
{
    // Unbuffer stdout
    setbuf(stdout, NULL);

    uint16_t nthreads = 1;
    const char *progname = argv[0];
    std::vector<NodeConfig> ingestion_nodes, storage_nodes;
    std::vector<uint16_t> storage_map;
    ++argv;

    // Parse options
    while (*argv && (*argv)[0] == '-') {
        if (!strcmp(*argv, "-t")) {
            if (argv[1] == NULL) {
                usage(progname);
            }
            nthreads = uint16_t(atoi(argv[1]));
            argv += 2;
        } else {
            usage(progname);
        }
    }

    printf("nthreads = %d\n", nthreads);
    // Read the config.json from the first line of stdin.  We have to do
    // this before outputting anything to avoid potential deadlock with
    // the launch program.
    std::string configstr;
    std::getline(std::cin, configstr);

    Config config;
    aes_key EMK, TMK;
    boost::asio::io_context io_context;
    boost::asio::ip::tcp::resolver resolver(io_context);

    if (!config_parse(config, configstr, ingestion_nodes,
        storage_nodes, storage_map)) {
       exit(1);
    }

    Client *clients = new Client[config.user_count];
    printf("Number of ingestion_nodes = %ld, Number of storage_node = %ld\n",
        ingestion_nodes.size(), storage_nodes.size());

    generateMasterKeys(config.master_secret, EMK, TMK);

    uint32_t num_clients_total = config.user_count;
    uint16_t num_ing_nodes = ingestion_nodes.size();
    uint32_t clients_per_ing = CEILDIV(num_clients_total, num_ing_nodes);
    uint32_t clients_per_thread = CEILDIV(num_clients_total, nthreads);
    uint16_t ing_with_additional = num_clients_total % num_ing_nodes;

    std::thread threads[nthreads];

    // Generate all the clients for the experiment
    for(int i=0; i<nthreads; i++) {
        uint32_t cstart, cstop;
        cstart = i * clients_per_thread;
        cstop = (i==nthreads-1)? num_clients_total: (i+1) * clients_per_thread;
        printf("Thread %d, cstart = %d, cstop = %d\n", i, cstart, cstop);
        threads[i] = std::thread(generateClients, std::ref(io_context),
            cstart, cstop, std::ref(clients), std::ref(EMK), std::ref(config),
            std::ref(ingestion_nodes), std::ref(storage_nodes),
            std::ref(storage_map), num_clients_total,
            clients_per_ing, ing_with_additional);
    }

    for(int i=0; i<nthreads; i++) {
        threads[i].join();
    }

    // Multithreaded client message bundle generation and send
    uint32_t epoch = 1;
    while(epoch <= 3) {
        struct timespec tp;
        clock_gettime(CLOCK_REALTIME_COARSE, &tp);
        unsigned long start = tp.tv_sec * 1000000 + tp.tv_nsec/1000;

        for(int i=0; i<nthreads; i++) {
            uint32_t cstart, cstop;
            cstart = i * clients_per_thread;
            cstop = (i==nthreads-1)? num_clients_total: (i+1) * clients_per_thread;
            threads[i] = std::thread(sendMessageBundles, cstart, cstop,
                std::ref(clients), std::ref(config));
        }
        for(int i=0; i<nthreads; i++) {
            threads[i].join();
        }

        clock_gettime(CLOCK_REALTIME_COARSE, &tp);
        unsigned long end = tp.tv_sec * 1000000 + tp.tv_nsec/1000;
        unsigned long time_diff = end - start;

        // Sleep for the rest of the epoch interval
        printf("Done with submissions for Epoch %d\n", epoch);
        if (time_diff < EPOCH_INTERVAL) {
            unsigned long time_to_sleep_in_us = (useconds_t) EPOCH_INTERVAL - (useconds_t) time_diff;
            //printf("tts_us = %ld\n", time_to_sleep_in_us);
            usleep(time_to_sleep_in_us);
        }
        epoch++;
    }

    delete [] clients;
}