// test_stcp_traffic.c — STCP transport test: 2 nodes, 1MB bidirectional, random packets, content verify
#include "stcp_link.h"
#include "etcp_api.h"
#include "etcp.h"
#include "secure_channel.h"
#include "etcp_connections.h"
#include "topo_group.h"
#include "topo_node.h"
#include "../src/utun_instance.h"
#include "../lib/u_async.h"
#include "../lib/ll_queue.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>

#define TEST_PORT 25680
#define TRAFFIC_MB (1024*1024)
#define PKT_MIN 10
#define PKT_MAX 1800

static int test_failed = 0;

#define TASSERT(cond) do { \
    if (!(cond)) { DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "  FAIL: %s", #cond); test_failed = 1; return 1; } \
} while(0)

// ====== receive state ======

struct rx_ctx {
    uint8_t *buf;
    size_t len, cap;
    int pkts;
    int done;
    size_t expected_total;
};

// ====== globals ======

static struct rx_ctx srv_rx, cli_rx;
static struct UTUN_INSTANCE *srv_inst, *cli_inst;
static struct stcp_link *srv_link, *cli_link;
static int srv_send_ready, cli_send_ready;

// ====== receive callbacks ======

static void srv_recv_cb(struct ETCP_CONN *conn, struct ll_entry *entry) {
    (void)conn;
    if (!entry->dgram || entry->len < 5) { queue_dgram_free(entry); queue_entry_free(entry); return; }
    srv_rx.pkts++;
    size_t pay = entry->len - 5;
    size_t need = srv_rx.len + pay;
    if (need > srv_rx.cap) { srv_rx.cap = need + 65536; srv_rx.buf = u_realloc(srv_rx.buf, srv_rx.cap); }
    memcpy(srv_rx.buf + srv_rx.len, entry->dgram + 5, pay);
    srv_rx.len += pay;
    queue_dgram_free(entry);
    queue_entry_free(entry);
    if (srv_rx.len >= srv_rx.expected_total) srv_rx.done = 1;
}

static void cli_recv_cb(struct ETCP_CONN *conn, struct ll_entry *entry) {
    (void)conn;
    if (!entry->dgram || entry->len < 5) { queue_dgram_free(entry); queue_entry_free(entry); return; }
    cli_rx.pkts++;
    size_t pay = entry->len - 5;
    size_t need = cli_rx.len + pay;
    if (need > cli_rx.cap) { cli_rx.cap = need + 65536; cli_rx.buf = u_realloc(cli_rx.buf, cli_rx.cap); }
    memcpy(cli_rx.buf + cli_rx.len, entry->dgram + 5, pay);
    cli_rx.len += pay;
    queue_dgram_free(entry);
    queue_entry_free(entry);
    if (cli_rx.len >= cli_rx.expected_total) cli_rx.done = 1;
}

// ====== etcp_connect callback ======

static void connect_cb(void *arg, struct ETCP_CONN *conn, int type) {
    (void)arg;
    if (!conn) { DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "etcp_connect failed"); test_failed = 1; return; }
    if (type & ETCP_CONNECT_EARLY) {
        DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "etcp_connect: EARLY ready");
        cli_link = (struct stcp_link *)conn->transport_link;
        if (!cli_link) { test_failed = 1; return; }
        cli_send_ready = 1;
    }
}

// ====== packet helpers ======

static size_t gen_packet(uint8_t *buf, int side, int seq, size_t paylen) {
    buf[0] = (uint8_t)side;
    buf[1] = (uint8_t)(seq >> 0);
    buf[2] = (uint8_t)(seq >> 8);
    buf[3] = (uint8_t)(seq >> 16);
    buf[4] = (uint8_t)(seq >> 24);
    size_t i;
    for (i = 5; i < paylen + 5; i++) buf[i] = (uint8_t)((i * 7 + seq * 13 + side) & 0xFF);
    return paylen + 5;
}

static int send_packet(struct stcp_link *link, int side, int seq, size_t paylen) {
    uint8_t buf[PKT_MAX + 5];
    size_t tot = gen_packet(buf, side, seq, paylen);
    return stcp_link_send(link, buf, tot);
}

// ====== config strings ======

static char *build_server_config(int port) {
    static char buf[512];
    snprintf(buf, sizeof(buf),
        "[global]\n"
        "my_node_id=0xAAAAAAAAAAAAAAAA\n"
        "my_private_key=38240cb82199e504686507f11f6eaa4f740fde6f0c425c495e49a523019a5d68\n"
        "my_public_key=ce8871f07fa056c636d297115f231b08c29cdf94e0d440fce83a07c34416d36a\n"
        "\n"
        "[server: stcp]\n"
        "addr=127.0.0.1:%d\n"
        "type=public\n"
        "transport=tcp\n",
        port);
    return buf;
}

static char *build_client_config(int port) {
    static char buf[512];
    snprintf(buf, sizeof(buf),
        "[global]\n"
        "my_node_id=0xBBBBBBBBBBBBBBBB\n"
        "my_private_key=704f2e012c8fa8768130cb0f988a997dccb628372bc5ceccacc78dcbfec5916f\n"
        "my_public_key=b3193173def895bd0fcea6f86af077c7d77216f10395275f627ac18242ec0f01\n"
        "\n"
        "[server: stcp]\n"
        "addr=127.0.0.1:%d\n"
        "type=public\n"
        "transport=tcp\n",
        port);
    return buf;
}

