#include <assert.h>
#include <string.h>
#include "utun_instance.h"
#include "topo_group.h"
#include "topo_recovery.h"
#include "etcp.h"
#include "etcp_api.h"
#include "node_conn_direct.h"
#include "../lib/mem.h"
#include "../lib/debug_config.h"
#include "../lib/platform_compat.h"

struct fixture {
    struct UASYNC* ua;
    struct UTUN_INSTANCE* inst;
    struct TOPO_GROUP* group;
    uint64_t ids[3];
    struct ETCP_CONN* conns[3];
};

static struct TOPO_GROUP_CONN_ITEM* peer(struct fixture* f, int index) {
    for (struct ll_entry* e = f->group->senders_list->head; e; e = e->next) {
        struct TOPO_GROUP_CONN_ITEM* p = (struct TOPO_GROUP_CONN_ITEM*)e->data;
        if (p->node_id == f->ids[index]) return p;
    }
    return NULL;
}

static void poll_events(struct fixture* f) { for (int i = 0; i < 4; i++) uasync_poll(f->ua, 0); }

static void create(struct fixture* f) {
    memset(f, 0, sizeof(*f));
    f->ua = uasync_create(); assert(f->ua);
    f->inst = utun_instance_create_from_str(f->ua,
        "[global]\n"
        "my_private_key=704f2e012c8fa8768130cb0f988a997dccb628372bc5ceccacc78dcbfec5916f\n"
        "my_public_key=b3193173def895bd0fcea6f86af077c7d77216f10395275f627ac18242ec0f01\n"
        "[server: udp]\naddr=127.0.0.1:0\ntype=public\n");
    assert(f->inst && utun_instance_init(f->inst) == 0);
    f->group = topo_groups_create_group(f->inst->topo_groups, 42, TOPO_GROUP_TYPE_UTUN, NULL); assert(f->group);
    for (int i = 0; i < 3; i++) {
        struct SC_MYKEYS keys; assert(sc_generate_keypair(&keys) == SC_OK);
        struct TOPO_NODE* node = u_calloc(1, sizeof(*node)); assert(node);
        memcpy(node->public_key, keys.public_key, SC_PUBKEY_SIZE);
        f->ids[i] = node->node_id = sc_derive_node_id_from_pubkey(node->public_key);
        assert(topo_node_registry_store(f->inst->topo_groups, node));
        struct ETCP_CONN* conn = f->conns[i] = etcp_connection_create(f->inst, "recovery_fixture"); assert(conn);
        conn->peer_node_id = node->node_id;
        queue_set_callback(conn->send_input_q, NULL, NULL);
        etcp_conn_ready(conn);
    }
}

static void destroy(struct fixture* f) {
    for (int i = 0; i < 3; i++) topo_node_registry_unref(f->inst->topo_groups, f->ids[i]);
    f->inst->running = 0; utun_instance_destroy(f->inst);
    uasync_poll(f->ua, 0); uasync_destroy(f->ua, 0);
}

static void up(struct fixture* f, int i) {
    f->conns[i]->links_up = 1;
    etcp_cbk_fire(f->conns[i], ETCP_CBK_EVENT_UP);
    assert(peer(f, i) && peer(f, i)->conn == f->conns[i]);
}

static void route(struct fixture* f, int target, int via) {
    struct ll_entry* entry = queue_entry_new(sizeof(struct TOPO_GROUP_NODE) - sizeof(struct ll_entry)); assert(entry);
    struct TOPO_GROUP_NODE* node = (struct TOPO_GROUP_NODE*)entry;
    node->node_id = f->ids[target];
    topo_node_registry_ref(f->inst->topo_groups, node->node_id);
    uint64_t hops[] = { node->node_id, f->ids[via] };
    assert(topo_group_add_path(node, f->conns[via], hops, target == via ? 1 : 2, 10) == 0);
    assert(queue_data_put_with_index(f->group->nodes, entry) == 0);
    topo_recovery_changed(f->group);
}

