#include <condition_variable>
#include <mutex>
#include <stdlib.h>

#include "Untrusted.hpp"
#include "start.hpp"

class Epoch {
    boost::asio::io_context &io_context;
    uint32_t epoch_num;
    std::mutex m;
    std::condition_variable cv;
    bool epoch_complete;

    void round_cb(uint32_t round_num) {
        if (round_num) {
            printf("Round %u complete\n", round_num);
            boost::asio::post(io_context, [this]{
                proceed();
            });
        } else {
            printf("Epoch %u complete\n", epoch_num);
            {
                std::lock_guard lk(m);
                epoch_complete = true;
            }
            cv.notify_one();
        }
    }

public:
    Epoch(boost::asio::io_context &context, uint32_t ep_num):
        io_context(context), epoch_num(ep_num),
        epoch_complete(false) {}

    void proceed() {
        ecall_routing_proceed([this](uint32_t round_num) {
            round_cb(round_num);
        });
    }

    void wait() {
        std::unique_lock lk(m);
        cv.wait(lk, [this]{ return epoch_complete; });
    }

};

static void epoch(NetIO &netio, char **args) {

    static uint32_t epoch_num = 1;

    uint16_t num_nodes = netio.num_nodes;
    uint32_t num_tokens[num_nodes];
    uint32_t tot_tokens = 0;
    for (nodenum_t j=0;j<num_nodes;++j) {
        num_tokens[j] = atoi(args[netio.me*num_nodes+j]);
        tot_tokens += num_tokens[j];
    }

    const Config &config = netio.config();
    uint16_t msg_size = config.msg_size;
    nodenum_t my_node_num = config.my_node_num;

    uint8_t *msgs = new uint8_t[tot_tokens * msg_size];
    uint8_t *nextmsg = msgs;
    uint32_t dest_uid_mask = (1 << DEST_UID_BITS) - 1;
    uint32_t rem_tokens = tot_tokens;
    while (rem_tokens > 0) {
        // Pick a random remaining token
        uint32_t r = uint32_t(lrand48()) % rem_tokens;
        for (nodenum_t j=0;j<num_nodes;++j) {
            if (r < num_tokens[j]) {
                // Use a token from node j
                *((uint32_t*)nextmsg) =
                    (j << DEST_UID_BITS) +
                        (((r<<8)+(my_node_num&0xff)) & dest_uid_mask);
                // Put a bunch of copies of r as the message body
                for (uint16_t i=1;i<msg_size/4;++i) {
                    ((uint32_t*)nextmsg)[i] = r;
                }
                num_tokens[j] -= 1;
                rem_tokens -= 1;
                nextmsg += msg_size;
            } else {
                r -= num_tokens[j];
            }
        }
    }
    /*
    for (uint32_t i=0;i<tot_tokens;++i) {
        for(uint16_t j=0;j<msg_size/4;++j) {
            printf("%08x ", ((uint32_t*)msgs)[i*msg_size/4+j]);
        }
        printf("\n");
    }
    */

    if (!ecall_ingest_raw(msgs, tot_tokens)) {
        printf("Ingestion failed\n");
        return;
    }

    Epoch epoch(netio.io_context(), epoch_num);
    epoch.proceed();
    epoch.wait();
    // Launch threads to refill the precomputed Waksman networks we
    // used, but just let them run in the background.
    size_t num_sizes = ecall_precompute_sort(-1);
    for (int i=0;i<int(num_sizes);++i) {
        boost::thread t([i] {
            ecall_precompute_sort(i);
        });
        t.detach();
    }
    ++epoch_num;
}

