#include "u_async.h"
#include <stdio.h>

/* Test statistics - same as in the original test */
static struct {
    int tests_run;
    int tests_passed;
    int tests_failed;
    
    /* Timer statistics */
    int timer_callbacks;
    int timer_cancellations;
    int immediate_timeouts;
    
    /* Socket statistics */
    int socket_events;
    int socket_errors;
    
    /* Error statistics */
    int memory_allocation_errors;
    int invalid_parameter_errors;
    int race_condition_errors;
} test_stats = {0};

/* Test result tracking */
#define TEST_START(name) do { \
    printf("TEST: %s... \n", name); \
    printf("  immediate_timeouts before: %d\n", test_stats.immediate_timeouts); \
    test_stats.tests_run++; \
} while(0)

#define TEST_PASS() do { \
    printf("  immediate_timeouts after: %d\n", test_stats.immediate_timeouts); \
    printf("PASS\n"); \
    test_stats.tests_passed++; \
} while(0)

#define TEST_FAIL(msg) do { \
    printf("  immediate_timeouts after: %d\n", test_stats.immediate_timeouts); \
    printf("FAIL: %s\n", msg); \
    test_stats.tests_failed++; \
} while(0)

#define ASSERT_EQ(a, b, msg) do { \
    if ((a) != (b)) { \
        printf("  Expected: %ld, Got: %ld\n", (long)(b), (long)(a)); \
        TEST_FAIL(msg); \
        return; \
    } \
} while(0)

#define ASSERT_NOT_NULL(ptr, msg) do { \
    if ((ptr) == NULL) { \
        TEST_FAIL(msg); \
        return; \
    } \
} while(0)

/* Test context for callbacks */
typedef struct {
    int callback_count;
    int expected_count;
    int callback_arg;
    int timeout_ms;
    uasync_t* ua;
    void* timer_id;
} test_context_t;

/* Timer callback for testing */
static void test_timer_callback(void* arg) {
    test_context_t* ctx = (test_context_t*)arg;
    ctx->callback_count++;
    test_stats.timer_callbacks++;
    
    if (ctx->timeout_ms == 0) {
        test_stats.immediate_timeouts++;
    }
}

/* Test 1: Basic timer functionality */
static void test_basic_timers(void) {
    TEST_START("Basic timer functionality");
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    test_context_t ctx = {0};
    ctx.expected_count = 3;
    ctx.timeout_ms = 10;  // Non-zero timeout
    
    /* Set multiple timers with different timeouts */
    void* timer1 = uasync_set_timeout(ua, 10, &ctx, test_timer_callback);  /* 1ms */
    void* timer2 = uasync_set_timeout(ua, 20, &ctx, test_timer_callback);  /* 2ms */
    void* timer3 = uasync_set_timeout(ua, 30, &ctx, test_timer_callback);  /* 3ms */
    
    ASSERT_NOT_NULL(timer1, "Failed to set timer 1");
    ASSERT_NOT_NULL(timer2, "Failed to set timer 2");
    ASSERT_NOT_NULL(timer3, "Failed to set timer 3");
    
    /* Poll and verify timers fire in order */
    int poll_count = 0;
    while (ctx.callback_count < ctx.expected_count && poll_count < 100) {
        uasync_poll(ua, 10);  /* 1ms poll */
        poll_count++;
    }
    
    ASSERT_EQ(ctx.callback_count, ctx.expected_count, "Not all timers fired");
    
    /* Cleanup */
    uasync_destroy(ua, 0);
    TEST_PASS();
}

/* Test 3: Immediate timeout handling */
static void test_immediate_timeouts(void) {
    TEST_START("Immediate timeout handling");
    
    uasync_t* ua = uasync_create();
    ASSERT_NOT_NULL(ua, "Failed to create uasync instance");
    
    test_context_t ctx = {0};
    ctx.expected_count = 5;
    ctx.timeout_ms = 0;  // This should trigger immediate timeout counting
    
    /* Set multiple immediate timeouts (0ms) */
    for (int i = 0; i < 5; i++) {
        void* timer = uasync_set_timeout(ua, 0, &ctx, test_timer_callback);
        ASSERT_NOT_NULL(timer, "Failed to set immediate timer");
    }
    
    /* Immediate timeouts should fire during next poll */
    ASSERT_EQ(ctx.callback_count, 0, "Callbacks fired too early");
    
    uasync_poll(ua, 1);  /* Minimal poll */
    
    ASSERT_EQ(ctx.callback_count, ctx.expected_count, "Immediate timeouts didn't fire correctly");
    ASSERT_EQ(test_stats.immediate_timeouts, 5, "Immediate timeout counter incorrect");
    
    uasync_destroy(ua, 0);
    TEST_PASS();
}

int main() {
    printf("=== Tracking immediate_timeouts counter ===\n");
    printf("Initial: immediate_timeouts = %d\n", test_stats.immediate_timeouts);
    
    test_basic_timers();
    printf("After basic_timers: immediate_timeouts = %d\n", test_stats.immediate_timeouts);
    
    test_immediate_timeouts();
    printf("After immediate_timeouts: immediate_timeouts = %d\n", test_stats.immediate_timeouts);
    
    return 0;
}