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

Config g_teems_config;

bool ecall_config_load(threadid_t nthreads, bool private_routing,
    EnclaveAPIParams *apiparams,
    EnclaveAPINodeConfig *apinodeconfigs,
    nodenum_t num_nodes, nodenum_t my_node_num)
{
    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_priv_out = apiparams->m_priv_out;
    g_teems_config.m_priv_in = apiparams->m_priv_in;
    g_teems_config.m_pub_out = apiparams->m_pub_out;
    g_teems_config.m_pub_in = apiparams->m_pub_in;
    g_teems_config.private_routing = private_routing;
    // 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();
    for (nodenum_t i=0; i<num_nodes; ++i) {
        NodeWeight nw;
        nw.startweight = cumul_weight;
        nw.weight = apinodeconfigs[i].weight;
        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) {
            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);
            }
        }
        cumul_weight += nw.weight;
        g_teems_config.weights.push_back(nw);
    }
    // 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);
    PRB_pool_init(nthreads);

    return comms_init_nodestate(apinodeconfigs, num_nodes, my_node_num);
}

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