#include "framework/framework.h"

#include "../../toxcore/tox_private.h"

#include <stdio.h>

#include <string.h>

#include <stdlib.h>

#define NUM_TOXES 30

typedef struct {

    uint8_t  public_key[TOX_DHT_NODE_PUBLIC_KEY_SIZE];

    char     ip[TOX_DHT_NODE_IP_STRING_SIZE];

    uint16_t port;

} Dht_Node;

typedef struct {

    Dht_Node *nodes[NUM_TOXES];

    size_t num_nodes;

    uint8_t public_key_list[NUM_TOXES][TOX_PUBLIC_KEY_SIZE];

} State;

static bool node_crawled(const State *state, const uint8_t *public_key)

{

    for (size_t i = 0; i < state->num_nodes; ++i) {

        if (memcmp(state->nodes[i]->public_key, public_key, TOX_DHT_NODE_PUBLIC_KEY_SIZE) == 0) {

            return true;

        }

    }

    return false;

}

static void on_dht_nodes_response(const Tox_Event_Dht_Nodes_Response *event, void *user_data)

{

    ToxNode *self = (ToxNode *)user_data;

    State *state = (State *)tox_node_get_script_ctx(self);

    const uint8_t *public_key = tox_event_dht_nodes_response_get_public_key(event);

    const char *ip = (const char *)tox_event_dht_nodes_response_get_ip(event);

    const uint16_t port = tox_event_dht_nodes_response_get_port(event);

    if (node_crawled(state, public_key)) {

        return;

    }

    if (state->num_nodes >= NUM_TOXES) {

        return;

    }

    Dht_Node *node = (Dht_Node *)calloc(1, sizeof(Dht_Node));

    ck_assert(node != nullptr);

    memcpy(node->public_key, public_key, TOX_DHT_NODE_PUBLIC_KEY_SIZE);

    snprintf(node->ip, sizeof(node->ip), "%s", ip);

    node->port = port;

    state->nodes[state->num_nodes] = node;

    ++state->num_nodes;

    // ask new node to give us their close nodes to every public key

    for (size_t i = 0; i < NUM_TOXES; ++i) {

        tox_dht_send_nodes_request(tox_node_get_tox(self), public_key, ip, port, state->public_key_list[i], nullptr);

    }

}

static void peer_script(ToxNode *self, void *ctx)

