// Реальные STCP/REALITY-линки: встречное подключение и рестарт пира без reset выжившего узла.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "etcp.h"
#include "etcp_api.h"
#include "etcp_connections.h"
#include "etcp_session.h"
#include "node_conn_direct.h"
#include "topo_group.h"
#include "topo_node.h"
#include "utun_instance.h"
#include "../lib/debug_config.h"
#include "../lib/mem.h"
#include "test_clock.h"

#define CHECK(x) do { if (!(x)) { DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "FAIL %s:%d: %s", __func__, __LINE__, #x); exit(1); } } while (0)
#define TEST_CMD 253
#define MESSAGE_COUNT 16
#define FULL_MASK ((1u << MESSAGE_COUNT) - 1)

struct endpoint {
    struct UTUN_INSTANCE* inst;
    struct NODE_CONN_DIRECT* owner;
    struct ETCP_CONN* conn;
    uint64_t id;
    uint16_t port;
    unsigned received, reinit, deleted;
};
struct message { uint8_t cmd, from, phase, seq, body[64]; };
static struct endpoint endpoints[2];
static struct UASYNC* ua;
static unsigned phase;
static int failed;
static uint64_t stream_packets, session_packets;
static const char* public_keys[2] = {
    "b3193173def895bd0fcea6f86af077c7d77216f10395275f627ac18242ec0f01",
    "ce8871f07fa056c636d297115f231b08c29cdf94e0d440fce83a07c34416d36a"
};
static const char* private_keys[2] = {
    "704f2e012c8fa8768130cb0f988a997dccb628372bc5ceccacc78dcbfec5916f",
    "38240cb82199e504686507f11f6eaa4f740fde6f0c425c495e49a523019a5d68"
};

#ifdef __linux__
// Считаем реальные отправки потока/control: тест ловит шторм независимо от скорости CPU.
int __real_etcp_encrypt_send(struct ETCP_DGRAM* packet);
int __wrap_etcp_encrypt_send(struct ETCP_DGRAM* packet) {
    if (packet && packet->data_len) {
        if (packet->data[0] <= ETCP_SECTION_ACK) stream_packets++;
        if (packet->data[0] >= ETCP_SESSION_HELLO && packet->data[0] <= ETCP_SESSION_CONFIRM) session_packets++;
    }
    return __real_etcp_encrypt_send(packet);
}
#endif

static void fill_message(struct message* message, int from, unsigned seq) {
    message->cmd = TEST_CMD; message->from = from; message->phase = phase; message->seq = seq;
    for (unsigned i = 0; i < sizeof(message->body); i++) message->body[i] = (phase * 17 + seq * 3 + i) ^ from;
}

static void receive_message(struct ETCP_CONN* conn, struct ll_entry* entry) {
    int to = conn->instance == endpoints[0].inst ? 0 : 1;
    struct message expected;
    int valid = entry->len == sizeof(expected) && entry->dgram;
    const struct message* received = valid ? (const struct message*)entry->dgram : NULL;
    if (valid) valid = received->phase == phase && received->from == 1 - to && received->seq < MESSAGE_COUNT;
    if (valid) {
        fill_message(&expected, 1 - to, received->seq);
        valid = !memcmp(received, &expected, sizeof(expected)) && !(endpoints[to].received & (1u << received->seq));
    }
    if (valid) endpoints[to].received |= 1u << received->seq;
    else { DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "Invalid/duplicate message: to=%d phase=%u bytes=%u", to, phase, entry->len); failed = 1; }
    queue_dgram_free(entry); queue_entry_free(entry);
}

static void connection_event(struct ETCP_CONN* conn, int event, void* arg) {
    (void)conn;
    struct endpoint* endpoint = arg;
    if (event == ETCP_CBK_EVENT_REINIT) endpoint->reinit++;
    if (event == ETCP_CBK_EVENT_DELETE) endpoint->deleted++;
}