static void table(struct fixture* f, int i, uint8_t subcmd, uint64_t id) {
    struct TOPOMSG_HEADER msg = { .cmd = ETCP_ID_TOPO_ENTRY, .subcmd = subcmd, .group_id = f->group->group_id,
                                 .src_epoch = 7, .dst_epoch = id };
    struct ll_entry* entry = ll_alloc_lldgram(sizeof(msg)); assert(entry);
    memcpy(entry->dgram, &msg, sizeof(msg)); entry->len = sizeof(msg);
    f->inst->api_bindings.callbacks[ETCP_ID_TOPO_ENTRY](f->conns[i], entry);
}

static void accept_join(struct fixture* f, int i) {
    struct TOPOMSG_JOIN_GROUP msg = { .h = { .cmd = ETCP_ID_TOPO_ENTRY, .subcmd = TOPO_SUBCMD_JOIN_ACCEPT,
        .group_id = f->group->group_id, .src_epoch = 7, .dst_epoch = peer(f, i)->local_epoch }, .group_type = TOPO_GROUP_TYPE_UTUN };
    struct ll_entry* e = ll_alloc_lldgram(sizeof(msg)); assert(e);
    memcpy(e->dgram, &msg, sizeof(msg)); e->len = sizeof(msg);
    f->inst->api_bindings.callbacks[ETCP_ID_TOPO_ENTRY](f->conns[i], e);
    assert(peer(f, i)->accepted);
}

static void ready(struct fixture* f, int i) {
    assert(peer(f, i) && peer(f, i)->local_epoch);
    accept_join(f, i);
    table(f, i, TOPO_SUBCMD_REQUEST_TABLE, peer(f, i)->local_epoch);
    table(f, i, TOPO_SUBCMD_TABLE_COMPLETE, peer(f, i)->local_epoch);
    assert(topo_group_peer_ready(f->group, f->ids[i]));
}

static void partial_and_external_routes(void) {
    struct fixture f; create(&f);
    topo_recovery_add_node(f.group, f.ids[0], f.ids[0], 100);
    topo_recovery_add_node(f.group, f.ids[1], f.ids[0], 1);
    struct TOPO_RECOVERY_CTX* context = f.group->recovery;
    topo_recovery_add_node(f.group, f.ids[1], f.ids[2], 1);
    assert(context == f.group->recovery);
    assert(topo_node_registry_find(f.inst->topo_groups, f.ids[1])->group_ref_count == 2);
    topo_recovery_start(f.group); poll_events(&f);
    assert(peer(&f, 0) && !peer(&f, 1) && !peer(&f, 2)); /* next-hop beats smaller RTT */
    topo_recovery_add_node(f.group, f.ids[2], f.ids[2], 200);
    topo_recovery_start(f.group); poll_events(&f);
    assert(f.group->recovery == context && peer(&f, 0) && !peer(&f, 2)); /* merge a new loss into the active round */
    struct NODE_CONN_DIRECT* ownership = peer(&f, 0)->handle;
    up(&f, 0); route(&f, 0, 0); poll_events(&f);
    assert(f.group->recovery == context && !peer(&f, 2)); /* UP and a route are insufficient */
    assert(peer(&f, 0)->handle == ownership); /* no handoff or second NCD handle on UP */
    ready(&f, 0); poll_events(&f);
    assert(f.group->recovery == context && peer(&f, 2) && !peer(&f, 1));
    assert(peer(&f, 0)->handle == ownership && topo_group_peer_ready(f.group, f.ids[0]));
    route(&f, 1, 0); poll_events(&f); /* another READY peer restores the remaining subtree */
    up(&f, 2); route(&f, 2, 2);
    table(&f, 2, TOPO_SUBCMD_TABLE_COMPLETE, peer(&f, 0)->local_epoch); poll_events(&f);
    assert(f.group->recovery && !topo_group_peer_ready(f.group, f.ids[2]));
    ready(&f, 2); poll_events(&f);
    assert(!f.group->recovery && !peer(&f, 1));
    assert(topo_group_peer_ready(f.group, f.ids[0]) && topo_group_peer_ready(f.group, f.ids[2]));
    destroy(&f);
}

