#define _GNU_SOURCE
#include <errno.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdint.h>
#include <time.h>

#include <unity/unity.h>

#include "slowlane.h"

/* ====================================================================
 * raft-drive-async task 1.2 — SlowLane executor unit tests (RED first).
 *
 * The SlowLane is a bounded pool of dedicated worker threads that runs
 * long blocking I/O off the Raft drive threads, with admission control
 * and per-key in-flight dedup.  These tests assert the four core
 * behavioral contracts (acceptance criteria 3.1-3.4, 1.5):
 *   (a) slowlane_submit is non-blocking
 *   (b) concurrent RUNNING tasks never exceed max_inflight
 *   (c) a submit with an already-in-flight non-zero key returns -EEXIST
 *   (d) fn runs exactly once on a lane worker
 *   (+) -EAGAIN is returned (but work still queued) when the cap is full
 *
 * Determinism: every blocking task parks on a test-owned condvar gate
 * (NOT a sleep), so the test controls exactly when a worker is occupied
 * and when it releases.  Concurrency peaks are observed via an atomic
 * running counter the tasks themselves bump.
 * ==================================================================== */

/* ── Test-owned gate: a controllable blocking task body ───────────── */

typedef struct {
    pthread_mutex_t  mu;
    pthread_cond_t   release_cv;  /* test broadcasts to release parked tasks */
    pthread_cond_t   entered_cv;  /* task broadcasts when it has entered/parked */
    int              gate_open;   /* 0 = parked tasks must wait */

    _Atomic int      running;     /* tasks currently inside fn (post-entry) */
    _Atomic int      max_running; /* high-water mark of `running` */
    _Atomic int      total_runs;  /* how many times any task fn executed */
    _Atomic int      entered;     /* tasks that have reached the gate */
} gate_t;

static void gate_init(gate_t *g)
{
    pthread_mutex_init(&g->mu, NULL);
    pthread_cond_init(&g->release_cv, NULL);
    pthread_cond_init(&g->entered_cv, NULL);
    g->gate_open = 0;
    atomic_store(&g->running, 0);
    atomic_store(&g->max_running, 0);
    atomic_store(&g->total_runs, 0);
    atomic_store(&g->entered, 0);
}

static void gate_destroy(gate_t *g)
{
    pthread_mutex_destroy(&g->mu);
    pthread_cond_destroy(&g->release_cv);
    pthread_cond_destroy(&g->entered_cv);
}

/* Blocking task body: bump running, record the peak, signal entry, then park
 * on the gate until the test opens it.  Runs on a lane worker. */
static void gated_task(void *arg)
{
    gate_t *g = arg;

    int now = atomic_fetch_add(&g->running, 1) + 1;
    int prev = atomic_load(&g->max_running);
    while (now > prev && !atomic_compare_exchange_weak(&g->max_running, &prev, now))
        ; /* CAS-max */
    atomic_fetch_add(&g->total_runs, 1);

    pthread_mutex_lock(&g->mu);
    atomic_fetch_add(&g->entered, 1);
    pthread_cond_broadcast(&g->entered_cv);
    while (!g->gate_open)
        pthread_cond_wait(&g->release_cv, &g->mu);
    pthread_mutex_unlock(&g->mu);

    atomic_fetch_sub(&g->running, 1);
}

static void gate_open(gate_t *g)
{
    pthread_mutex_lock(&g->mu);
    g->gate_open = 1;
    pthread_cond_broadcast(&g->release_cv);
    pthread_mutex_unlock(&g->mu);
}

/* Wait (bounded) until at least `n` tasks have entered the gate. */
static int wait_entered_at_least(gate_t *g, int n, int timeout_ms)
{
    struct timespec dl;
    clock_gettime(CLOCK_REALTIME, &dl);
    dl.tv_sec += timeout_ms / 1000;
    dl.tv_nsec += (timeout_ms % 1000) * 1000000L;
    if (dl.tv_nsec >= 1000000000L) { dl.tv_sec++; dl.tv_nsec -= 1000000000L; }

    pthread_mutex_lock(&g->mu);
    while (atomic_load(&g->entered) < n) {
        if (pthread_cond_timedwait(&g->entered_cv, &g->mu, &dl) == ETIMEDOUT)
            break;
    }
    int got = atomic_load(&g->entered);
    pthread_mutex_unlock(&g->mu);
    return got;
}

/* ── (a) submit is non-blocking ───────────────────────────────────── */

/* A submit must return promptly even when the task it submits will block
 * indefinitely.  We submit a gated (never-released until we open) task and
 * assert submit returned without us opening the gate. */
