#include <pthread.h>
#include "Enclave_t.h"
#include "config.hpp"
#include "utils.hpp"
#include "sort.hpp"
#include "route.hpp"

#define PROFILE_ROUTING

struct MsgBuffer {
    pthread_mutex_t mutex;
    uint8_t *buf;
    // The number of messages (not bytes) in, or on their way into, the
    // buffer
    uint32_t reserved;
    // The number of messages definitely in the buffer
    uint32_t inserted;

    MsgBuffer() : buf(NULL), reserved(0), inserted(0) {
        pthread_mutex_init(&mutex, NULL);
    }

    ~MsgBuffer() {
        delete[] buf;
    }

    // The number passed is messages, not bytes
    void alloc(uint32_t msgs) {
        delete[] buf;
        buf = NULL;
        reserved = 0;
        inserted = 0;
        // This may throw bad_alloc, but we'll catch it higher up
        buf = new uint8_t[size_t(msgs) * g_teems_config.msg_size];
    }

    // Reset the contents of the buffer
    void reset() {
        memset(buf, 0, inserted * g_teems_config.msg_size);
        reserved = 0;
        inserted = 0;
    }

    // You can't copy a MsgBuffer
    MsgBuffer(const MsgBuffer&) = delete;
    MsgBuffer &operator=(const MsgBuffer&) = delete;
};

enum RouteStep {
    ROUTE_NOT_STARTED,
    ROUTE_ROUND_1,
    ROUTE_ROUND_2
};

// The round1 MsgBuffer stores messages we ingest while waiting for
// round 1 to start, which will be sorted and sent out in round 1.  The
// round2 MsgBuffer stores messages we receive in round 1, which will be
// padded, sorted, and sent out in round 2.

static struct RouteState {
    MsgBuffer round1;
    MsgBuffer round2;
    RouteStep step;
    uint32_t tot_msg_per_ing;
    uint32_t max_msg_to_each_str;
    uint32_t max_round2_msgs;
} route_state;

// Computes ceil(x/y) where x and y are integers, x>=0, y>0.

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

// Call this near the end of ecall_config_load, but before
// comms_init_nodestate. Returns true on success, false on failure.
bool route_init()
{
    // Compute the maximum number of messages we could receive by direct
    // ingestion

    // Each ingestion node will have at most
    // ceil(user_count/num_ingestion_nodes) users, and each user will
    // send at most m_priv_out messages.
    uint32_t users_per_ing = CEILDIV(g_teems_config.user_count,
        g_teems_config.num_ingestion_nodes);
    uint32_t tot_msg_per_ing = users_per_ing * g_teems_config.m_priv_out;

    // Compute the maximum number of messages we could receive in round 1
    // Each ingestion node will send us an our_weight/tot_weight
    // fraction of the messages they hold
    uint32_t max_msg_from_each_ing = CEILDIV(tot_msg_per_ing,
        g_teems_config.tot_weight) * g_teems_config.my_weight;

    // And the maximum number we can receive in total is that times the
    // number of ingestion nodes
    uint32_t max_round1_msgs = max_msg_from_each_ing *
        g_teems_config.num_ingestion_nodes;

    // Compute the maximum number of messages we could send in round 2

    // Each storage node has at most this many users
    uint32_t users_per_str = CEILDIV(g_teems_config.user_count,
        g_teems_config.num_storage_nodes);

    // And so can receive at most this many messages
    uint32_t tot_msg_per_str = users_per_str *
        g_teems_config.m_priv_in;

    // Which will be at most this many from us
    uint32_t max_msg_to_each_str = CEILDIV(tot_msg_per_str,
        g_teems_config.tot_weight) * g_teems_config.my_weight;

    // But we can't send more messages to each storage server than we
    // could receive in total
    if (max_msg_to_each_str > max_round1_msgs) {
        max_msg_to_each_str = max_round1_msgs;
    }

    // And the max total number of outgoing messages in round 2 is then
    uint32_t max_round2_msgs = max_msg_to_each_str *
        g_teems_config.num_storage_nodes;

    // In case we have a weird configuration where users can send more
    // messages per epoch than they can receive, ensure the round 2
    // buffer is large enough to hold the incoming messages as well
    if (max_round2_msgs < max_round1_msgs) {
        max_round2_msgs = max_round1_msgs;
    }

    /*
    printf("round1_msgs = %u, round2_msgs = %u\n",
        max_round1_msgs, max_round2_msgs);
    */

    // Create the route state
    try {
        route_state.round1.alloc(tot_msg_per_ing);
        route_state.round2.alloc(max_round2_msgs);
    } catch (std::bad_alloc&) {
        printf("Memory allocation failed in route_init\n");
        return false;
    }
    route_state.step = ROUTE_NOT_STARTED;
    route_state.tot_msg_per_ing = tot_msg_per_ing;
    route_state.max_msg_to_each_str = max_msg_to_each_str;
    route_state.max_round2_msgs = max_round2_msgs;

    threadid_t nthreads = g_teems_config.nthreads;
#ifdef PROFILE_ROUTING
    unsigned long start = printf_with_rtclock("begin precompute evalplans (%u,%hu) (%u,%hu)\n", tot_msg_per_ing, nthreads, max_round2_msgs, nthreads);
#endif
    sort_precompute_evalplan(tot_msg_per_ing, nthreads);
    sort_precompute_evalplan(max_round2_msgs, nthreads);
#ifdef PROFILE_ROUTING
    printf_with_rtclock_diff(start, "end precompute evalplans\n");
#endif
    return true;
}

