#include "naut/buf.h"
#include "test.h"
#include <pthread.h>
#include <sched.h>
#include <string.h>

/* Concurrency test: the owner thread allocates buffers and hands them to N
 * worker threads which put() them back. Mirrors recv-on-reactor / free-on-hash
 * -worker. At the end every buffer must be back on the freelist. */

#define NBLOCKS 1024
#define NWORKERS 8

static naut_bufpool *pool;
static _Atomic int bad_cap;   /* worker-thread failures (harness CHECK is not MT-safe) */

struct chan { _Atomic(naut_buf *) slot[NBLOCKS]; _Atomic int head, tail; };
static struct chan ch;

static void *worker(void *arg) {
    (void)arg;
    for (;;) {
        int t = atomic_load(&ch.tail);
        if (t >= NBLOCKS) break;
        if (!atomic_compare_exchange_weak(&ch.tail, &t, t + 1)) continue;
        /* spin until producer publishes slot t */
        naut_buf *b;
        while (!(b = atomic_load_explicit(&ch.slot[t], memory_order_acquire)))
            sched_yield();
        if (b->cap != NAUT_BLOCK) atomic_fetch_add(&bad_cap, 1);
        memset(b->data, 0xab, b->cap);   /* touch the pages */
        naut_buf_put(b);
    }
    return NULL;
}

int main(void) {
    pool = naut_bufpool_create(NAUT_BLOCK, NBLOCKS, false);
    CHECK(pool != NULL);
    CHECK_EQ(naut_bufpool_capacity(pool), NBLOCKS);
    CHECK_EQ(naut_bufpool_available(pool), NBLOCKS);

    /* basic single-thread get/put */
    naut_buf *a = naut_buf_get(pool);
    CHECK(a != NULL);
    CHECK_EQ(naut_bufpool_available(pool), NBLOCKS - 1);
    naut_buf_ref(a);                       /* refcnt 1 -> 2 */
    naut_buf_put(a);                       /* 2 -> 1, stays out */
    CHECK_EQ(naut_bufpool_available(pool), NBLOCKS - 1);
    naut_buf_put(a);                       /* 1 -> 0, returns */
    CHECK_EQ(naut_bufpool_available(pool), NBLOCKS);

    /* exhaustion */
    naut_buf *held[NBLOCKS];
    for (int i = 0; i < NBLOCKS; i++) { held[i] = naut_buf_get(pool); CHECK(held[i]); }
    CHECK_EQ(naut_bufpool_available(pool), 0);
    CHECK(naut_buf_get(pool) == NULL);     /* empty => NULL, no crash */

    /* concurrent producer/consumer */
    atomic_store(&ch.head, 0);
    atomic_store(&ch.tail, 0);
    for (int i = 0; i < NBLOCKS; i++)
        atomic_store_explicit(&ch.slot[i], NULL, memory_order_relaxed);

    pthread_t th[NWORKERS];
    for (int i = 0; i < NWORKERS; i++) pthread_create(&th[i], NULL, worker, NULL);
    for (int i = 0; i < NBLOCKS; i++)                      /* publish */
        atomic_store_explicit(&ch.slot[i], held[i], memory_order_release);
    for (int i = 0; i < NWORKERS; i++) pthread_join(th[i], NULL);

    CHECK_EQ(atomic_load(&bad_cap), 0);
    CHECK_EQ(naut_bufpool_available(pool), NBLOCKS);
    /* and we can drain the whole pool again afterwards */
    int got = 0;
    while (naut_buf_get(pool)) got++;
    CHECK_EQ(got, NBLOCKS);

    naut_bufpool_destroy(pool);
    TEST_MAIN_END();
}