/*
 * test_ha_vep_teardown.c — bug-1197: deferred vote-ep teardown.
 *
 * When g_drive_async_peerctrl_dispatch_enabled is ON, raft_cb_become_follower
 * runs on a vtask pool worker and must endpoint_fini() every loser vote-ep while
 * another pool worker may be mid-vtask (or have a queued vtask) on one of those
 * same eps reading ep->data → use-after-free.  The fix defers each vote-ep
 * teardown onto the SAME per-endpoint FIFO the pool already enforces (so the
 * free runs strictly after the last live use) and gates late dispatch on a
 * tearing_down flag.
 *
 * Tests (all run with ASAN; the runner uses detect_leaks=0 for the pre-existing
 * suite-wide leak, but these tests are use-after-free / double-free clean):
 *   (a) ORDERING  — teardown runs strictly after an in-flight same-ep vtask.
 *   (b) LATE-DROP — a late same-batch event on a tearing-down ep is dropped,
 *                   its packet freed, no new task enqueued.
 *   (c) DISARM    — a teardown vtask run through the real pool is ASAN-clean
 *                   (iod->ep = NULL prevents run_task touching the freed lock).
 *   (d) OFF-PATH  — with the flag OFF the loser loop uses plain endpoint_fini
 *                   and never touches the pool (behavior unchanged).
 */
#define _GNU_SOURCE
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#include <pthread.h>

#include <unity/unity.h>
#include "ha.h"
#include "netio.h"
#include "vtaskmgr.h"

/* ── XCE_TESTING seams declared in ha.c ──────────────────────────────────── */
extern int  peer_ctrl_msg_handler_for_test(endpoint_t *ep, io_packet_t *pkt);
extern io_packet_t *repl_catchup_done_pkt_for_test(void);
extern void raft_cb_become_follower_for_test(ha_state_t *ha, uint64_t term);

/* ── Fixtures ────────────────────────────────────────────────────────────── */

static ha_config_t g_cfg = {
    .heartbeat_interval_ms     = 3000,
    .heartbeat_timeout_ms      = 6000,
    .election_timeout_ms       = 500,
    .probe_connect_timeout_ms  = 50,
    .max_promotion_retries     = 3,
    .election_jitter_max_ms    = 0,
    .election_priority_step_ms = 500,
};

/* Sequence oracle: the in-flight vtask body and the teardown probe each grab a
 * monotonically increasing ticket; teardown's ticket must be the larger. */
static _Atomic int g_seq           = 0;
static _Atomic int g_inflight_seq  = -1;
static _Atomic int g_teardown_seq  = -1;
static _Atomic int g_teardown_done = 0;
static _Atomic int g_body_calls    = 0;
static _Atomic int g_body_seq      = -1;
static _Atomic int g_body_ran      = 0;

/* Gate the in-flight vtask body parks on (so the test controls when it ends). */
static struct {
    pthread_mutex_t mu;
    pthread_cond_t  entered_cv;
    pthread_cond_t  release_cv;
    int             open;
    _Atomic int     entered;
} g_gate;

static void gate_init(void)
{
    pthread_mutex_init(&g_gate.mu, NULL);
    pthread_cond_init(&g_gate.entered_cv, NULL);
    pthread_cond_init(&g_gate.release_cv, NULL);
    g_gate.open = 0;
    atomic_store(&g_gate.entered, 0);
}

static void gate_destroy(void)
{
    pthread_mutex_destroy(&g_gate.mu);
    pthread_cond_destroy(&g_gate.entered_cv);
    pthread_cond_destroy(&g_gate.release_cv);
}

static void gate_open(void)
{
    pthread_mutex_lock(&g_gate.mu);
    g_gate.open = 1;
    pthread_cond_broadcast(&g_gate.release_cv);
    pthread_mutex_unlock(&g_gate.mu);
}