// Precompute the WaksmanNetworks needed for the sorts.  If you pass -1,
// it will return the number of different sizes it needs.  If you pass
// [0,sizes-1], it will compute one WaksmanNetwork with that size index
// and return the number of available WaksmanNetworks of that size.

size_t ecall_precompute_sort(int sizeidx)
{
    size_t ret = 0;

    switch(sizeidx) {
    case 0:
#ifdef PROFILE_ROUTING
    {unsigned long start = printf_with_rtclock("begin precompute WaksmanNetwork (%u)\n", route_state.tot_msg_per_ing);
#endif
        ret = sort_precompute(route_state.tot_msg_per_ing);
#ifdef PROFILE_ROUTING
    printf_with_rtclock_diff(start, "end precompute Waksman Network (%u)\n", route_state.tot_msg_per_ing);}
#endif
        break;
    case 1:
#ifdef PROFILE_ROUTING
    {unsigned long start = printf_with_rtclock("begin precompute WaksmanNetwork (%u)\n", route_state.max_round2_msgs);
#endif
        ret = sort_precompute(route_state.max_round2_msgs);
#ifdef PROFILE_ROUTING
    printf_with_rtclock_diff(start, "end precompute Waksman Network (%u)\n", route_state.max_round2_msgs);}
#endif
        break;
    default:
        ret = 2;
        break;
    }

    return ret;
}

// Directly ingest a buffer of num_msgs messages into the round1 buffer.
// Return true on success, false on failure.
bool ecall_ingest_raw(uint8_t *msgs, uint32_t num_msgs)
{
    uint16_t msg_size = g_teems_config.msg_size;
    MsgBuffer &round1 = route_state.round1;

    pthread_mutex_lock(&round1.mutex);
    uint32_t start = round1.reserved;
    if (start + num_msgs > route_state.tot_msg_per_ing) {
        pthread_mutex_unlock(&round1.mutex);
        printf("Max %u messages exceeded\n",
            route_state.tot_msg_per_ing);
        return false;
    }
    round1.reserved += num_msgs;
    pthread_mutex_unlock(&round1.mutex);

    memmove(round1.buf + start * msg_size,
        msgs, num_msgs * msg_size);

    pthread_mutex_lock(&round1.mutex);
    round1.inserted += num_msgs;
    pthread_mutex_unlock(&round1.mutex);

    return true;
}

// Send the round 1 messages
static void send_round1_msgs(const uint8_t *msgs, const uint64_t *indices,
    uint32_t N)
{
    uint16_t msg_size = g_teems_config.msg_size;
    uint16_t tot_weight = g_teems_config.tot_weight;

    /*
    for (uint32_t i=0;i<N;++i) {
        const uint8_t *msg = msgs + indices[i]*msg_size;
        for (uint16_t j=0;j<msg_size/4;++j) {
            printf("%08x ", ((const uint32_t*)msg)[j]);
        }
        printf("\n");
    }
    */
}

// Perform the next round of routing.  The callback pointer will be
// passed to ocall_routing_round_complete when the round is complete.
void ecall_routing_proceed(void *cbpointer)
{
    if (route_state.step == ROUTE_NOT_STARTED) {

        MsgBuffer &round1 = route_state.round1;

        pthread_mutex_lock(&round1.mutex);
        // Ensure there are no pending messages currently being inserted
        // into the buffer
        while (round1.reserved != round1.inserted) {
            pthread_mutex_unlock(&round1.mutex);
            pthread_mutex_lock(&round1.mutex);
        }
        // Sort the messages we've received
#ifdef PROFILE_ROUTING
        uint32_t inserted = round1.inserted;
        unsigned long start = printf_with_rtclock("begin oblivious sort (%u,%u)\n", inserted, route_state.tot_msg_per_ing);
#endif
        sort_mtobliv(g_teems_config.nthreads, round1.buf,
            g_teems_config.msg_size, round1.inserted,
            route_state.tot_msg_per_ing, send_round1_msgs);
        round1.reset();
#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start, "end oblivious sort (%u,%u)\n", inserted, route_state.tot_msg_per_ing);
#endif
        route_state.step = ROUTE_ROUND_1;
        ocall_routing_round_complete(cbpointer, 1);
    }
}