/*
 * test_ha_connection_reuse.c — Unity tests for ha-peer-connection-reuse spec.
 *
 * Tests verify:
 *   - peer_ctrl_msg_handler routes by role at dispatch time
 *   - raft_cb_become_leader promotes vote endpoints to standby_endpoints
 *   - raft_cb_become_follower closes vote endpoints and arms watchdog
 *   - handle_ctrlhello sets primary_ctrl_fd on first ctrlhello (CTRL_STANDBY)
 *   - vote_ep_err_handler removes promoted endpoint from standby_endpoints
 *
 * Requirements: 1.2, 1.3, 1.4, 2.1, 2.3, 3.1, 3.2, 3.3, 4.1, 4.2, 5.5
 */
#define _GNU_SOURCE
#include <errno.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <unity/unity.h>

#include "controller.capnp.h"
#include "ha.h"
#include "netio.h"
#include "xce/cluster_def.h"

/* ── Test seams exposed from ha.c under XCE_TESTING ────────────────────── */

/* Vote reply hook (xce_vote_reply_hook): fires instead of raft_on_vote_reply */
extern int (*xce_vote_reply_hook)(ha_state_t *ha, ha_elect_ctx_t *ctx, int peer_idx,
                                  uint64_t term, int vote_granted,
                                  uint64_t peer_tag, uint64_t peer_lsn);

/* peer_ctrl_msg_handler_for_test: sets ep->fsm_in.pkt then calls the handler */
extern int peer_ctrl_msg_handler_for_test(endpoint_t *ep, io_packet_t *pkt);

/* vote_ep_err_handler_for_test: call the static vote_ep_err_handler directly */
extern int vote_ep_err_handler_for_test(endpoint_t *ep);

/* vote_reply_build_pkt_for_test: encode a RequestVoteReply capnp packet */
extern io_packet_t *vote_reply_build_pkt_for_test(uint64_t term, int vote_granted,
                                                   uint64_t peer_tag, uint64_t peer_lsn);

/* repl_catchup_done_pkt_for_test: encode a replCatchUpDone capnp packet */
extern io_packet_t *repl_catchup_done_pkt_for_test(void);

/* vote_ep_ctx_build_for_test: allocate a vote_ep_ctx_t and wire ep->data */
extern endpoint_t *vote_ep_ctx_build_for_test(endpoint_t    *ep,
                                               ha_state_t    *ha,
                                               ha_elect_ctx_t *ctx,
                                               int            peer_idx,
                                               uint64_t       my_term,
                                               uint64_t       my_tag,
                                               uint64_t       my_lsn,
                                               const char    *my_name);

/* raft callbacks wrapped for test (ha->raft_c set to NULL to suppress production deps) */
extern void raft_cb_become_leader_for_test(ha_state_t *ha, uint64_t term);
extern void raft_cb_become_follower_for_test(ha_state_t *ha, uint64_t term);

/* The replication-ack promotion handler was DELETED in production
 * (ha-single-leader-role task 4.1: the ack-counting force-promotion path is gone —
 * promotion is driven by the heartbeat decision loop + per-entry majority ack).
 * The ack-counting cases that depended on it were deleted; no shim. */

/* ha_phase2_hook: suppress cluster_init_phase2 in callbacks */
extern int (*ha_phase2_hook)(void *c);
extern int (*ha_push_config_hook)(void *c);
extern int (*ha_send_ctrlhello_hook)(void *c);

/* ── Minimal config ─────────────────────────────────────────────────────── */

static ha_config_t g_cfg = {
    .heartbeat_interval_ms     = 3000,
    .heartbeat_timeout_ms      = 10000,
    .election_timeout_ms       = 2000,
    .probe_connect_timeout_ms  = 500,
    .max_promotion_retries     = 5,
    .election_jitter_max_ms    = 0,
    .election_priority_step_ms = 500,
};

/* ── Stubs for production hooks ─────────────────────────────────────────── */

static int g_phase2_calls = 0;
static int g_push_config_calls = 0;
static int g_ctrlhello_calls = 0;
static int g_vote_reply_calls = 0;