static int gate_wait_entered(int timeout_ms)
{
    struct timespec d;
    clock_gettime(CLOCK_REALTIME, &d);
    d.tv_sec += timeout_ms / 1000;
    d.tv_nsec += (long)(timeout_ms % 1000) * 1000000L;
    if (d.tv_nsec >= 1000000000L) { d.tv_sec++; d.tv_nsec -= 1000000000L; }

    pthread_mutex_lock(&g_gate.mu);
    int rc = 0;
    while (atomic_load(&g_gate.entered) < 1) {
        rc = pthread_cond_timedwait(&g_gate.entered_cv, &g_gate.mu, &d);
        if (rc != 0) break;
    }
    int entered = atomic_load(&g_gate.entered);
    pthread_mutex_unlock(&g_gate.mu);
    return entered >= 1 ? 0 : -1;
}

/* In-flight peer-control vtask body: record a ticket, signal entry, park. */
static int inflight_vtask_body(vtask_t *task)
{
    (void)task;
    atomic_store(&g_inflight_seq, atomic_fetch_add(&g_seq, 1));

    pthread_mutex_lock(&g_gate.mu);
    atomic_fetch_add(&g_gate.entered, 1);
    pthread_cond_broadcast(&g_gate.entered_cv);
    while (!g_gate.open)
        pthread_cond_wait(&g_gate.release_cv, &g_gate.mu);
    pthread_mutex_unlock(&g_gate.mu);
    return 0;
}

/* Teardown probe (fires at the top of the teardown vtask body). */
static void teardown_probe(endpoint_t *ep)
{
    (void)ep;
    atomic_store(&g_teardown_seq, atomic_fetch_add(&g_seq, 1));
    atomic_store(&g_teardown_done, 1);
}

/* Body probe (fires at the top of the inline routing body). */
static void body_probe(endpoint_t *ep)
{
    (void)ep;
    atomic_fetch_add(&g_body_calls, 1);
}

/* Body probe variant that also records a sequence ticket (drain test). */
static void body_seq_probe(endpoint_t *ep)
{
    (void)ep;
    atomic_store(&g_body_seq, atomic_fetch_add(&g_seq, 1));
    atomic_store(&g_body_ran, 1);
    atomic_fetch_add(&g_body_calls, 1);
}

/* Predispatch probe (drain test): fires INSIDE the dispatch_inflight bracket,
 * after the tearing_down check passed but before the pool enqueue.  Signals
 * entry, then parks on the gate so the dispatch stays in-flight (counter == 1,
 * not yet enqueued) while the test drives endpoint_mark_tearing_down's drain. */
static void predispatch_park_probe(endpoint_t *ep)
{
    (void)ep;
    pthread_mutex_lock(&g_gate.mu);
    atomic_fetch_add(&g_gate.entered, 1);
    pthread_cond_broadcast(&g_gate.entered_cv);
    while (!g_gate.open)
        pthread_cond_wait(&g_gate.release_cv, &g_gate.mu);
    pthread_mutex_unlock(&g_gate.mu);
}

/* Heap ep with a heap ep->data (so endpoint_fini's FREEUP has something to free,
 * matching the production vote_ep_ctx_t allocation). */
static endpoint_t *make_ep(iomgr_t *iomgr)
{
    endpoint_t *ep = endpoint_get();
    TEST_ASSERT_NOT_NULL(ep);
    ep->iomgr = iomgr;
    ep->fd    = -1;
    ep->data  = calloc(1, 128);   /* stand-in for vote_ep_ctx_t */
    TEST_ASSERT_NOT_NULL(ep->data);
    return ep;
}

static void reset_oracles(void)
{
    atomic_store(&g_seq, 0);
    atomic_store(&g_inflight_seq, -1);
    atomic_store(&g_teardown_seq, -1);
    atomic_store(&g_teardown_done, 0);
    atomic_store(&g_body_calls, 0);
    atomic_store(&g_body_seq, -1);
    atomic_store(&g_body_ran, 0);
}

