#define _GNU_SOURCE
#include "slowlane.h"

#include <errno.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>

/*
 * slowlane.c — bounded, admission-capped executor with per-key in-flight
 * dedup, running long blocking I/O off the Raft drive threads.
 *
 * raft-drive-async task 1.2 — Requirements 3.1, 3.2, 3.3, 3.4, 1.5.
 *
 * Concurrency model (single mutex `lock`, one condvar `cond`):
 *   - `q_head/q_tail` is the FIFO of admitted-but-not-yet-running tasks.
 *   - `n_running` counts tasks executing on a worker; admission holds
 *     `n_running <= max_inflight` at all times.
 *   - `inflight_keys[]` holds every non-zero key that is queued OR running;
 *     a key is added at submit and removed when its `fn` returns, so a second
 *     submit of the same key in that window dedups (-EEXIST).
 *   - Workers wait on `cond` for "work pending AND a running slot free".
 *     submit and task-completion both signal `cond`.
 *
 * Lock discipline: `fn` ALWAYS runs with `lock` released (it blocks for I/O);
 * the lane never holds `lock` across user work — so a slow/stuck `fn` cannot
 * stall submit() or other workers' admission decisions.
 */

/* ── dedup key set helpers (caller holds sl->lock) ────────────────── */

static bool key_present_locked(const slowlane_t *sl, uint64_t key)
{
    for (int i = 0; i < sl->n_inflight_keys; i++)
        if (sl->inflight_keys[i] == key)
            return true;
    return false;
}

static int key_add_locked(slowlane_t *sl, uint64_t key)
{
    if (sl->n_inflight_keys == sl->cap_inflight_keys) {
        int newcap = sl->cap_inflight_keys ? sl->cap_inflight_keys * 2 : 8;
        uint64_t *grown =
            realloc(sl->inflight_keys, (size_t)newcap * sizeof(uint64_t));
        if (grown == NULL)
            return -ENOMEM;
        sl->inflight_keys     = grown;
        sl->cap_inflight_keys = newcap;
    }
    sl->inflight_keys[sl->n_inflight_keys++] = key;
    return 0;
}

static void key_remove_locked(slowlane_t *sl, uint64_t key)
{
    if (key == 0)
        return;
    for (int i = 0; i < sl->n_inflight_keys; i++) {
        if (sl->inflight_keys[i] == key) {
            /* swap-remove (order irrelevant for a membership set) */
            sl->inflight_keys[i] =
                sl->inflight_keys[--sl->n_inflight_keys];
            return;
        }
    }
}

/* ── worker thread ────────────────────────────────────────────────── */

static void *slowlane_worker(void *arg)
{
    slowlane_t *sl = arg;

    pthread_mutex_lock(&sl->lock);
    for (;;) {
        /* Wait for: a queued task AND a free running slot.  On shutdown we
         * still drain the queue (so already-submitted work runs exactly once),
         * but stop waiting for new arrivals. */
        while (!sl->stopping &&
               (sl->q_head == NULL || sl->n_running >= sl->max_inflight))
            pthread_cond_wait(&sl->cond, &sl->lock);

        if (sl->q_head == NULL) {
            /* Nothing left to run. If stopping, exit; otherwise (can only be a
             * spurious wake on a non-empty-but-cap-full path that then emptied)
             * loop back to wait. */
            if (sl->stopping)
                break;
            continue;
        }

        /* Even while stopping, respect the admission cap on the drain. */
        if (sl->n_running >= sl->max_inflight) {
            if (sl->stopping) {
                /* A free slot will appear when a running task completes and
                 * signals; wait for it so we can finish draining. */
                pthread_cond_wait(&sl->cond, &sl->lock);
            }
            continue;
        }

        /* Pop FIFO head, claim a running slot. */
        slowlane_task_t *t = sl->q_head;
        sl->q_head = t->next;
        if (sl->q_head == NULL)
            sl->q_tail = NULL;
        sl->n_running++;

        pthread_mutex_unlock(&sl->lock);
        /* Run user work with NO lock held (it blocks for I/O). Exactly once. */
        t->fn(t->arg);
        uint64_t done_key = t->key;
        free(t);
        pthread_mutex_lock(&sl->lock);

        sl->n_running--;
        key_remove_locked(sl, done_key);
        /* A slot freed (admission) and a key cleared (dedup): wake a waiter
         * that may now be admissible or a submit blocked on nothing here. */
        pthread_cond_broadcast(&sl->cond);
    }
    pthread_mutex_unlock(&sl->lock);
    return NULL;
}