static void connect_done(void* arg, struct ETCP_CONN* conn, int type) {
    struct endpoint* endpoint = arg;
    if (!conn) { DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "Connection failed: local=%016llx", (unsigned long long)endpoint->id); failed = 1; }
    else if (type & ETCP_CONNECT_EARLY) endpoint->conn = conn;
}

static void start_endpoint(int i, int reality) {
    struct endpoint* endpoint = &endpoints[i];
    struct SC_MYKEYS keys;
    CHECK(sc_init_local_keys(&keys, public_keys[i], private_keys[i]) == SC_OK);
    uint64_t id = sc_derive_node_id_from_pubkey(keys.public_key);
    char config[2048], reality_config[512] = "";
    if (reality) snprintf(reality_config, sizeof(reality_config),
        "[reality]\nenabled=1\nprivate_key=%s\nshort_ids=0102030405060708\nserver_name=localhost.test\n"
        "dest=127.0.0.1:1\nfingerprint=openssl\n", private_keys[i]);
    snprintf(config, sizeof(config),
        "[global]\nmy_node_id=0x%016llx\nmy_node_name=transport-%d\nmy_public_key=%s\nmy_private_key=%s\n"
        "[server:stcp]\naddr=127.0.0.1:%u\ntype=public\ntransport=tcp\nreality=%d\n%s",
        (unsigned long long)id, i, public_keys[i], private_keys[i], endpoint->port, reality, reality_config);
    endpoint->inst = utun_instance_create_from_str(ua, config); CHECK(endpoint->inst);
    CHECK(utun_instance_init(endpoint->inst) >= 0);
    CHECK(etcp_bind(endpoint->inst, TEST_CMD, receive_message) == 0);
    endpoint->id = endpoint->inst->node_id;
    struct ETCP_SOCKET* socket = endpoint->inst->etcp_sockets;
    while (socket && !socket->is_tcp) socket = socket->next;
    CHECK(socket && socket->local_addr.ss_family == AF_INET);
    endpoint->port = ntohs(((struct sockaddr_in*)&socket->local_addr)->sin_port); CHECK(endpoint->port);
    endpoint->conn = NULL; endpoint->owner = NULL; endpoint->received = endpoint->reinit = endpoint->deleted = 0;
}

static void stop_endpoint(int i) {
    struct endpoint* endpoint = &endpoints[i];
    if (endpoint->owner) { node_conn_direct_close(endpoint->owner); endpoint->owner = NULL; }
    endpoint->inst->running = 0;
    utun_instance_destroy(endpoint->inst); endpoint->inst = NULL; endpoint->conn = NULL;
}

static struct TOPO_NODE* import_peer(int to) {
    struct UTUN_INSTANCE* from = endpoints[1-to].inst;
    struct TOPO_GROUP_NODE* local = topo_groups_get_default(from->topo_groups)->local_node;
    struct TOPO_NODE* info = topo_node_registry_find(from->topo_groups, from->node_id); CHECK(info);
    uint8_t wire[4096]; int len = topo_node_serialize(info, local, 0, 0, wire, sizeof(wire), 0); CHECK(len > 0);
    struct TOPO_NODE* imported = NULL;
    CHECK(topo_node_deserialize(topo_groups_get_default(endpoints[to].inst->topo_groups), wire, len,
        &imported, NULL, NULL, NULL, NULL) == 0);
    CHECK(topo_node_registry_store(endpoints[to].inst->topo_groups, imported));
    return imported;
}

static int pair_ready(void) {
    for (int i = 0; i < 2; i++) {
        endpoints[i].conn = instance_find_conn(endpoints[i].inst, endpoints[1-i].id);
        struct ETCP_CONN* conn = endpoints[i].conn;
        if (!conn || !conn->initialized || !conn->links_up || conn->session_required || conn->reinit_pending) return 0;
    }
    return endpoints[0].conn->peer_reset_id == endpoints[1].conn->reset_id &&
        endpoints[1].conn->peer_reset_id == endpoints[0].conn->reset_id;
}