static void test_submit_is_non_blocking(void)
{
    slowlane_t sl;
    gate_t     g;
    gate_init(&g);
    TEST_ASSERT_EQUAL_INT(0, slowlane_init(&sl, 2, 2));

    struct timespec t0, t1;
    clock_gettime(CLOCK_MONOTONIC, &t0);
    int rc = slowlane_submit(&sl, 0, gated_task, &g);
    clock_gettime(CLOCK_MONOTONIC, &t1);
    TEST_ASSERT_EQUAL_INT(0, rc);

    /* submit returned while the task is (or will be) parked: a blocking call
     * would not have returned.  Bound the wall time generously (<500ms). */
    long ms = (t1.tv_sec - t0.tv_sec) * 1000 +
              (t1.tv_nsec - t0.tv_nsec) / 1000000;
    TEST_ASSERT_TRUE_MESSAGE(ms < 500, "submit blocked the caller");

    gate_open(&g);
    slowlane_fini(&sl);   /* drains + joins */
    gate_destroy(&g);
}

/* ── (b) concurrent running tasks never exceed the cap ────────────── */

static void test_admission_cap_bounds_concurrency(void)
{
    slowlane_t sl;
    gate_t     g;
    gate_init(&g);
    const int n_workers   = 4;
    const int max_inflight = 2;
    TEST_ASSERT_EQUAL_INT(0, slowlane_init(&sl, n_workers, max_inflight));

    /* Submit more tasks than the cap; all must enqueue, but only `cap` may run
     * concurrently even though 4 workers exist. */
    for (int i = 0; i < 6; i++) {
        int rc = slowlane_submit(&sl, 0, gated_task, &g);
        /* 0 while a slot is free, -EAGAIN once the cap is full (still queued). */
        TEST_ASSERT_TRUE_MESSAGE(rc == 0 || rc == -EAGAIN, "submit dropped work");
    }

    /* Let the first wave run and park at the gate. */
    int entered = wait_entered_at_least(&g, max_inflight, 2000);
    TEST_ASSERT_EQUAL_INT(max_inflight, entered);

    /* Give any (incorrect) over-admission a chance to manifest, then assert the
     * peak never exceeded the cap. */
    struct timespec nap = { .tv_sec = 0, .tv_nsec = 100 * 1000000L };
    nanosleep(&nap, NULL);
    TEST_ASSERT_EQUAL_INT(max_inflight, atomic_load(&g.running));
    TEST_ASSERT_TRUE_MESSAGE(atomic_load(&g.max_running) <= max_inflight,
                             "concurrent running tasks exceeded the cap");

    /* Release everything; all 6 must eventually run and the peak stays <= cap. */
    gate_open(&g);
    slowlane_fini(&sl);
    TEST_ASSERT_EQUAL_INT(6, atomic_load(&g.total_runs));
    TEST_ASSERT_TRUE_MESSAGE(atomic_load(&g.max_running) <= max_inflight,
                             "cap violated across the full run");
    gate_destroy(&g);
}

/* ── (+) -EAGAIN when the cap is full at submit (work still queued) ── */

static void test_submit_returns_eagain_when_cap_full(void)
{
    slowlane_t sl;
    gate_t     g;
    gate_init(&g);
    const int max_inflight = 1;
    TEST_ASSERT_EQUAL_INT(0, slowlane_init(&sl, 4, max_inflight));

    /* First task takes the only running slot and parks. */
    TEST_ASSERT_EQUAL_INT(0, slowlane_submit(&sl, 0, gated_task, &g));
    TEST_ASSERT_EQUAL_INT(1, wait_entered_at_least(&g, 1, 2000));

    /* Cap is now full: the next submit is deferred -> -EAGAIN, but queued. */
    int rc = slowlane_submit(&sl, 0, gated_task, &g);
    TEST_ASSERT_EQUAL_INT_MESSAGE(-EAGAIN, rc,
                                  "over-cap submit must return -EAGAIN");

    gate_open(&g);
    slowlane_fini(&sl);
    /* The deferred task was NOT dropped: both ran. */
    TEST_ASSERT_EQUAL_INT(2, atomic_load(&g.total_runs));
    gate_destroy(&g);
}

/* ── (c) per-key dedup returns -EEXIST for an in-flight key ────────── */