static int stub_phase2(void *c)       { (void)c; g_phase2_calls++; return 0; }
static int stub_push_config(void *c)  { (void)c; g_push_config_calls++; return 0; }
static int stub_ctrlhello(void *c)    { (void)c; g_ctrlhello_calls++; return 0; }
static int stub_vote_reply_hook(ha_state_t *ha, ha_elect_ctx_t *ctx, int peer_idx,
                                 uint64_t term, int vote_granted,
                                 uint64_t peer_tag, uint64_t peer_lsn)
{
    (void)ha; (void)ctx; (void)peer_idx; (void)term;
    (void)vote_granted; (void)peer_tag; (void)peer_lsn;
    g_vote_reply_calls++;
    return 0;
}

/* ── Helper: build ha with standby_endpoints array ─────────────────────── */

static void init_ha_with_standby_eps(ha_state_t *ha)
{
    ha_init(ha, &g_cfg, 0);
    ha->standby_endpoints = calloc(1, sizeof(zarray_t));
    zarray_init(ha->standby_endpoints, 4, NULL);
}

static void fini_ha_with_standby_eps(ha_state_t *ha)
{
    if (ha->standby_endpoints) {
        zarray_fini(ha->standby_endpoints);
        free(ha->standby_endpoints);
        ha->standby_endpoints = NULL;
    }
    ha_fini(ha);
}

/* ── Helper: build a minimal ha_elect_ctx_t with n entries ─────────────── */

static ha_elect_ctx_t *alloc_elect_ctx(int n_peers)
{
    ha_elect_ctx_t *ctx = calloc(1, sizeof(*ctx) +
                                 (size_t)(n_peers > 0 ? n_peers : 1) *
                                 sizeof(ctx->votes[0]));
    if (!ctx) return NULL;
    ctx->n_peers = n_peers;
    pthread_mutex_init(&ctx->lock, NULL);
    sem_init(&ctx->done, 0, 0);
    return ctx;
}

static void free_elect_ctx(ha_elect_ctx_t *ctx)
{
    if (!ctx) return;
    pthread_mutex_destroy(&ctx->lock);
    sem_destroy(&ctx->done);
    free(ctx);
}

/* ── Test 1: dispatcher routes CTRL_ELECTING → vote_reply_decode ────────── */

static void test_dispatcher_electing_calls_vote_reply_hook(void)
{
    ha_state_t ha;
    ha_init(&ha, &g_cfg, 0);
    atomic_store_explicit(&ha.role, (int32_t)CTRL_ELECTING, memory_order_release);

    ha_elect_ctx_t *ctx = alloc_elect_ctx(1);
    TEST_ASSERT_NOT_NULL_MESSAGE(ctx, "alloc_elect_ctx must succeed");

    endpoint_t ep;
    memset(&ep, 0, sizeof(ep));
    vote_ep_ctx_build_for_test(&ep, &ha, ctx, 0, 1, 0, 0, "test");

    g_vote_reply_calls = 0;
    xce_vote_reply_hook = stub_vote_reply_hook;

    io_packet_t *pkt = vote_reply_build_pkt_for_test(1, 1, 0, 0);
    TEST_ASSERT_NOT_NULL_MESSAGE(pkt, "vote_reply_build_pkt_for_test must succeed");

    int rc = peer_ctrl_msg_handler_for_test(&ep, pkt);
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "handler must return 0");
    TEST_ASSERT_EQUAL_INT_MESSAGE(1, g_vote_reply_calls,
        "CTRL_ELECTING: peer_ctrl_msg_handler must invoke xce_vote_reply_hook (Req 1.2)");

    xce_vote_reply_hook = NULL;
    free(ep.data);
    free_elect_ctx(ctx);
    ha_fini(&ha);
    printf("test_dispatcher_electing_calls_vote_reply_hook: PASS\n");
}

