#include "framework/framework.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#define NUM_TOXES 40
#define NUM_FRIEND_PAIRS 50
#define FR_MESSAGE "Gentoo"
#define RELAY_TCP_PORT 33448

typedef struct {
    uint32_t expected_friends;
} State;

static void on_friend_request(const Tox_Event_Friend_Request *event, void *user_data)
{
    ToxNode *self = (ToxNode *)user_data;
    const uint8_t *msg = tox_event_friend_request_get_message(event);
    size_t len = tox_event_friend_request_get_message_length(event);

    if (len == 6 && memcmp(msg, FR_MESSAGE, 6) == 0) {
        tox_friend_add_norequest(tox_node_get_tox(self), tox_event_friend_request_get_public_key(event), nullptr);
    }
}

static void peer_script(ToxNode *self, void *ctx)
{
    State *state = (State *)ctx;
    tox_events_callback_friend_request(tox_node_get_dispatch(self), on_friend_request);

    tox_node_log(self, "Waiting for connection to relay...");
    tox_node_wait_for_self_connected(self);
    tox_node_log(self, "Connected. Expected friends: %u", state->expected_friends);

    // Wait until we have the expected number of friends connected via TCP
    uint32_t last_connected = 0;
    while (1) {
        uint32_t connected_friends = 0;
        uint32_t total_friends = tox_self_get_friend_list_size(tox_node_get_tox(self));
        for (uint32_t i = 0; i < total_friends; ++i) {
            if (tox_node_get_friend_connection_status(self, i) == TOX_CONNECTION_TCP) {
                connected_friends++;
            }
        }

        if (connected_friends != last_connected) {
            tox_node_log(self, "Connected friends: %u/%u", connected_friends, state->expected_friends);
            last_connected = connected_friends;
        }

        if (connected_friends >= state->expected_friends) {
            break;
        }
        tox_scenario_yield(self);
    }

    tox_node_log(self, "All %u friends connected via TCP.", state->expected_friends);
}

static void relay_script(ToxNode *self, void *ctx)
{
    (void)ctx;
    ToxScenario *s = tox_node_get_scenario(self);

    while (tox_scenario_is_running(self)) {
        bool all_finished = true;
        for (uint32_t i = 0; i < NUM_TOXES; ++i) {
            ToxNode *peer = tox_scenario_get_node(s, i + 1); // Peers start at index 1
            if (!tox_node_is_finished(peer)) {
                all_finished = false;
                break;
            }
        }
        if (all_finished) {
            break;
        }
        tox_scenario_yield(self);
    }
}

int main(int argc, char *argv[])
{
    ToxScenario *s = tox_scenario_new(argc, argv, 300000);

    struct Tox_Options *opts_relay = tox_options_new(nullptr);
    tox_options_set_tcp_port(opts_relay, RELAY_TCP_PORT);
    ToxNode *relay = tox_scenario_add_node_ex(s, "Relay", relay_script, nullptr, 0, opts_relay);
    tox_options_free(opts_relay);

    uint8_t relay_dht_id[TOX_PUBLIC_KEY_SIZE];
    tox_self_get_dht_id(tox_node_get_tox(relay), relay_dht_id);

    uint16_t relay_tcp_port = tox_self_get_tcp_port(tox_node_get_tox(relay), nullptr);

    State states[NUM_TOXES] = {0};
    ToxNode *nodes[NUM_TOXES];

    struct Tox_Options *opts_peer = tox_options_new(nullptr);
    tox_options_set_udp_enabled(opts_peer, false);
    tox_options_set_local_discovery_enabled(opts_peer, false);

    for (int i = 0; i < NUM_TOXES; ++i) {
        char alias[16];
        snprintf(alias, sizeof(alias), "Tox%d", i);
        nodes[i] = tox_scenario_add_node_ex(s, alias, peer_script, &states[i], sizeof(State), opts_peer);

        // All peers use the Relay for TCP and DHT bootstrapping
        tox_add_tcp_relay(tox_node_get_tox(nodes[i]), "127.0.0.1", relay_tcp_port, relay_dht_id, nullptr);
        tox_node_bootstrap(nodes[i], relay);
    }
    tox_options_free(opts_peer);

    struct {
        uint16_t t1;
        uint16_t t2;
    } pairs[NUM_FRIEND_PAIRS];

    // Generate friend pairs
    for (int i = 0; i < NUM_FRIEND_PAIRS; ++i) {
        bool unique;
        do {
            unique = true;
            pairs[i].t1 = rand() % NUM_TOXES;
            pairs[i].t2 = (pairs[i].t1 + rand() % (NUM_TOXES - 1) + 1) % NUM_TOXES;

            // Avoid reciprocal or duplicate pairs
            for (int j = 0; j < i; ++j) {
                if ((pairs[j].t1 == pairs[i].t1 && pairs[j].t2 == pairs[i].t2) ||
                        (pairs[j].t1 == pairs[i].t2 && pairs[j].t2 == pairs[i].t1)) {
                    unique = false;
                    break;
                }
            }
        } while (!unique);

        uint8_t addr[TOX_ADDRESS_SIZE];
        tox_self_get_address(tox_node_get_tox(nodes[pairs[i].t1]), addr);

        Tox_Err_Friend_Add err;
        tox_friend_add(tox_node_get_tox(nodes[pairs[i].t2]), addr, (const uint8_t *)FR_MESSAGE, 6, &err);
        ck_assert(err == TOX_ERR_FRIEND_ADD_OK);

        states[pairs[i].t1].expected_friends++;
        states[pairs[i].t2].expected_friends++;
    }

    ToxScenarioStatus res = tox_scenario_run(s);
    if (res != TOX_SCENARIO_DONE) {
        tox_scenario_log(s, "Test failed with status %u", res);
        return 1;
    }

    tox_scenario_free(s);
    return 0;
}

#undef FR_MESSAGE
#undef NUM_TOXES
#undef NUM_FRIEND_PAIRS
#undef RELAY_TCP_PORT