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

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) {}

    ~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];
    }

    // 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;

    return true;
}

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