/* Enqueue an arbitrary VTASK_IO body bound to ep (drives per-ep serialization). */
static void enqueue_io_vtask(vtaskmgr_t *pool, endpoint_t *ep, vtask_func_t fn)
{
    vtask_data_t *data = calloc(1, sizeof(*data));
    vtask_t      *task = calloc(1, sizeof(*task));
    TEST_ASSERT_NOT_NULL(data);
    TEST_ASSERT_NOT_NULL(task);
    data->vtype        = VTASK_IO;
    data->u.iodata.ep  = ep;
    data->u.iodata.pkt = NULL;
    vtask_init(task, fn, data);
    TEST_ASSERT_EQUAL_INT(0, vtask_enq_classified(pool, task, VTASK_LIVENESS));
}

static int poll_until(_Atomic int *flag, int want, int timeout_ms)
{
    for (int i = 0; i < timeout_ms / 5; i++) {
        if (atomic_load(flag) == want)
            return 0;
        struct timespec ts = { .tv_sec = 0, .tv_nsec = 5 * 1000000L };
        nanosleep(&ts, NULL);
    }
    return atomic_load(flag) == want ? 0 : -1;
}

/* ── Test (a): ORDERING ──────────────────────────────────────────────────── */

static void test_vep_teardown_runs_after_inflight(void)
{
    reset_oracles();
    gate_init();

    ha_state_t ha;
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &g_cfg, 1));

    vtaskmgr_t pool;
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_init(&pool, 2));
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_start(&pool));
    ha.pool = &pool;

    iomgr_t iomgr;
    TEST_ASSERT_EQUAL_INT(0, iomgr_init(&iomgr));

    endpoint_t *ep = make_ep(&iomgr);

    ha_set_peerctrl_dispatch_for_test(1);
    ha_vep_teardown_probe = teardown_probe;

    /* Park an in-flight peer-control vtask on this ep (busy=1 while gated). */
    enqueue_io_vtask(&pool, ep, (vtask_func_t)inflight_vtask_body);
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, gate_wait_entered(2000),
        "in-flight vtask must run and park (busy=1)");

    /* Now tear the SAME ep down: marks tearing_down + enqueues the teardown
     * vtask, which is busy-blocked behind the parked in-flight vtask. */
    ha_vote_ep_teardown_for_test(&ha, ep);

    /* While the in-flight vtask is still parked, teardown must NOT have run. */
    TEST_ASSERT_TRUE_MESSAGE(endpoint_is_tearing_down(ep),
        "ep must be marked tearing_down once teardown is initiated");
    struct timespec ts = { .tv_sec = 0, .tv_nsec = 80 * 1000000L };
    nanosleep(&ts, NULL);
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, atomic_load(&g_teardown_done),
        "teardown must NOT run while an in-flight same-ep vtask is parked");

    /* Release the in-flight vtask; teardown becomes eligible and runs. */
    gate_open();
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, poll_until(&g_teardown_done, 1, 2000),
        "teardown must run once the in-flight vtask completes");

    TEST_ASSERT_GREATER_THAN_INT_MESSAGE(atomic_load(&g_inflight_seq),
        atomic_load(&g_teardown_seq),
        "teardown_seq must be strictly after inflight_seq");

    /* Join the pool (workers idle; ep already freed by the teardown body). */
    vtaskmgr_fini(&pool);
    iomgr_fini(&iomgr);
    ha_vep_teardown_probe = NULL;
    ha_set_peerctrl_dispatch_for_test(0);
    ha.pool = NULL;
    ha_fini(&ha);
    gate_destroy();
    printf("test_vep_teardown_runs_after_inflight: PASS\n");
}

/* ── Test (b): LATE-DROP ─────────────────────────────────────────────────── */

