#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);
    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;
    uint64_t header = *((uint64_t*) ptr);
    ptr+=sizeof(uint64_t);

    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");
}


#define HEADER_SIZE 8
static inline uint32_t encMsgBundleSize(uint16_t priv_out, uint16_t msg_size) {
    return(HEADER_SIZE + 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(HEADER_SIZE + (priv_out * msg_size));
}

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

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

    read_json(configstream, conftree);

    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>();
                // Currently hardcoding an AES key for client -> server communication
                } 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));
                }
            }
        } 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 encryption keys,
    and the client aes_key to be generated.
*/
int generateClientEncryptionKey(clientid_t client_number, aes_key &EMK, aes_key &client_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_key[i]);
    }
    printf("\n");
    */

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

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

    return 1;
}


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 retry.\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;
    uint64_t header = (id << 8) + CLIENT_MESSAGE_BUNDLE;

    // Setup header
    memcpy(ptr, (uint8_t*) &header, sizeof(header));
    ptr+=sizeof(header);

    // 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(i));
        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)
{
    // Copy the header
    memcpy(enc_msgbundle, pt_msgbundle, HEADER_SIZE);

    // Encrypt the rest of the pt_msgbundle
    unsigned char *pt_msgbundle_start = pt_msgbundle + HEADER_SIZE;
    unsigned char *enc_msgbundle_start = enc_msgbundle + HEADER_SIZE + 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 - HEADER_SIZE - SGX_AESGCM_IV_SIZE;
    if (bytes_to_encrypt != gcm_encrypt(pt_msgbundle_start, bytes_to_encrypt,
        NULL, 0, key, iv, SGX_AESGCM_IV_SIZE, enc_msgbundle_start, enc_tag)) {
            printf("Client: encryptMessageBundle FAIL\n");
            return 0;
    }

    // Copy the IV into the bundle
    unsigned char *enc_msgbundle_iv = enc_msgbundle + HEADER_SIZE;
    memcpy(enc_msgbundle_iv, iv, SGX_AESGCM_IV_SIZE);

    // Update IV
    uint64_t *iv_ctr = (uint64_t*) 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);

    displayPtMessageBundle(pt_msgbundle, priv_out, msg_size);

    //displayEncMessageBundle(enc_msgbundle, priv_out, msg_size);

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

int Client::sendAuthMessage()
{
    uint32_t auth_size = sizeof(uint64_t) + SGX_AESGCM_KEY_SIZE;
    unsigned char *auth_string = (unsigned char*) malloc(auth_size);
    unsigned char *as_ptr = auth_string;
    uint64_t header = (id << 8) + CLIENT_AUTHENTICATE;
    memcpy(as_ptr, &header, sizeof(header));
    as_ptr+=sizeof(header);

    unsigned char zeroes[SGX_AESGCM_KEY_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);

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

    // Update IV
    uint64_t *iv_ctr = (uint64_t*) iv;
    (*iv_ctr)+=1;

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

    boost::asio::write(*ingestion_sock,
        boost::asio::buffer(auth_string, 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,
    uint32_t num_clients_total, uint32_t clients_per_ing,
    uint32_t ing_with_additional)
{
    aes_key client_key;

    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 = generateClientEncryptionKey(i, EMK, client_key);
        clients[i].initClient(i, client_key);

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

        /*
        // 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");
        */
    }
}


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;
    ++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)) {
       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), 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 = 0;
    while(epoch < 1) {
        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();
        }
        epoch++;
    }
    delete [] clients;
}