#include <stdlib.h>

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

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.

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

    // Precompute some WaksmanNetworks
    size_t num_sizes = ecall_precompute_sort(-1);
    for (int i=0;i<int(num_sizes);++i) {
        ecall_precompute_sort(i);
    }

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

    ecall_routing_proceed([&](uint32_t round_num) {
        printf("Round %u complete\n", round_num);
        if (round_num == 1) {
            boost::asio::post(netio.io_context(), []{
                ecall_routing_proceed([&](uint32_t round_num2) {
                    printf("Round %u complete\n", round_num2);
                });
            });
        }
    });
}

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