static void test_vep_late_event_dropped(void)
{
    reset_oracles();

    ha_state_t ha;
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &g_cfg, 1));

    vtaskmgr_t pool;
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_init(&pool, 2));
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_start(&pool));
    ha.pool = &pool;

    iomgr_t iomgr;
    TEST_ASSERT_EQUAL_INT(0, iomgr_init(&iomgr));

    endpoint_t *ep = make_ep(&iomgr);

    ha_set_peerctrl_dispatch_for_test(1);
    ha_peerctrl_body_probe = body_probe;

    /* mark + unmonitor (the teardown caller's first two steps) — but do NOT
     * enqueue a teardown vtask, so the ep stays alive for this call. */
    endpoint_mark_tearing_down(ep);
    iomgr_unmonitor(ep->iomgr, ep);

    io_packet_t *pkt = repl_catchup_done_pkt_for_test();
    TEST_ASSERT_NOT_NULL(pkt);

    int rc = peer_ctrl_msg_handler_for_test(ep, pkt);  /* takes ownership of pkt */
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc,
        "late event on a tearing-down ep must return 0 (dropped)");

    /* Give any (erroneously) enqueued task time to run the body. */
    struct timespec ts = { .tv_sec = 0, .tv_nsec = 100 * 1000000L };
    nanosleep(&ts, NULL);
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, atomic_load(&g_body_calls),
        "a dropped late event must NOT enqueue/run the routing body");

    ha_peerctrl_body_probe = NULL;
    ha_set_peerctrl_dispatch_for_test(0);

    /* ep was not torn down by a vtask here (late event dropped, refcnt==1) —
     * drop the owner ref ourselves; endpoint_unref fini's + frees it (pkt was
     * freed by the drop path). bug-1298: use unref, not raw fini+free. */
    endpoint_unref(ep);

    vtaskmgr_fini(&pool);
    iomgr_fini(&iomgr);
    ha.pool = NULL;
    ha_fini(&ha);
    printf("test_vep_late_event_dropped: PASS\n");
}

/* ── Test (c): DISARM ────────────────────────────────────────────────────── */

static void test_vep_teardown_vtask_disarms_unblock(void)
{
    reset_oracles();

    ha_state_t ha;
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &g_cfg, 1));

    vtaskmgr_t pool;
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_init(&pool, 2));
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_start(&pool));
    ha.pool = &pool;

    iomgr_t iomgr;
    TEST_ASSERT_EQUAL_INT(0, iomgr_init(&iomgr));

    endpoint_t *ep = make_ep(&iomgr);

    ha_set_peerctrl_dispatch_for_test(1);
    ha_vep_teardown_probe = teardown_probe;

    /* Run a teardown vtask straight through the pool (no in-flight task). If the
     * iod->ep = NULL disarm were missing, run_task would call
     * endpoint_unblock_tasks on the freed ep->lock — ASAN would flag a UAF. */
    ha_vote_ep_teardown_for_test(&ha, ep);
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, poll_until(&g_teardown_done, 1, 2000),
        "teardown vtask must run through the pool");

    vtaskmgr_fini(&pool);
    iomgr_fini(&iomgr);
    ha_vep_teardown_probe = NULL;
    ha_set_peerctrl_dispatch_for_test(0);
    ha.pool = NULL;
    ha_fini(&ha);
    printf("test_vep_teardown_vtask_disarms_unblock: PASS\n");
}

/* ── Test (d): OFF-PATH ──────────────────────────────────────────────────── */