static void stalled_and_exhausted(void) {
    struct fixture f; create(&f);
    topo_recovery_add_node(f.group, f.ids[0], f.ids[0], 10);
    topo_recovery_add_node(f.group, f.ids[1], f.ids[0], 20);
    topo_recovery_start(f.group); poll_events(&f);
    up(&f, 0);
    accept_join(&f, 0);
    uint64_t old_progress = get_time_tb() - TOPO_RECOVERY_SYNC_TIMEOUT_MS * 10ULL - 1;
    peer(&f, 0)->progress = old_progress;
    /* Реальный шаг обмена обновляет срок, пустая повторная проверка — нет. */
    table(&f, 0, TOPO_SUBCMD_REQUEST_TABLE, peer(&f, 0)->local_epoch); poll_events(&f);
    assert(peer(&f, 0) && !peer(&f, 1) && peer(&f, 0)->progress > old_progress);
    peer(&f, 0)->progress = old_progress;
    topo_recovery_changed(f.group); poll_events(&f);
    assert(!peer(&f, 0) && peer(&f, 1));
    up(&f, 1); ready(&f, 1); poll_events(&f); /* READY without routes cannot claim success */
    assert(!f.group->recovery && !topo_node_find_by_id(f.group, f.ids[0]) && !topo_node_find_by_id(f.group, f.ids[1]));
    topo_recovery_changed(f.group); topo_recovery_start(f.group); poll_events(&f);
    assert(!f.group->recovery && !peer(&f, 0)); /* exhausted round does not restart itself */
    destroy(&f);
}

static void cancel_and_group_stop(void) {
    struct fixture f; create(&f);
    struct TOPO_PEER_REQUEST* other = NULL;
    assert(topo_group_peer_open(f.group, f.ids[0], &other) == 0);
    struct NODE_CONN_DIRECT* ownership = peer(&f, 0)->handle;
    topo_recovery_add_node(f.group, f.ids[0], f.ids[0], 10);
    topo_recovery_start(f.group); poll_events(&f);
    assert(peer(&f, 0)->handle == ownership);
    topo_recovery_cancel_all(f.group); poll_events(&f);
    assert(!f.group->recovery && peer(&f, 0)->handle == ownership);
    assert(topo_group_peer_phase(other) == TOPO_PEER_CONNECTING);
    topo_recovery_add_node(f.group, f.ids[1], f.ids[1], 10);
    topo_recovery_start(f.group); poll_events(&f);
    topo_groups_remove_group(f.inst->topo_groups, f.group->group_id); f.group = NULL;
    assert(topo_group_peer_phase(other) == TOPO_PEER_FAILED);
    topo_group_peer_close(other); poll_events(&f);
    destroy(&f);
}

static void restored_during_connect(void) {
    struct fixture f; create(&f);
    f.conns[1]->links_up = 1;
    assert(topo_group_new_conn(f.group, f.conns[1]) == 0);
    route(&f, 1, 1); ready(&f, 1);
    topo_recovery_add_node(f.group, f.ids[0], f.ids[0], 10);
    topo_recovery_start(f.group); poll_events(&f);
    assert(peer(&f, 0) && !peer(&f, 0)->conn);
    route(&f, 0, 1); poll_events(&f);
    assert(!f.group->recovery && !peer(&f, 0) && topo_group_peer_ready(f.group, f.ids[1]));
    destroy(&f);
}

static void connect_deadline(void) {
    struct fixture f; create(&f);
    f.inst->etcp_connect_timeout_tb = 100000; /* recovery must time out before NCD's own 10-second timer */
    topo_recovery_add_node(f.group, f.ids[0], f.ids[0], 10);
    topo_recovery_start(f.group); poll_events(&f);
    assert(peer(&f, 0));
    uint64_t deadline = get_time_tb() + (TOPO_RECOVERY_CONNECT_TIMEOUT_MS + 1000) * 10ULL;
    while (f.group->recovery && get_time_tb() < deadline) uasync_poll(f.ua, 100);
    assert(!f.group->recovery && !peer(&f, 0));
    destroy(&f);
}

int main(void) {
    debug_config_init(); debug_set_level(DEBUG_LEVEL_INFO);
    utun_instance_set_tun_init_enabled(0);
    partial_and_external_routes();
    stalled_and_exhausted();
    cancel_and_group_stop();
    restored_during_connect();
    connect_deadline();
    DEBUG_INFO(DEBUG_CATEGORY_BGP, "recovery tests: partial routes, readiness, progress, exhaustion and shared cancellation passed");
    return 0;
}