static void wait_pair(void) {
    uint64_t end = get_time_tb() + 50000;
    while (!failed && get_time_tb() < end) {
        test_clock_poll(ua, 100);
        if (pair_ready()) return;
    }
    DEBUG_ERROR(DEBUG_CATEGORY_GENERAL, "Session recovery timed out: phase=%u stream=%llu control=%llu", phase,
        (unsigned long long)stream_packets, (unsigned long long)session_packets);
    CHECK(0);
}

static void check_unique(int i) {
    unsigned count = 0;
    for (struct ll_entry* e = endpoints[i].inst->connections->head; e; e = e->next)
        if (((struct conn_queue_entry*)e->data)->peer_node_id == endpoints[1-i].id) count++;
    CHECK(count == 1);
}

static void hold_connection(int i) {
    struct endpoint* endpoint = &endpoints[i];
    struct TOPO_NODE* peer = topo_node_registry_find(endpoint->inst->topo_groups, endpoints[1-i].id); CHECK(peer);
    CHECK(node_conn_direct_open_node(endpoint->inst, peer->node_id, NULL, NULL, &endpoint->owner, peer, NULL) == NCD_REUSED);
    CHECK(node_conn_direct_get_conn(endpoint->owner) == endpoint->conn);
    etcp_conn_add_cbk(endpoint->conn, connection_event, endpoint, ETCP_CBK_EVENT_REINIT | ETCP_CBK_EVENT_DELETE);
}

static void exchange_messages(void) {
    for (int i = 0; i < 2; i++) endpoints[i].received = 0;
    for (int i = 0; i < 2; i++) for (unsigned seq = 0; seq < MESSAGE_COUNT; seq++) {
        struct ll_entry* entry = ll_alloc_lldgram(sizeof(struct message)); CHECK(entry);
        fill_message((struct message*)entry->dgram, i, seq); entry->len = sizeof(struct message);
        CHECK(etcp_send(endpoints[i].conn, entry) == 0);
    }
    uint64_t end = get_time_tb() + 30000;
    while (!failed && get_time_tb() < end) {
        test_clock_poll(ua, 100);
        if (endpoints[0].received != FULL_MASK || endpoints[1].received != FULL_MASK) continue;
        int drained = 1;
        for (int i = 0; i < 2; i++) {
            struct ETCP_CONN* conn = endpoints[i].conn;
            if (conn->input_wait_ack->count || conn->input_send_q->count || conn->input_queue->count || conn->send_input_q->count) drained = 0;
        }
        if (drained) break;
    }
    CHECK(!failed && endpoints[0].received == FULL_MASK && endpoints[1].received == FULL_MASK);
    for (int i = 0; i < 2; i++) CHECK(!endpoints[i].conn->input_wait_ack->count && !endpoints[i].conn->input_send_q->count);
}