static void test_vep_off_path_uses_plain_fini(void)
{
    reset_oracles();

    ha_state_t ha;
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &g_cfg, 1));
    ha.pool = NULL;   /* off path must not need the pool */

    ha_set_peerctrl_dispatch_for_test(0);   /* flag OFF */
    ha_vep_teardown_probe = teardown_probe;

    endpoint_t *ep = endpoint_get();
    TEST_ASSERT_NOT_NULL(ep);
    ep->iomgr = NULL;          /* OFF path: endpoint_fini only, no unmonitor */
    ep->fd    = -1;
    ep->data  = calloc(1, 128);
    TEST_ASSERT_NOT_NULL(ep->data);

    /* One-peer election context with this ep as the loser vote-ep. */
    ha_elect_ctx_t *ectx = calloc(1, sizeof(*ectx) + sizeof(ha_vote_entry_t));
    TEST_ASSERT_NOT_NULL(ectx);
    ectx->n_peers      = 1;
    ectx->votes[0].ep  = ep;
    atomic_store_explicit(&ha.elect_ctx, ectx, memory_order_release);

    raft_cb_become_follower_for_test(&ha, 5);

    TEST_ASSERT_EQUAL_INT_MESSAGE(0, atomic_load(&g_teardown_done),
        "OFF path must NOT run the async teardown vtask");
    TEST_ASSERT_NULL_MESSAGE(ectx->votes[0].ep,
        "loser loop must clear votes[i].ep after endpoint_unref");

    /* bug-1298: the OFF path now drops the owner ref via endpoint_unref — with no
     * pool vtask referencing the ep (refcnt==1) that fini's AND frees it here,
     * fixing the legacy per-lost-election struct leak.  So the test must NOT
     * free(ep) again (that was the old leak-cleanup; it is now a double-free). */
    atomic_store_explicit(&ha.elect_ctx, NULL, memory_order_release);
    free(ectx);

    /* become_follower armed the heartbeat watchdog (c=NULL) — release its ctx. */
    ha_disarm_heartbeat_timeout_vtask(&ha);

    ha_vep_teardown_probe = NULL;
    ha_fini(&ha);
    printf("test_vep_off_path_uses_plain_fini: PASS\n");
}

/* ── Test (e): DRAIN — mark blocks until an in-flight dispatch finishes ───────
 *
 * Closes the TOCTOU window: a dispatch that read tearing_down as false but is
 * preempted before its enqueue must NOT be able to enqueue after the teardown
 * vtask.  endpoint_mark_tearing_down drains ep->dispatch_inflight to 0 (after
 * setting the flag) before returning, so the teardown caller enqueues the
 * teardown vtask only after every such in-flight dispatch has finished enqueuing
 * → the peer-control vtask is strictly FIFO-before the teardown vtask.
 *
 * Deterministic: a predispatch probe parks thread A inside the dispatch_inflight
 * bracket (counter == 1, not yet enqueued).  Thread B runs the teardown seam; we
 * assert its mark/drain does NOT return while A is parked, then release A and
 * assert mark returns and the peer-control body ran before the teardown probe. */

typedef struct { endpoint_t *ep; _Atomic int done; } dispatch_arg_t;
typedef struct { ha_state_t *ha; endpoint_t *ep; _Atomic int done; } teardown_arg_t;

/* Layout-compatible stand-in for vote_ep_ctx_t: the inline routing body reads
 * only the first member (ha) before the default-drop path frees the packet. */
struct mini_vctx { ha_state_t *ha; char pad[160]; };

static void *drain_dispatch_thread(void *p)
{
    dispatch_arg_t *a = (dispatch_arg_t *)p;
    io_packet_t    *pkt = repl_catchup_done_pkt_for_test();
    peer_ctrl_msg_handler_for_test(a->ep, pkt);  /* parks in predispatch probe */
    atomic_store(&a->done, 1);
    return NULL;
}

static void *drain_teardown_thread(void *p)
{
    teardown_arg_t *b = (teardown_arg_t *)p;
    ha_vote_ep_teardown_for_test(b->ha, b->ep);  /* mark drains, then enqueues */
    atomic_store(&b->done, 1);
    return NULL;
}

