#include "framework/framework.h"

#include <stdio.h>

#include <string.h>

#define NUM_MSGS 40000

typedef struct {

    uint32_t recv_count;

} State;

static void on_friend_message(const Tox_Event_Friend_Message *event, void *user_data)

{

    ToxNode *self = (ToxNode *)user_data;

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

    state->recv_count++;

}

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

{

    State *state = (State *)ctx;

    Tox_Dispatch *dispatch = tox_node_get_dispatch(self);

    tox_events_callback_friend_message(dispatch, on_friend_message);

    tox_node_wait_for_self_connected(self);

    tox_node_wait_for_friend_connected(self, 0);

    tox_node_wait_for_friend_connected(self, 1);

    tox_node_log(self, "Ready to receive...");

    tox_scenario_barrier_wait(self); // Barrier 1: Connected

    // Wait until we get many messages or scenario timeout.

    // We don't expect ALL 80k messages to arrive in a short time,

    // but we want to see if it survives the flood.

    uint32_t last_count = 0;

    uint32_t stable_ticks = 0;

    while (tox_scenario_is_running(self)) {

        tox_scenario_yield(self);

        if (state->recv_count == last_count && state->recv_count > 0) {

            stable_ticks++;

            if (stable_ticks > 100) {

                break;    // No new messages for 5 seconds

            }

        } else {

            stable_ticks = 0;

        }

        last_count = state->recv_count;

        if (state->recv_count >= NUM_MSGS * 2) {

            break;

        }

    }

    tox_node_log(self, "Received %u messages.", state->recv_count);

}

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

{

    tox_node_wait_for_self_connected(self);

    tox_node_wait_for_friend_connected(self, 0);

    tox_scenario_barrier_wait(self); // Barrier 1: Connected

    tox_node_log(self, "Sending messages...");

    Tox *tox = tox_node_get_tox(self);

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

        uint8_t message[128] = {0};

        snprintf((char *)message, sizeof(message), "%u-%u", tox_node_get_index(self), i);

        Tox_Err_Friend_Send_Message err;

        tox_friend_send_message(tox, 0, TOX_MESSAGE_TYPE_NORMAL, message, sizeof(message), &err);

        if (err == TOX_ERR_FRIEND_SEND_MESSAGE_SENDQ) {

            // This is expected when flooding

            if (i % 1000 == 0) {

                tox_node_log(self, "Send queue full at message %u, yielding...", i);

            }

            tox_scenario_yield(self);

            i--; // Retry this message

            continue;

        }

        if (err == TOX_ERR_FRIEND_SEND_MESSAGE_FRIEND_NOT_CONNECTED) {

            if (i % 1000 == 0) {

                tox_node_log(self, "Friend not connected at message %u, waiting...", i);

            }

            tox_node_wait_for_friend_connected(self, 0);

            i--; // Retry this message

            continue;

        }

        ck_assert_msg(err == TOX_ERR_FRIEND_SEND_MESSAGE_OK, "send_message failed with %s",

                      tox_err_friend_send_message_to_string(err));

        if (i % 500 == 0) {

            tox_scenario_yield(self);

        }

    }

    tox_node_log(self, "Finished sending.");

}

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

{

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

    State receiver_state = {0};

    ToxNode *receiver = tox_scenario_add_node(s, "Receiver", receiver_script, &receiver_state, sizeof(State));

    ToxNode *sender1 = tox_scenario_add_node(s, "Sender1", sender_script, nullptr, 0);

    ToxNode *sender2 = tox_scenario_add_node(s, "Sender2", sender_script, nullptr, 0);

    tox_node_bootstrap(sender1, receiver);

    tox_node_bootstrap(sender2, receiver);

    tox_node_friend_add(receiver, sender1);

    tox_node_friend_add(sender1, receiver);

    tox_node_friend_add(receiver, sender2);

    tox_node_friend_add(sender2, receiver);

    ToxScenarioStatus res = tox_scenario_run(s);

    if (res != TOX_SCENARIO_DONE && res != TOX_SCENARIO_TIMEOUT) {

        return 1;

    }

    tox_scenario_free(s);

    return 0;

}

#undef NUM_MSGS