static void test_dedup_returns_eexist_for_inflight_key(void)
{
    slowlane_t sl;
    gate_t     g;
    gate_init(&g);
    TEST_ASSERT_EQUAL_INT(0, slowlane_init(&sl, 4, 4));

    const uint64_t KEY = 0x51A8;

    /* First submit of KEY is accepted and runs/parks. */
    TEST_ASSERT_EQUAL_INT(0, slowlane_submit(&sl, KEY, gated_task, &g));
    TEST_ASSERT_EQUAL_INT(1, wait_entered_at_least(&g, 1, 2000));

    /* Second submit of the SAME in-flight key is a deduped no-op. */
    TEST_ASSERT_EQUAL_INT(-EEXIST,
                          slowlane_submit(&sl, KEY, gated_task, &g));

    /* A DIFFERENT key is still admitted. */
    TEST_ASSERT_EQUAL_INT(0, slowlane_submit(&sl, KEY + 1, gated_task, &g));
    TEST_ASSERT_EQUAL_INT(2, wait_entered_at_least(&g, 2, 2000));

    gate_open(&g);
    slowlane_fini(&sl);
    /* Exactly 2 runs: the duplicate KEY submit did not run a third task. */
    TEST_ASSERT_EQUAL_INT(2, atomic_load(&g.total_runs));
    gate_destroy(&g);
}

/* After an in-flight key completes, the SAME key is admittable again. */
static void test_dedup_key_reusable_after_completion(void)
{
    slowlane_t sl;
    gate_t     g;
    gate_init(&g);
    TEST_ASSERT_EQUAL_INT(0, slowlane_init(&sl, 2, 2));

    const uint64_t KEY = 7;

    TEST_ASSERT_EQUAL_INT(0, slowlane_submit(&sl, KEY, gated_task, &g));
    TEST_ASSERT_EQUAL_INT(1, wait_entered_at_least(&g, 1, 2000));
    /* Release and let it fully complete (running drops to 0, key cleared). */
    gate_open(&g);

    /* Re-close the gate for the second wave. */
    pthread_mutex_lock(&g.mu);
    /* wait until the first task left fn */
    while (atomic_load(&g.running) != 0) {
        pthread_mutex_unlock(&g.mu);
        struct timespec s = { 0, 1 * 1000000L };
        nanosleep(&s, NULL);
        pthread_mutex_lock(&g.mu);
    }
    g.gate_open = 0;
    atomic_store(&g.entered, 0);
    pthread_mutex_unlock(&g.mu);

    /* Same key, after completion -> admitted again (0, not -EEXIST). */
    TEST_ASSERT_EQUAL_INT(0, slowlane_submit(&sl, KEY, gated_task, &g));
    TEST_ASSERT_EQUAL_INT(1, wait_entered_at_least(&g, 1, 2000));

    gate_open(&g);
    slowlane_fini(&sl);
    TEST_ASSERT_EQUAL_INT(2, atomic_load(&g.total_runs));
    gate_destroy(&g);
}

/* ── (d) fn runs exactly once ──────────────────────────────────────── */

static void test_fn_runs_exactly_once(void)
{
    slowlane_t sl;
    gate_t     g;
    gate_init(&g);
    TEST_ASSERT_EQUAL_INT(0, slowlane_init(&sl, 4, 4));

    /* The gate starts open so non-blocking tasks run straight through. */
    gate_open(&g);

    const int N = 32;
    for (int i = 0; i < N; i++)
        TEST_ASSERT_EQUAL_INT(0, slowlane_submit(&sl, 0, gated_task, &g));

    slowlane_fini(&sl);   /* drains + joins -> all submitted work has run */

    /* Exactly N executions: no task dropped, none double-run. */
    TEST_ASSERT_EQUAL_INT(N, atomic_load(&g.total_runs));
    gate_destroy(&g);
}

/* ── init/fini hardening: fini on an empty lane, double-fini safety ── */

static void test_fini_on_idle_lane(void)
{
    slowlane_t sl;
    TEST_ASSERT_EQUAL_INT(0, slowlane_init(&sl, 2, 2));
    slowlane_fini(&sl);   /* no submits — must drain/join cleanly */
    TEST_PASS();
}

void test_slowlane(void)
{
    RUN_TEST(test_submit_is_non_blocking);
    RUN_TEST(test_admission_cap_bounds_concurrency);
    RUN_TEST(test_submit_returns_eagain_when_cap_full);
    RUN_TEST(test_dedup_returns_eexist_for_inflight_key);
    RUN_TEST(test_dedup_key_reusable_after_completion);
    RUN_TEST(test_fn_runs_exactly_once);
    RUN_TEST(test_fini_on_idle_lane);
}