#include "Enclave_t.h"
#include "comms.hpp"
#include "config.hpp"
#include "utils.hpp"
#include "route.hpp"
#include "ingest.hpp"

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

unsigned long ingestion_epoch;
unsigned long storage_epoch;

Config g_teems_config;

static int generateMasterKeys(sgx_aes_gcm_128bit_key_t master_secret,
    sgx_aes_gcm_128bit_key_t &ESK, sgx_aes_gcm_128bit_key_t &TSK)
{
    unsigned char zeroes[SGX_AESGCM_KEY_SIZE];
    unsigned char iv[SGX_AESGCM_IV_SIZE];
    sgx_aes_gcm_128bit_tag_t mac;
    memset(zeroes, 0, SGX_AESGCM_KEY_SIZE);
    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memcpy(iv, "Encryption", sizeof("Encryption"));
    sgx_status_t ret = SGX_SUCCESS;

    ret = sgx_rijndael128GCM_encrypt((const sgx_aes_gcm_128bit_key_t *)
        master_secret, zeroes, SGX_AESGCM_KEY_SIZE,
        (uint8_t*) ESK, iv, SGX_AESGCM_IV_SIZE, NULL, 0, &mac);
    if(ret!=SGX_SUCCESS) {
        return -1;
    }

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

    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memcpy(iv, "Token", sizeof("Token"));
    ret = sgx_rijndael128GCM_encrypt((const sgx_aes_gcm_128bit_key_t *)
        master_secret, zeroes, SGX_AESGCM_KEY_SIZE,
        (uint8_t*) TSK, iv, SGX_AESGCM_IV_SIZE, NULL, 0, &mac);
    if(ret!=SGX_SUCCESS) {
        return -1;
    }

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

    return 1;
}

bool ecall_config_load(threadid_t nthreads,
    EnclaveAPIParams *apiparams,
    EnclaveAPINodeConfig *apinodeconfigs,
    nodenum_t num_nodes, nodenum_t my_node_num)
{
#ifdef TRACK_HEAP_USAGE
    printf("ecall_config_load begin heap %u\n", g_peak_heap_used);
#endif
    g_teems_config.nthreads = nthreads;
    g_teems_config.num_nodes = num_nodes;
    g_teems_config.num_ingestion_nodes = 0;
    g_teems_config.num_routing_nodes = 0;
    g_teems_config.num_storage_nodes = 0;
    g_teems_config.my_node_num = my_node_num;
    g_teems_config.user_count = apiparams->user_count;
    g_teems_config.msg_size = apiparams->msg_size;
    g_teems_config.m_token_out = apiparams->m_token_out;
    g_teems_config.m_token_in = apiparams->m_token_in;
    g_teems_config.m_id_out = apiparams->m_id_out;
    g_teems_config.m_id_in = apiparams->m_id_in;
    memcpy(g_teems_config.master_secret, apiparams->master_secret,
        SGX_AESGCM_KEY_SIZE);
    g_teems_config.token_channel = apiparams->token_channel;
    // Temporary vectors to store node numbers for nodes of different
    // types, where the node numbers are smaller than our own node
    // number
    std::vector<nodenum_t> ing_smaller, rte_smaller, str_smaller;
    uint16_t cumul_weight = 0;
    g_teems_config.weights.clear();
    g_teems_config.ingestion_nodes.clear();
    g_teems_config.routing_nodes.clear();
    g_teems_config.storage_nodes.clear();
    g_teems_config.storage_map.clear();
    for (nodenum_t i=0; i<num_nodes; ++i) {
        NodeWeight nw;
        nw.startweight = cumul_weight;
        // Weights only matter for routing nodes
        nw.weight = 0;
        if (apinodeconfigs[i].roles & ROLE_INGESTION) {
            g_teems_config.num_ingestion_nodes += 1;
            if (i < my_node_num) {
                ing_smaller.push_back(i);
            } else {
                g_teems_config.ingestion_nodes.push_back(i);
            }
        }
        if (apinodeconfigs[i].roles & ROLE_ROUTING) {
            // Only use weights in token channel routing
            if (g_teems_config.token_channel) {
                nw.weight = apinodeconfigs[i].weight;
            } else  {
                nw.weight = 1;
            }
            g_teems_config.num_routing_nodes += 1;
            if (i < my_node_num) {
                rte_smaller.push_back(i);
            } else {
                g_teems_config.routing_nodes.push_back(i);
            }
        }
        if (apinodeconfigs[i].roles & ROLE_STORAGE) {
            g_teems_config.num_storage_nodes += 1;
            if (i < my_node_num) {
                str_smaller.push_back(i);
            } else {
                g_teems_config.storage_nodes.push_back(i);
            }
            g_teems_config.storage_map.push_back(i);
        }
        cumul_weight += nw.weight;
        g_teems_config.weights.push_back(nw);
        g_teems_config.roles.push_back(apinodeconfigs[i].roles);
        if (i == my_node_num) {
            g_teems_config.my_weight = nw.weight;
        }
    }
    g_teems_config.tot_weight = cumul_weight;
    // Concatenate the *_smaller vectors to the ends of the
    // g_teems_config.*_nodes vectors.  This way, each node has a list
    // of nodes of each role starting with itself and "looping around".
    // This should make the communication pattern have less of a
    // bottleneck.
    g_teems_config.ingestion_nodes.insert(
        g_teems_config.ingestion_nodes.end(),
        ing_smaller.begin(),
        ing_smaller.end());
    g_teems_config.routing_nodes.insert(
        g_teems_config.routing_nodes.end(),
        rte_smaller.begin(),
        rte_smaller.end());
    g_teems_config.storage_nodes.insert(
        g_teems_config.storage_nodes.end(),
        str_smaller.begin(),
        str_smaller.end());

    // Initialize the threadpool and the pseudorandom bytes pools
    threadpool_init(nthreads);

    uint8_t my_role = apinodeconfigs[my_node_num].roles;
    if ( (my_role & ROLE_INGESTION) || (my_role & ROLE_STORAGE) ) {
        generateMasterKeys(g_teems_config.master_secret,
            g_teems_config.ESK, g_teems_config.TSK);

        uint32_t num_clients_total = g_teems_config.user_count;

        if (my_role & ROLE_INGESTION) {
            uint32_t num_ing_nodes = g_teems_config.num_ingestion_nodes;
            uint32_t clients_per_server =
                CEILDIV(num_clients_total, num_ing_nodes);
            uint32_t num_clients_this_ing = clients_per_server;
            uint32_t client_start = ing_smaller.size();
            g_ing.initialize(num_clients_this_ing, client_start,
                g_teems_config.ESK);
        }
    }

    ingestion_epoch = 0;
    storage_epoch = 0;

#ifdef TRACK_HEAP_USAGE
    printf("ecall_config_load H1 heap %u\n", g_peak_heap_used);
#endif
    if (!route_init()) {
        return false;
    }
#ifdef TRACK_HEAP_USAGE
    printf("ecall_config_load H2 heap %u\n", g_peak_heap_used);
#endif

    bool ret = comms_init_nodestate(apinodeconfigs, num_nodes, my_node_num);
#ifdef TRACK_HEAP_USAGE
    printf("ecall_config_load end heap %u\n", g_peak_heap_used);
#endif
    return ret;
}

void ecall_close()
{
    route_close();
    threadpool_shutdown();
}