static void test_vep_mark_drains_inflight_dispatch(void)
{
    reset_oracles();
    gate_init();

    ha_state_t ha;
    TEST_ASSERT_EQUAL_INT(0, ha_init(&ha, &g_cfg, 1));
    /* CTRL_INIT → the inline body hits the default case (drop + free pkt) so it
     * never decodes, keeping the stand-in vote_ep_ctx_t sufficient. */
    atomic_store_explicit(&ha.role, (int32_t)CTRL_INIT, memory_order_release);

    vtaskmgr_t pool;
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_init(&pool, 2));
    TEST_ASSERT_EQUAL_INT(0, vtaskmgr_start(&pool));
    ha.pool = &pool;

    iomgr_t iomgr;
    TEST_ASSERT_EQUAL_INT(0, iomgr_init(&iomgr));

    endpoint_t *ep = endpoint_get();
    TEST_ASSERT_NOT_NULL(ep);
    ep->iomgr = &iomgr;
    ep->fd    = -1;
    struct mini_vctx *v = calloc(1, sizeof(*v));
    TEST_ASSERT_NOT_NULL(v);
    v->ha    = &ha;
    ep->data = v;

    ha_set_peerctrl_dispatch_for_test(1);
    ha_peerctrl_predispatch_probe = predispatch_park_probe;
    ha_peerctrl_body_probe        = body_seq_probe;
    ha_vep_teardown_probe         = teardown_probe;

    /* Thread A: dispatch parks in the drained window (inflight == 1). */
    dispatch_arg_t aarg = { .ep = ep };
    atomic_store(&aarg.done, 0);
    pthread_t ta;
    TEST_ASSERT_EQUAL_INT(0, pthread_create(&ta, NULL, drain_dispatch_thread, &aarg));
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, gate_wait_entered(2000),
        "dispatch must reach the predispatch window and park (inflight=1)");

    /* Thread B: teardown mark must BLOCK draining on A. */
    teardown_arg_t barg = { .ha = &ha, .ep = ep };
    atomic_store(&barg.done, 0);
    pthread_t tb;
    TEST_ASSERT_EQUAL_INT(0, pthread_create(&tb, NULL, drain_teardown_thread, &barg));

    TEST_ASSERT_EQUAL_INT_MESSAGE(-1, poll_until(&barg.done, 1, 150),
        "endpoint_mark_tearing_down must block while a dispatch is in-flight");
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, atomic_load(&g_teardown_done),
        "teardown vtask must not be enqueued/run while mark is still draining");

    /* Release A → it completes its enqueue + decrements → B's drain ends. */
    gate_open();
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, poll_until(&barg.done, 1, 2000),
        "mark must return once the in-flight dispatch finishes its enqueue");

    pthread_join(ta, NULL);
    pthread_join(tb, NULL);

    /* Both vtasks run on the pool; peer-control is FIFO-before teardown. */
    TEST_ASSERT_EQUAL_INT_MESSAGE(0, poll_until(&g_teardown_done, 1, 2000),
        "teardown vtask must run");
    TEST_ASSERT_EQUAL_INT_MESSAGE(1, atomic_load(&g_body_ran),
        "the in-flight peer-control body must have run");
    TEST_ASSERT_GREATER_THAN_INT_MESSAGE(atomic_load(&g_body_seq),
        atomic_load(&g_teardown_seq),
        "peer-control vtask must be FIFO-before the teardown vtask");

    vtaskmgr_fini(&pool);
    iomgr_fini(&iomgr);
    ha_peerctrl_predispatch_probe = NULL;
    ha_peerctrl_body_probe        = NULL;
    ha_vep_teardown_probe         = NULL;
    ha_set_peerctrl_dispatch_for_test(0);
    ha.pool = NULL;
    ha_fini(&ha);
    gate_destroy();
    printf("test_vep_mark_drains_inflight_dispatch: PASS\n");
}

/* ── Entry point ─────────────────────────────────────────────────────────── */

void test_ha_vep_teardown(void)
{
    RUN_TEST(test_vep_teardown_runs_after_inflight);
    RUN_TEST(test_vep_late_event_dropped);
    RUN_TEST(test_vep_teardown_vtask_disarms_unblock);
    RUN_TEST(test_vep_off_path_uses_plain_fini);
    RUN_TEST(test_vep_mark_drains_inflight_dispatch);
}