/* ── public API ───────────────────────────────────────────────────── */

int slowlane_init(slowlane_t *sl, int n_workers, int max_inflight)
{
    if (sl == NULL || n_workers < 1 || max_inflight < 1)
        return -EINVAL;

    memset(sl, 0, sizeof(*sl));
    sl->n_workers    = n_workers;
    sl->max_inflight = max_inflight;

    if (pthread_mutex_init(&sl->lock, NULL) != 0)
        return -errno;
    if (pthread_cond_init(&sl->cond, NULL) != 0) {
        pthread_mutex_destroy(&sl->lock);
        return -errno;
    }

    sl->workers = calloc((size_t)n_workers, sizeof(pthread_t));
    if (sl->workers == NULL) {
        pthread_cond_destroy(&sl->cond);
        pthread_mutex_destroy(&sl->lock);
        return -ENOMEM;
    }

    for (int i = 0; i < n_workers; i++) {
        int rc = pthread_create(&sl->workers[i], NULL, slowlane_worker, sl);
        if (rc != 0) {
            /* Roll back the partially-started pool: stop and join what exists. */
            pthread_mutex_lock(&sl->lock);
            sl->stopping = 1;
            pthread_cond_broadcast(&sl->cond);
            pthread_mutex_unlock(&sl->lock);
            for (int j = 0; j < i; j++)
                pthread_join(sl->workers[j], NULL);
            free(sl->workers);
            sl->workers = NULL;
            pthread_cond_destroy(&sl->cond);
            pthread_mutex_destroy(&sl->lock);
            return -rc;
        }
        sl->started++;
    }
    return 0;
}

int slowlane_submit(slowlane_t *sl, uint64_t key, slowlane_fn_t fn, void *arg)
{
    if (sl == NULL || fn == NULL)
        return -EINVAL;

    pthread_mutex_lock(&sl->lock);

    if (sl->stopping) {
        pthread_mutex_unlock(&sl->lock);
        return -ESHUTDOWN;
    }

    /* Per-key in-flight dedup: a non-zero key already queued/running is a
     * no-op (R1.5 snapshot-per-standby, R2.5 reconnect-per-peer). */
    if (key != 0 && key_present_locked(sl, key)) {
        pthread_mutex_unlock(&sl->lock);
        return -EEXIST;
    }

    slowlane_task_t *t = calloc(1, sizeof(*t));
    if (t == NULL) {
        pthread_mutex_unlock(&sl->lock);
        return -ENOMEM;
    }
    t->key = key;
    t->fn  = fn;
    t->arg = arg;
    t->next = NULL;

    if (key != 0) {
        int rc = key_add_locked(sl, key);
        if (rc != 0) {
            free(t);
            pthread_mutex_unlock(&sl->lock);
            return rc;
        }
    }

    /* Append to FIFO tail. */
    if (sl->q_tail != NULL)
        sl->q_tail->next = t;
    else
        sl->q_head = t;
    sl->q_tail = t;

    /* Admission status at submit time: a free running slot means the task can
     * start immediately (0); a full cap means it is DEFERRED but still queued
     * (-EAGAIN) — never dropped, never run on a drive thread. */
    int rc = (sl->n_running >= sl->max_inflight) ? -EAGAIN : 0;

    pthread_cond_signal(&sl->cond);
    pthread_mutex_unlock(&sl->lock);
    return rc;
}

void slowlane_fini(slowlane_t *sl)
{
    if (sl == NULL || sl->workers == NULL)
        return;

    /* Request drain+stop; workers finish the queued tasks, then exit. */
    pthread_mutex_lock(&sl->lock);
    sl->stopping = 1;
    pthread_cond_broadcast(&sl->cond);
    pthread_mutex_unlock(&sl->lock);

    for (int i = 0; i < sl->started; i++)
        pthread_join(sl->workers[i], NULL);

    /* All workers joined: no concurrent access remains. Free any tasks that
     * were never run (only possible if a worker create failed mid-init; the
     * normal drain path empties the queue). */
    slowlane_task_t *t = sl->q_head;
    while (t != NULL) {
        slowlane_task_t *next = t->next;
        free(t);
        t = next;
    }
    sl->q_head = sl->q_tail = NULL;

    free(sl->inflight_keys);
    sl->inflight_keys     = NULL;
    sl->n_inflight_keys   = 0;
    sl->cap_inflight_keys = 0;

    free(sl->workers);
    sl->workers = NULL;
    sl->started = 0;

    pthread_cond_destroy(&sl->cond);
    pthread_mutex_destroy(&sl->lock);
}