/* ── Test 2 (dispatcher transitional-leader ack count) DELETED ─────────
 * ha-single-leader-role task 4.1.  Subject was the deleted ack-counting
 * promotion machinery (replCatchUpDone incrementing the per-term ACK tally to
 * drive a transitional leader state to PRIMARY).  That path is gone — nothing
 * left to verify.  No shim (design §"Transitional-State Deletions").
 */

/* ── Test 3: dispatcher drops packet for CTRL_STANDBY ───────────────────── */

static void test_dispatcher_standby_drops_packet(void)
{
    ha_state_t ha;
    ha_init(&ha, &g_cfg, 0);
    atomic_store_explicit(&ha.role, (int32_t)CTRL_STANDBY, memory_order_release);

    ha_elect_ctx_t *ctx = alloc_elect_ctx(1);
    TEST_ASSERT_NOT_NULL_MESSAGE(ctx, "alloc_elect_ctx must succeed");

    endpoint_t ep;
    memset(&ep, 0, sizeof(ep));
    vote_ep_ctx_build_for_test(&ep, &ha, ctx, 0, 1, 0, 0, "test");

    g_vote_reply_calls = 0;
    xce_vote_reply_hook = stub_vote_reply_hook;

    io_packet_t *pkt = vote_reply_build_pkt_for_test(1, 1, 0, 0);
    TEST_ASSERT_NOT_NULL_MESSAGE(pkt, "vote_reply_build_pkt_for_test must succeed");

    int rc = peer_ctrl_msg_handler_for_test(&ep, pkt);
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "handler must return 0");
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, g_vote_reply_calls,
        "CTRL_STANDBY: packet must be dropped, vote hook must NOT be called (Req 1.2)");

    xce_vote_reply_hook = NULL;
    free(ep.data);
    free_elect_ctx(ctx);
    ha_fini(&ha);
    printf("test_dispatcher_standby_drops_packet: PASS\n");
}

/* ── Test 4: raft_cb_become_leader populates standby_endpoints ──────────── */

static void test_become_leader_promotes_vote_endpoints(void)
{
    ha_state_t ha;
    init_ha_with_standby_eps(&ha);

    /* Suppress production-only hooks */
    ha_phase2_hook      = stub_phase2;
    ha_push_config_hook = stub_push_config;
    ha_send_ctrlhello_hook = stub_ctrlhello;

    /* Build an elect_ctx with 2 peers, each with a mock endpoint */
    ha_elect_ctx_t *ctx = alloc_elect_ctx(2);
    TEST_ASSERT_NOT_NULL_MESSAGE(ctx, "alloc_elect_ctx must succeed");

    endpoint_t ep0, ep1;
    memset(&ep0, 0, sizeof(ep0));
    memset(&ep1, 0, sizeof(ep1));
    ctx->votes[0].ep = &ep0;
    ctx->votes[1].ep = &ep1;

    atomic_store_explicit(&ha.elect_ctx, (void *)ctx, memory_order_release);

    raft_cb_become_leader_for_test(&ha, 1);

    int count = ha.standby_endpoints ? ha.standby_endpoints->count : -1;
    TEST_ASSERT_EQUAL_INT_MESSAGE(2, count,
        "become_leader: both surviving vote eps must be in standby_endpoints (Req 1.3)");

    /* Cleanup */
    ha_phase2_hook         = NULL;
    ha_push_config_hook    = NULL;
    ha_send_ctrlhello_hook = NULL;
    atomic_store_explicit(&ha.elect_ctx, NULL, memory_order_release);
    free_elect_ctx(ctx);
    fini_ha_with_standby_eps(&ha);
    printf("test_become_leader_promotes_vote_endpoints: PASS\n");
}

/* ── Test 5: raft_cb_become_leader skips NULL vote endpoints ─────────────── */

