#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();
        }
    }

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