static void route_test(NetIO &netio, char **args)
{
    // Count the number of arguments
    size_t nargs = 0;
    while (args[nargs]) {
        ++nargs;
    }

    uint16_t num_nodes = netio.num_nodes;
    size_t sq_nodes = num_nodes;
    sq_nodes *= sq_nodes;

    if (nargs != sq_nodes) {
        printf("Expecting %lu arguments, found %lu\n", sq_nodes, nargs);
        return;
    }

    // The arguments are num_nodes sets of num_nodes values.  The jth
    // value in the ith set is the number of private routing tokens
    // ingestion node i holds for storage node j.

    // We are node i = netio.me, so ignore the other sets of values.

    // Precompute some WaksmanNetworks
    const Config &config = netio.config();
    size_t num_sizes = ecall_precompute_sort(-1);
    for (int i=0;i<int(num_sizes);++i) {
        std::vector<boost::thread> ts;
        for (int j=0; j<config.nthreads; ++j) {
            ts.emplace_back([i] {
                ecall_precompute_sort(i);
            });
        }
        for (auto& t: ts) {
            t.join();
        }
    }

    // The epoch interval, in microseconds
    uint32_t epoch_interval_us = 1000000;

    // Run 10 epochs
    for (int i=0; i<10; ++i) {
        struct timespec tp;
        clock_gettime(CLOCK_REALTIME_COARSE, &tp);
        unsigned long start = tp.tv_sec * 1000000 + tp.tv_nsec/1000;
        epoch(netio, args);
        clock_gettime(CLOCK_REALTIME_COARSE, &tp);
        unsigned long end = tp.tv_sec * 1000000 + tp.tv_nsec/1000;
        unsigned long diff = end - start;
        printf("Epoch time: %lu.%06lu s\n", diff/1000000, diff%1000000);
        // Sleep for the rest of the epoch interval
        if (diff < epoch_interval_us) {
            usleep(epoch_interval_us - (useconds_t)diff);
        }
    }
    netio.close();
}

// Once all the networking is set up, start doing whatever we were asked
// to do on the command line
void start(NetIO &netio, char **args)
{
    if (*args && !strcmp(*args, "route")) {
        ++args;
        route_test(netio, args);

        return;
    }
}

/*
    Handler for received client messages.
*/
void handle_async_clients(NetIO &netio, std::shared_ptr<tcp::socket> csocket,
    const boost::system::error_code& error, size_t auth_size,
    size_t msgbundle_size)
{
    if(!error) {
#ifdef VERBOSE_NET
        printf("Accept handler success\n");
#endif
        // Read header (1 uint64_t) from the socket and extract the client ID
        size_t header;
        clientid_t cid;
        boost::asio::read(*csocket,
             boost::asio::buffer(&header, sizeof(uint64_t)));

        if((header & 0xff) == CLIENT_AUTHENTICATE) {
            // Read the authentication token
            boost::asio::read(*csocket,
               boost::asio::buffer(&header, auth_size));

        } else if ((header & 0xff) == CLIENT_MESSAGE_BUNDLE) {
            unsigned char *msgbundle = (unsigned char*) malloc(msgbundle_size);
            cid = (clientid_t)(header >> 8);

            // Read the message_bundle
            boost::asio::read(*csocket,
               boost::asio::buffer(msgbundle, msgbundle_size));

            //Ingest the message_bundle
            bool ret = ecall_ingest_msgbundle(cid, msgbundle, apiparams.m_priv_out);
            free(msgbundle);
        }

        start_accept(netio, auth_size, msgbundle_size);
    } else {
        printf("Accept handler failed\n");
    }
}

/*
    Asynchronously accept client connections
*/
void start_accept(NetIO &netio, size_t auth_size, size_t msgbundle_size)
{
    boost::asio::io_context &io_context = netio.io_context();
    const NodeConfig &myconf = netio.myconfig();
    std::shared_ptr<tcp::socket> csocket(new tcp::socket(netio.io_context()));
    tcp::resolver resolver(io_context);
    std::shared_ptr<tcp::acceptor> client_acceptor;
    client_acceptor = std::shared_ptr<tcp::acceptor>(
        new tcp::acceptor(io_context,
            resolver.resolve(myconf.clistenhost,
            myconf.clistenport)->endpoint()));
#ifdef VERBOSE_NET
    std::cout << "Accepting on " << myconf.clistenhost << ":"
        << myconf.clistenport << "\n";
#endif
    client_acceptor->async_accept(*csocket,
        boost::bind(&handle_async_clients, netio, csocket,
        boost::asio::placeholders::error, auth_size, msgbundle_size));
}

void runAsyncListeners(NetIO &netio)
{
    const Config &conf = netio.config();
    const NodeConfig &myconf = netio.myconfig();
    if(myconf.roles & ROLE_INGESTION) {
        size_t auth_size, msgbundle_size;
        auth_size = SGX_AESGCM_MAC_SIZE;
        msgbundle_size = SGX_AESGCM_IV_SIZE +
            (conf.m_priv_out * conf.msg_size) + SGX_AESGCM_MAC_SIZE;
        start_accept(netio, auth_size, msgbundle_size);
    }
}