static void test_become_leader_skips_null_endpoints(void)
{
    ha_state_t ha;
    init_ha_with_standby_eps(&ha);

    ha_phase2_hook         = stub_phase2;
    ha_push_config_hook    = stub_push_config;
    ha_send_ctrlhello_hook = stub_ctrlhello;

    ha_elect_ctx_t *ctx = alloc_elect_ctx(2);
    TEST_ASSERT_NOT_NULL_MESSAGE(ctx, "alloc_elect_ctx must succeed");

    endpoint_t ep0;
    memset(&ep0, 0, sizeof(ep0));
    ctx->votes[0].ep = &ep0; /* only first peer connected */
    ctx->votes[1].ep = NULL; /* second peer failed connect */

    atomic_store_explicit(&ha.elect_ctx, (void *)ctx, memory_order_release);

    raft_cb_become_leader_for_test(&ha, 1);

    int count = ha.standby_endpoints ? ha.standby_endpoints->count : -1;
    TEST_ASSERT_EQUAL_INT_MESSAGE(1, count,
        "become_leader: NULL vote ep must be skipped, only 1 promoted (Req 1.3)");

    ha_phase2_hook         = NULL;
    ha_push_config_hook    = NULL;
    ha_send_ctrlhello_hook = NULL;
    atomic_store_explicit(&ha.elect_ctx, NULL, memory_order_release);
    free_elect_ctx(ctx);
    fini_ha_with_standby_eps(&ha);
    printf("test_become_leader_skips_null_endpoints: PASS\n");
}

/* ── Test 6: raft_cb_become_follower closes all vote endpoints ──────────── */

static void test_become_follower_closes_vote_endpoints(void)
{
    ha_state_t ha;
    ha_init(&ha, &g_cfg, 0);

    ha_phase2_hook = stub_phase2;

    ha_elect_ctx_t *ctx = alloc_elect_ctx(2);
    TEST_ASSERT_NOT_NULL_MESSAGE(ctx, "alloc_elect_ctx must succeed");

    /* We cannot call endpoint_fini on stack-allocated endpoints since it would
     * try to close fd=-1.  Instead use heap-allocated endpoints with fd=-1
     * which endpoint_fini handles gracefully (it checks fd). */
    endpoint_t *ep0 = calloc(1, sizeof(*ep0));
    endpoint_t *ep1 = calloc(1, sizeof(*ep1));
    TEST_ASSERT_NOT_NULL(ep0);
    TEST_ASSERT_NOT_NULL(ep1);
    ep0->fd = -1;
    ep1->fd = -1;
    ctx->votes[0].ep = ep0;
    ctx->votes[1].ep = ep1;

    atomic_store_explicit(&ha.elect_ctx, (void *)ctx, memory_order_release);

    raft_cb_become_follower_for_test(&ha, 1);

    /* All eps must be NULL after follower transition */
    TEST_ASSERT_NULL_MESSAGE(ctx->votes[0].ep,
        "become_follower: votes[0].ep must be NULL after close (Req 1.4, 4.1)");
    TEST_ASSERT_NULL_MESSAGE(ctx->votes[1].ep,
        "become_follower: votes[1].ep must be NULL after close (Req 1.4, 4.1)");

    int role = (int)atomic_load_explicit(&ha.role, memory_order_acquire);
    TEST_ASSERT_EQUAL_INT_MESSAGE((int)CTRL_STANDBY, role,
        "become_follower: role must be CTRL_STANDBY after follower transition (Req 2.1)");

    ha_phase2_hook = NULL;
    atomic_store_explicit(&ha.elect_ctx, NULL, memory_order_release);
    free_elect_ctx(ctx);
    free(ep0);
    free(ep1);
    ha_fini(&ha);
    printf("test_become_follower_closes_vote_endpoints: PASS\n");
}

/* ── Test 7: vote_ep_err_handler removes ep from standby_endpoints ────────── */