static void run_case(int reality, int simultaneous, int restarted) {
    DEBUG_INFO(DEBUG_CATEGORY_DEBUG, "Transport case: reality=%d simultaneous=%d restart=%d", reality, simultaneous, restarted);
    memset(endpoints, 0, sizeof(endpoints)); phase = 0; failed = 0; stream_packets = session_packets = 0;
    ua = uasync_create(); CHECK(ua);
    start_endpoint(0, reality); start_endpoint(1, reality);
    struct TOPO_NODE* peers[2] = {import_peer(0), import_peer(1)};
    for (int i = 0; i <= simultaneous; i++) {
        struct TOPO_GROUP_NODE node = {0}; node.node_id = peers[i]->node_id;
        CHECK(etcp_connect(endpoints[i].inst, &node, connect_done, &endpoints[i], ETCP_CONNECT_EARLY | ETCP_CONNECT_LATE) == 0);
    }
    wait_pair();
    for (int i = 0; i < 2; i++) { check_unique(i); hold_connection(i); }
    // Повторные владельцы и обновление адресов не создают дополнительные линки.
    unsigned links = 0; for (struct ETCP_LINK* l = endpoints[0].conn->links; l; l = l->next) links++;
    struct NODE_CONN_DIRECT* extra = NULL;
    CHECK(node_conn_direct_open_node(endpoints[0].inst, peers[0]->node_id, NULL, NULL, &extra, peers[0], NULL) == NCD_REUSED);
    CHECK(node_conn_direct_update_node(endpoints[0].inst, peers[0]->node_id) == 0);
    unsigned updated = 0; for (struct ETCP_LINK* l = endpoints[0].conn->links; l; l = l->next) updated++;
    CHECK(updated == links); node_conn_direct_close(extra);
    exchange_messages();

    if (restarted < 2) {
        int survivor = 1 - restarted;
        struct ETCP_CONN* retained = endpoints[survivor].conn;
        CHECK(etcp_conn_ref_take(retained) == 0);
        uint64_t local_epoch = retained->reset_id, peer_epoch = retained->peer_reset_id;
        unsigned before = endpoints[survivor].reinit;
        uint64_t packets_before = stream_packets;
        stop_endpoint(restarted); start_endpoint(restarted, reality);
        phase++; wait_pair();
        CHECK(!endpoints[survivor].deleted && endpoints[survivor].conn == retained);
        CHECK(retained->reset_id == local_epoch && retained->peer_reset_id != peer_epoch);
        CHECK(endpoints[survivor].reinit == before + 1);
        CHECK(node_conn_direct_get_conn(endpoints[survivor].owner) == retained);
#ifdef __linux__
        CHECK(stream_packets - packets_before < 256);
#endif
        DEBUG_INFO(DEBUG_CATEGORY_DEBUG, "Recovery complete: stream=%llu control=%llu REINIT=1 survivor_epoch=%016llx",
            (unsigned long long)(stream_packets - packets_before), (unsigned long long)session_packets,
            (unsigned long long)local_epoch);
        etcp_conn_ref_free(retained);
        check_unique(0); check_unique(1); exchange_messages();
    } else {
        uint64_t old_epochs[2] = {endpoints[0].conn->reset_id, endpoints[1].conn->reset_id};
        stop_endpoint(0); stop_endpoint(1); start_endpoint(0, reality); start_endpoint(1, reality);
        peers[0] = import_peer(0); peers[1] = import_peer(1);
        for (int i = 0; i < 2; i++) {
            struct TOPO_GROUP_NODE node = {0}; node.node_id = peers[i]->node_id;
            CHECK(etcp_connect(endpoints[i].inst, &node, connect_done, &endpoints[i], ETCP_CONNECT_EARLY) == 0);
        }
        phase++; wait_pair();
        for (int i = 0; i < 2; i++) { CHECK(endpoints[i].conn->reset_id != old_epochs[i]); check_unique(i); }
        exchange_messages();
    }
    stop_endpoint(0); stop_endpoint(1);
    uasync_poll(ua, 0);
    CHECK(ua->timer_alloc_count == ua->timer_free_count && ua->socket_alloc_count == ua->socket_free_count);
    uasync_destroy(ua, 0);
    printf("PASS transport reality=%d simultaneous=%d restart=%d: epochs, bidirectional delivery and resource balance\n",
        reality, simultaneous, restarted);
}

int main(void) {
    debug_config_init(); debug_set_level(DEBUG_LEVEL_WARN);
    if (getenv("UTUN_TEST_DEBUG")) {
        debug_set_category_level(DEBUG_CATEGORY_ETCP, DEBUG_LEVEL_DEBUG);
        debug_set_category_level(DEBUG_CATEGORY_GENERAL, DEBUG_LEVEL_INFO);
    }
    test_clock_init(100); utun_instance_set_tun_init_enabled(0);
    for (int reality = 0; reality < 2; reality++) {
        for (int simultaneous = 0; simultaneous < 2; simultaneous++)
            for (int restarted = 0; restarted < 2; restarted++) run_case(reality, simultaneous, restarted);
        run_case(reality, 1, 2);
    }
    test_clock_finish();
    return 0;
}