{

    State *state = (State *)ctx;

    tox_events_callback_dht_nodes_response(tox_node_get_dispatch(self), on_dht_nodes_response);

    // Initial bootstrap: ask the node we bootstrapped from for all other nodes

    // Wait for self connection first

    WAIT_UNTIL(tox_node_is_self_connected(self));

    // After connecting, we should start receiving responses because we bootstrapped

    // but the original test calls tox_dht_send_nodes_request for all nodes.

    // In our case, we bootstrap linearly, so Peer-i bootstraps from Peer-(i-1).

    // Peer-0 doesn't bootstrap from anyone.

    // To kick off crawling, each node can ask its bootstrap source for all nodes.

    uint32_t my_index = tox_node_get_index(self);

    if (my_index > 0) {

        ToxNode *bootstrap_source = tox_scenario_get_node(tox_node_get_scenario(self), my_index - 1);

        uint8_t bs_pk[TOX_PUBLIC_KEY_SIZE];

        tox_self_get_dht_id(tox_node_get_tox(bootstrap_source), bs_pk);

        Tox_Err_Get_Port port_err;

        uint16_t bs_port = tox_self_get_udp_port(tox_node_get_tox(bootstrap_source), &port_err);

        for (size_t i = 0; i < NUM_TOXES; ++i) {

            tox_dht_send_nodes_request(tox_node_get_tox(self), bs_pk, "127.0.0.1", bs_port, state->public_key_list[i], nullptr);

        }

    }

    // Wait until we have crawled all nodes

    uint64_t last_log_time = 0;

    uint64_t last_request_time = 0;

    while (state->num_nodes < NUM_TOXES) {

        tox_scenario_yield(self);

        uint64_t now = tox_scenario_get_time(tox_node_get_scenario(self));

        if (now - last_request_time >= 2000) {

            last_request_time = now;

            // Retry strategy: ask a random known node for missing nodes

            if (state->num_nodes > 0) {

                int idx = rand() % state->num_nodes;

                const Dht_Node *n = state->nodes[idx];

                for (int i = 0; i < NUM_TOXES; ++i) {

                    if (!node_crawled(state, state->public_key_list[i])) {

                        tox_dht_send_nodes_request(tox_node_get_tox(self),

                                                   n->public_key, n->ip, n->port, state->public_key_list[i], nullptr);

                    }

                }

            } else if (my_index > 0) {

                // Fallback: if we haven't found anyone yet, ask bootstrap peer again

                ToxNode *bootstrap_source = tox_scenario_get_node(tox_node_get_scenario(self), my_index - 1);

                uint8_t bs_pk[TOX_PUBLIC_KEY_SIZE];

                tox_self_get_dht_id(tox_node_get_tox(bootstrap_source), bs_pk);

                Tox_Err_Get_Port port_err;

                uint16_t bs_port = tox_self_get_udp_port(tox_node_get_tox(bootstrap_source), &port_err);

                for (int i = 0; i < NUM_TOXES; ++i) {

                    tox_dht_send_nodes_request(tox_node_get_tox(self), bs_pk, "127.0.0.1", bs_port, state->public_key_list[i], nullptr);

                }

            }

        }

        if (now - last_log_time >= 5000) {

            last_log_time = now;

            tox_node_log(self, "Still crawling... found %zu/%d nodes", state->num_nodes, NUM_TOXES);

            if (state->num_nodes < NUM_TOXES) {

                char missing[256] = {0};

                size_t pos = 0;

                for (int i = 0; i < NUM_TOXES; ++i) {

                    if (!node_crawled(state, state->public_key_list[i])) {

                        pos += snprintf(missing + pos, sizeof(missing) - pos, "%d ", i);

                        if (pos >= sizeof(missing) - 1) {

                            break;

                        }

                    }

                }

                tox_node_log(self, "Missing nodes: %s", missing);

            }

        }

    }

    tox_node_log(self, "Finished crawling all %u nodes.", (unsigned int)state->num_nodes);

}

int main(int argc, char *argv[])

{

    ToxScenario *s = tox_scenario_new(argc, argv, 60000);

    State *states = (State *)calloc(NUM_TOXES, sizeof(State));

    ToxNode *nodes[NUM_TOXES];

    for (uint32_t i = 0; i < NUM_TOXES; ++i) {

        char alias[32];

        snprintf(alias, sizeof(alias), "Peer-%u", i);

        nodes[i] = tox_scenario_add_node(s, alias, peer_script, &states[i], sizeof(State));

    }

    // Fill the public_key_list for all nodes

    uint8_t public_keys[NUM_TOXES][TOX_PUBLIC_KEY_SIZE];

    for (uint32_t i = 0; i < NUM_TOXES; ++i) {

        tox_self_get_dht_id(tox_node_get_tox(nodes[i]), public_keys[i]);

    }

    for (uint32_t i = 0; i < NUM_TOXES; ++i) {

        memcpy(states[i].public_key_list, public_keys, sizeof(public_keys));

    }

    // Create a linear graph

    for (uint32_t i = 1; i < NUM_TOXES; ++i) {

        tox_node_bootstrap(nodes[i], nodes[i - 1]);

    }

    ToxScenarioStatus res = tox_scenario_run(s);

    if (res != TOX_SCENARIO_DONE) {

        fprintf(stderr, "Scenario failed with status %u\n", res);

        return 1;

    }

    for (uint32_t i = 0; i < NUM_TOXES; ++i) {

        for (size_t j = 0; j < states[i].num_nodes; ++j) {

            free(states[i].nodes[j]);

        }

    }

    free(states);

    tox_scenario_free(s);

    return 0;

}

#undef NUM_TOXES