static void test_err_handler_removes_promoted_ep_from_standby(void)
{
    ha_state_t ha;
    init_ha_with_standby_eps(&ha);
    ha.n_peers = 2;
    atomic_store_explicit(&ha.role, (int32_t)CTRL_PRIMARY,
                          memory_order_release);

    /* Populate standby_endpoints with a heap-allocated mock ep pointer */
    endpoint_t *ep = calloc(1, sizeof(*ep));
    TEST_ASSERT_NOT_NULL_MESSAGE(ep, "calloc ep must succeed");
    ep->fd = -1;
    zarray_store(ha.standby_endpoints, ep);

    TEST_ASSERT_EQUAL_INT_MESSAGE(1, ha.standby_endpoints->count,
        "standby_endpoints must have 1 entry before err_handler");

    /* Build a vote_ep_ctx_t wired to ep->data */
    ha_elect_ctx_t *ctx = alloc_elect_ctx(1);
    TEST_ASSERT_NOT_NULL_MESSAGE(ctx, "alloc_elect_ctx must succeed");
    vote_ep_ctx_build_for_test(ep, &ha, ctx, 0, 1, 0, 0, "test");

    /* Call vote_ep_err_handler directly via its test wrapper */
    int rc = vote_ep_err_handler_for_test(ep);
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "err_handler must return 0");

    TEST_ASSERT_EQUAL_INT_MESSAGE(0, ha.standby_endpoints->count,
        "err_handler: ep must be removed from standby_endpoints when role=CTRL_PRIMARY (Req 4.2)");

    free(ep->data);
    free(ep);
    free_elect_ctx(ctx);
    fini_ha_with_standby_eps(&ha);
    printf("test_err_handler_removes_promoted_ep_from_standby: PASS\n");
}

/* ── Test 8: handle_ctrlhello sets primary_ctrl_fd on first call ──────────── */

extern int handle_ctrlhello(cluster_t *c, endpoint_t *ep, xce_request_t *req);

static void test_handle_ctrlhello_sets_primary_ctrl_fd(void)
{
    /* Build a minimal cluster_t with role=CTRL_STANDBY, primary_ctrl_fd=-1 */
    cluster_t c;
    memset(&c, 0, sizeof(c));
    ha_init(&c.ha, &g_cfg, 0);
    c.ha.primary_ctrl_fd = -1;
    atomic_store_explicit(&c.ha.role, (int32_t)CTRL_STANDBY, memory_order_release);
    /* Provide a minimal myconf so handle_ctrlhello does not crash on NULL deref */
    myconf_t conf;
    memset(&conf, 0, sizeof(conf));
    c.myconf = &conf;

    endpoint_t ep;
    memset(&ep, 0, sizeof(ep));
    ep.fd = 42;

    /* Build a minimal CtrlHelloRequest body */
    ctrl_hello_req_t hello;
    memset(&hello, 0, sizeof(hello));
    hello.name       = "standby";
    hello.applied_lsn = 0;
    hello.config_tag  = 0;

    xce_request_t req;
    memset(&req, 0, sizeof(req));
    req.body_which    = Request_body_ctrlhellomsg;
    req.body.ctrlhellomsg = &hello;

    int rc = handle_ctrlhello(&c, &ep, &req);
    /* rc may be non-zero if repl_add_standby fails without full cluster state;
     * we only care that primary_ctrl_fd was set. */
    (void)rc;

    TEST_ASSERT_EQUAL_INT_MESSAGE(42, c.ha.primary_ctrl_fd,
        "handle_ctrlhello: primary_ctrl_fd must be set to ep->fd on first call (Req 3.1)");

    /* Second call with same ep must be a no-op (idempotency) */
    ep.fd = 99; /* different fd to prove it is not overwritten */
    handle_ctrlhello(&c, &ep, &req);
    TEST_ASSERT_EQUAL_INT_MESSAGE(42, c.ha.primary_ctrl_fd,
        "handle_ctrlhello: second call must not overwrite primary_ctrl_fd (Req 3.1)");

    ha_fini(&c.ha);
    printf("test_handle_ctrlhello_sets_primary_ctrl_fd: PASS\n");
}

/* ── Test 9 (Task 7.2): follower path is passive — no outbound connect ──────
 *
 * Req 2.1, 2.4: after raft_cb_become_follower the role is CTRL_STANDBY and
 * no fetch_config_from_primary is scheduled (verified by checking the role
 * and that votes[i].ep are cleared — the old blocking connect path is gone).
 */