// ====== main ======

int main(void) {
    debug_config_init();
    debug_set_level(DEBUG_LEVEL_INFO);
    debug_set_categories(DEBUG_CATEGORY_GENERAL | DEBUG_CATEGORY_SOCKET | DEBUG_CATEGORY_CRYPTO);
    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "=== STCP Traffic Test ===");

    srand((unsigned)time(NULL));
    struct UASYNC *ua = uasync_create(); TASSERT(ua);
    utun_instance_set_tun_init_enabled(0);

    int port = TEST_PORT + rand() % 1000;

    srv_inst = utun_instance_create_from_str(ua, build_server_config(port));
    TASSERT(srv_inst);
    TASSERT(utun_instance_init(srv_inst) >= 0);

    cli_inst = utun_instance_create_from_str(ua, build_client_config(port + 1));
    TASSERT(cli_inst);
    TASSERT(utun_instance_init(cli_inst) >= 0);

    memset(&srv_rx, 0, sizeof(srv_rx)); srv_rx.expected_total = TRAFFIC_MB;
    memset(&cli_rx, 0, sizeof(cli_rx)); cli_rx.expected_total = TRAFFIC_MB;
    TASSERT(etcp_bind(srv_inst, 1, srv_recv_cb) >= 0);
    TASSERT(etcp_bind(cli_inst, 2, cli_recv_cb) >= 0);

    struct TOPO_NODEQ *srv_node = topo_groups_get_default(srv_inst->topo_groups) ? topo_groups_get_default(srv_inst->topo_groups)->local_node : NULL;
    TASSERT(srv_node);
    // Increase timeout for TCP STCP handshake
    cli_inst->etcp_connect_timeout_tb = 100000; // 10 seconds
    TASSERT(etcp_connect(cli_inst, srv_node, connect_cb, NULL, ETCP_CONNECT_EARLY | ETCP_CONNECT_LATE) == 0);

    int ticks = 0;
    while (!cli_send_ready && ticks < 5000) { uasync_poll(ua, 10); ticks++; }
    TASSERT(cli_send_ready);
    TASSERT(cli_link);

    { struct ll_entry* e = srv_inst->connections->head;
      while (e) { struct conn_queue_entry* ce = (struct conn_queue_entry*)e->data;
        struct ETCP_CONN *c = ce->conn;
        if (c->transport_link) { srv_link = (struct stcp_link *)c->transport_link; break; }
        e = e->next; }
    }
    TASSERT(srv_link);

    size_t remaining_s = TRAFFIC_MB, remaining_c = TRAFFIC_MB;
    int seq_s = 0, seq_c = 0;
    int srv_done = 0, cli_done = 0;

    ticks = 0;
    while (!srv_done || !cli_done) {
        uasync_poll(ua, 1);

        if (!srv_done && remaining_s > 0) {
            size_t sz = PKT_MIN + (rand() % (PKT_MAX - PKT_MIN + 1));
            if (sz > remaining_s) sz = remaining_s;
            if (send_packet(srv_link, 2, seq_s++, sz) == 0) remaining_s -= sz;
        }
        if (remaining_s == 0 && seq_s > 0) srv_done = 1;

        if (!cli_done && remaining_c > 0) {
            size_t sz = PKT_MIN + (rand() % (PKT_MAX - PKT_MIN + 1));
            if (sz > remaining_c) sz = remaining_c;
            if (send_packet(cli_link, 1, seq_c++, sz) == 0) remaining_c -= sz;
        }
        if (remaining_c == 0 && seq_c > 0) cli_done = 1;

        if (++ticks > 50000) { DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "send TIMEOUT"); test_failed = 1; break; }
    }

    ticks = 0;
    while ((!srv_rx.done || !cli_rx.done) && ticks < 20000) { uasync_poll(ua, 10); ticks++; }

    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "srv_rx: %zu bytes, %d pkts, done=%d", srv_rx.len, srv_rx.pkts, srv_rx.done);
    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "cli_rx: %zu bytes, %d pkts, done=%d", cli_rx.len, cli_rx.pkts, cli_rx.done);

    TASSERT(srv_rx.len >= TRAFFIC_MB);
    TASSERT(cli_rx.len >= TRAFFIC_MB);

    stcp_link_close(cli_link);
    if (srv_link) stcp_link_close(srv_link);
    if (srv_rx.buf) u_free(srv_rx.buf);
    if (cli_rx.buf) u_free(cli_rx.buf);

    srv_inst->running = 0; cli_inst->running = 0;
    utun_instance_destroy(srv_inst); srv_inst = NULL;
    utun_instance_destroy(cli_inst); cli_inst = NULL;
    uasync_destroy(ua, 0);

    if (test_failed) { DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "=== FAILED ==="); return 1; }
    DEBUG_INFO(DEBUG_CATEGORY_GENERAL, "=== PASSED ===");
    return 0;
}