#include "unity/unity.h"
#include "zlib.h"

#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <pthread.h>

/* Detect availability of C11 atomics similarly to the source, to decide
   whether concurrent initialization is supported. */
#if defined(__STDC__) && __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_ATOMICS__)
#define HAVE_C11_ATOMICS 1
#else
#define HAVE_C11_ATOMICS 0
#endif

/* Number of threads used in concurrency tests */
#ifndef ONCE_TEST_THREADS
#define ONCE_TEST_THREADS 8
#endif

/* Simple shared state for thread start synchronization */
typedef struct {
    volatile int start_flag;
} start_sync_t;

typedef struct {
    start_sync_t *sync;
    const unsigned char *buf;
    size_t len;
    const void *table_ptr;
    unsigned long crc_out;
} thread_arg_t;

static void busy_wait_until_start(start_sync_t *sync) {
    /* Simple spin-wait on a volatile flag to try to start calls together. */
    while (!sync->start_flag) { /* spin */ }
}

static void *thread_get_table(void *arg) {
    thread_arg_t *a = (thread_arg_t *)arg;
    busy_wait_until_start(a->sync);
    /* Call get_crc_table() to exercise once() initialization path. */
    a->table_ptr = (const void *)get_crc_table();
    /* Also do a CRC to ensure the table is actually usable. */
    if (a->buf && a->len > 0) {
        a->crc_out = crc32(0L, a->buf, (uInt)a->len);
    } else {
        a->crc_out = 0UL;
    }
    return NULL;
}

static void *thread_do_crc(void *arg) {
    thread_arg_t *a = (thread_arg_t *)arg;
    busy_wait_until_start(a->sync);
    /* Just compute CRC many times to apply load. */
    unsigned long crc = 0L;
    for (int i = 0; i < 1000; i++) {
        crc = crc32(0L, a->buf, (uInt)a->len);
    }
    a->crc_out = crc;
    /* Read the table pointer once to compare pointers across threads too. */
    a->table_ptr = (const void *)get_crc_table();
    return NULL;
}

void setUp(void) {
    /* No pre-initialization here so that the concurrent init test can
       actually race on first initialization when atomics are available. */
}

void tearDown(void) {
    /* Nothing to clean up */
}

/* Test 1: Concurrent initialization via get_crc_table() (only if atomics available).
   Validates that all threads observe the same table pointer and no crashes occur. */
void test_once_concurrent_initialization_via_get_crc_table(void) {
#if HAVE_C11_ATOMICS
    pthread_t tids[ONCE_TEST_THREADS];
    thread_arg_t args[ONCE_TEST_THREADS];
    start_sync_t sync;
    sync.start_flag = 0;

    static const unsigned char data[] = "zlib once() concurrent init";
    for (int i = 0; i < ONCE_TEST_THREADS; i++) {
        args[i].sync = &sync;
        args[i].buf = data;
        args[i].len = strlen((const char *)data);
        args[i].table_ptr = NULL;
        args[i].crc_out = 0;
        int rc = pthread_create(&tids[i], NULL, thread_get_table, &args[i]);
        TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "pthread_create failed");
    }

    /* Start all threads at roughly the same time. */
    sync.start_flag = 1;

    for (int i = 0; i < ONCE_TEST_THREADS; i++) {
        pthread_join(tids[i], NULL);
    }

    /* All should have the same non-null table pointer. */
    const void *ptr0 = args[0].table_ptr;
    TEST_ASSERT_NOT_NULL(ptr0);
    for (int i = 1; i < ONCE_TEST_THREADS; i++) {
        TEST_ASSERT_EQUAL_PTR(ptr0, args[i].table_ptr);
    }

    /* A subsequent call should return the exact same pointer. */
    const void *ptr_after = (const void *)get_crc_table();
    TEST_ASSERT_EQUAL_PTR(ptr0, ptr_after);
#else
    TEST_IGNORE_MESSAGE("Concurrent initialization is not supported without C11 atomics; skipping.");
#endif
}

/* Test 2: Known CRC value, single-thread sanity check. */
void test_once_crc_known_value_single_thread(void) {
    const unsigned char msg[] = "123456789";
    unsigned long crc = crc32(0L, msg, (uInt)strlen((const char *)msg));
    TEST_ASSERT_EQUAL_HEX32(0xCBF43926UL, crc);
}

/* Test 3: Repeated get_crc_table() returns stable pointer (single-thread). */
void test_once_pointer_stability_single_thread(void) {
    const void *p1 = (const void *)get_crc_table();
    TEST_ASSERT_NOT_NULL(p1);

    for (int i = 0; i < 1000; i++) {
        const void *pi = (const void *)get_crc_table();
        TEST_ASSERT_EQUAL_PTR(p1, pi);
    }
}

/* Test 4: Concurrent CRC computations after single-thread pre-initialization.
   This is allowed even without atomics as long as get_crc_table() was called first. */
void test_once_concurrent_crc_after_preinit(void) {
    /* Pre-initialize tables single-threaded to avoid concurrent init in non-atomics builds. */
    const void *ptable = (const void *)get_crc_table();
    TEST_ASSERT_NOT_NULL(ptable);

    pthread_t tids[ONCE_TEST_THREADS];
    thread_arg_t args[ONCE_TEST_THREADS];
    start_sync_t sync;
    sync.start_flag = 0;

    static const unsigned char msg[] = "The quick brown fox jumps over the lazy dog";
    /* Compute expected CRC single-threaded. */
    unsigned long expected = crc32(0L, msg, (uInt)strlen((const char *)msg));

    for (int i = 0; i < ONCE_TEST_THREADS; i++) {
        args[i].sync = &sync;
        args[i].buf = msg;
        args[i].len = strlen((const char *)msg);
        args[i].table_ptr = NULL;
        args[i].crc_out = 0;
        int rc = pthread_create(&tids[i], NULL, thread_do_crc, &args[i]);
        TEST_ASSERT_EQUAL_INT_MESSAGE(0, rc, "pthread_create failed");
    }

    /* Start threads together. */
    sync.start_flag = 1;

    for (int i = 0; i < ONCE_TEST_THREADS; i++) {
        pthread_join(tids[i], NULL);
    }

    for (int i = 0; i < ONCE_TEST_THREADS; i++) {
        TEST_ASSERT_EQUAL_HEX32(expected, args[i].crc_out);
        TEST_ASSERT_EQUAL_PTR(ptable, args[i].table_ptr);
    }
}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_once_concurrent_initialization_via_get_crc_table);
    RUN_TEST(test_once_crc_known_value_single_thread);
    RUN_TEST(test_once_pointer_stability_single_thread);
    RUN_TEST(test_once_concurrent_crc_after_preinit);
    return UNITY_END();
}