static void test_follower_passive_no_outbound_connect(void)
{
    ha_state_t ha;
    ha_init(&ha, &g_cfg, 0);

    ha_phase2_hook = stub_phase2;

    ha_elect_ctx_t *ctx = alloc_elect_ctx(2);
    TEST_ASSERT_NOT_NULL_MESSAGE(ctx, "alloc_elect_ctx must succeed");

    /* Populate two heap-allocated mock endpoints */
    endpoint_t *ep0 = calloc(1, sizeof(*ep0));
    endpoint_t *ep1 = calloc(1, sizeof(*ep1));
    TEST_ASSERT_NOT_NULL(ep0);
    TEST_ASSERT_NOT_NULL(ep1);
    ep0->fd = -1;
    ep1->fd = -1;
    ctx->votes[0].ep = ep0;
    ctx->votes[1].ep = ep1;

    atomic_store_explicit(&ha.elect_ctx, (void *)ctx, memory_order_release);

    raft_cb_become_follower_for_test(&ha, 1);

    /* Role must be CTRL_STANDBY */
    TEST_ASSERT_EQUAL_INT_MESSAGE((int)CTRL_STANDBY,
        (int)atomic_load_explicit(&ha.role, memory_order_acquire),
        "Task 7.2: role must be CTRL_STANDBY after follower transition (Req 2.1)");

    /* All vote endpoints must be closed (NULL) — primary drives connection */
    TEST_ASSERT_NULL_MESSAGE(ctx->votes[0].ep,
        "Task 7.2: votes[0].ep must be NULL — standby never reconnects (Req 2.4)");
    TEST_ASSERT_NULL_MESSAGE(ctx->votes[1].ep,
        "Task 7.2: votes[1].ep must be NULL — standby never reconnects (Req 2.4)");

    ha_phase2_hook = NULL;
    atomic_store_explicit(&ha.elect_ctx, NULL, memory_order_release);
    free_elect_ctx(ctx);
    free(ep0);
    free(ep1);
    ha_fini(&ha);
    printf("test_follower_passive_no_outbound_connect: PASS\n");
}

/* ── Test 10 (Task 7.3: promoted ep ack promotes to PRIMARY) DELETED ─────
 * ha-single-leader-role task 4.1.  Subject was the deleted ack-counting
 * promotion machinery (replCatchUpDone incrementing the per-term ACK tally until
 * the majority threshold promoted a transitional leader state to CTRL_PRIMARY via
 * ha_promote_to_primary).
 * The mechanism is gone — promotion is now the heartbeat decision loop +
 * per-entry majority ack.  Nothing left to verify; no shim (design
 * §"Transitional-State Deletions").
 */

/* ── Test suite entry point ─────────────────────────────────────────────── */

void test_ha_connection_reuse(void)
{
    /* Req 1.2 — dispatcher routes by role */
    RUN_TEST(test_dispatcher_electing_calls_vote_reply_hook);
    /* test_dispatcher_committing_handles_repl_ack deleted (ha-single-leader-role
     * 4.1: ack-counting promotion machinery removed, no shim) */
    RUN_TEST(test_dispatcher_standby_drops_packet);

    /* Req 1.3 — winner promotes vote endpoints */
    RUN_TEST(test_become_leader_promotes_vote_endpoints);
    RUN_TEST(test_become_leader_skips_null_endpoints);

    /* Req 1.4, 2.1, 4.1 — loser closes vote endpoints */
    RUN_TEST(test_become_follower_closes_vote_endpoints);

    /* Req 4.2, 5.5 — promoted ep disconnect removes from standby_endpoints */
    RUN_TEST(test_err_handler_removes_promoted_ep_from_standby);

    /* Req 3.1 — handle_ctrlhello designates primary control channel */
    RUN_TEST(test_handle_ctrlhello_sets_primary_ctrl_fd);

    /* Task 7.2 — Req 2.1, 2.4: follower path is fully passive */
    RUN_TEST(test_follower_passive_no_outbound_connect);

    /* Task 7.3 (test_promoted_ep_repl_ack_promotes_to_primary) deleted
     * (ha-single-leader-role 4.1: ack-counting promotion machinery removed, no shim) */
}