#include "seobeo/seobeo_worker.h"

#include <assert.h>
#include <stdatomic.h>
#include <stdlib.h>
#include <unistd.h>

typedef struct {
  atomic_int executed;
  atomic_int cleaned;
  atomic_int entered;
  atomic_int release;
  atomic_int cleanup_entered;
  atomic_int hold_cleanup;
} Worker_Test_Context;

typedef struct {
  Seobeo_Worker_Pool *p_pool;
  Seobeo_Worker_Result result;
} Pool_Shutdown_Context;

typedef struct {
  Seobeo_Worker_Pool *p_pool;
  atomic_int returned;
} Pool_Wait_Context;

typedef struct {
  Seobeo_Worker_Pool *p_pool;
  Seobeo_Worker_Result wait_result;
  Seobeo_Worker_Result shutdown_result;
  atomic_int cleaned;
} Recursive_Cleanup_Context;

static void count_task(void *p_context)
{
  Worker_Test_Context *p_test = p_context;
  atomic_fetch_add(&p_test->executed, 1);
}

static void no_op_task(void *p_context)
{
  (void)p_context;
}

static void blocking_task(void *p_context)
{
  Worker_Test_Context *p_test = p_context;
  atomic_fetch_add(&p_test->executed, 1);
  atomic_store(&p_test->entered, 1);
  while (!atomic_load(&p_test->release))
    usleep(1000);
}

static void cleanup_task(void *p_context)
{
  Worker_Test_Context *p_test = p_context;
  if (atomic_load(&p_test->hold_cleanup))
  {
    atomic_store(&p_test->cleanup_entered, 1);
    while (!atomic_load(&p_test->release))
      usleep(1000);
  }
  atomic_fetch_add(&p_test->cleaned, 1);
}

static void shutdown_pool_task(void *p_context)
{
  Pool_Shutdown_Context *p_shutdown = p_context;
  p_shutdown->result =
      Seobeo_Worker_Pool_Shutdown(p_shutdown->p_pool, FALSE);
}

static void wait_pool_task(void *p_context)
{
  Pool_Wait_Context *p_wait = p_context;
  assert(Seobeo_Worker_Pool_Wait(p_wait->p_pool) == SEOBEO_WORKER_OK);
  atomic_store(&p_wait->returned, 1);
}

static void recursive_pool_cleanup(void *p_context)
{
  Recursive_Cleanup_Context *p_cleanup = p_context;
  p_cleanup->wait_result =
      Seobeo_Worker_Pool_Wait(p_cleanup->p_pool);
  p_cleanup->shutdown_result =
      Seobeo_Worker_Pool_Shutdown(p_cleanup->p_pool, FALSE);
  Seobeo_Worker_Pool_Destroy(p_cleanup->p_pool);
  atomic_store(&p_cleanup->cleaned, 1);
}

static void wait_for_value(atomic_int *p_value, int expected)
{
  for (int i = 0; i < 2000; i++)
  {
    if (atomic_load(p_value) == expected)
      return;
    usleep(1000);
  }
  assert(FALSE && "worker operation timed out");
}

static void test_joinable_thread(void)
{
  Worker_Test_Context context = {0};
  Seobeo_Thread *p_thread =
      Seobeo_Thread_Start(count_task, &context, cleanup_task);
  assert(p_thread);
  assert(Seobeo_Thread_Join(p_thread) == SEOBEO_WORKER_OK);
  assert(atomic_load(&context.executed) == 1);
  assert(atomic_load(&context.cleaned) == 1);
}

static void test_detached_thread(void)
{
  Worker_Test_Context context = {0};
  assert(Seobeo_Thread_Start_Detached(
      count_task,
      &context,
      cleanup_task) == SEOBEO_WORKER_OK);
  wait_for_value(&context.cleaned, 1);
  assert(atomic_load(&context.executed) == 1);
}

static void test_pool_drain(void)
{
  Worker_Test_Context context = {0};
  Seobeo_Worker_Pool *p_pool =
      Seobeo_Worker_Pool_Create(2, 8);
  assert(p_pool);
  for (int i = 0; i < 8; i++)
  {
    assert(Seobeo_Worker_Pool_Submit(
        p_pool,
        count_task,
        &context,
        cleanup_task) == SEOBEO_WORKER_OK);
  }
  assert(Seobeo_Worker_Pool_Wait(p_pool) == SEOBEO_WORKER_OK);
  assert(atomic_load(&context.executed) == 8);
  assert(atomic_load(&context.cleaned) == 8);
  assert(Seobeo_Worker_Pool_Pending(p_pool) == 0);
  assert(Seobeo_Worker_Pool_Active(p_pool) == 0);
  assert(Seobeo_Worker_Pool_Shutdown(p_pool, TRUE) == SEOBEO_WORKER_OK);
  assert(Seobeo_Worker_Pool_Submit(
      p_pool,
      count_task,
      &context,
      cleanup_task) == SEOBEO_WORKER_STOPPED);
  Seobeo_Worker_Pool_Destroy(p_pool);
}

static void test_pool_queue_limit(void)
{
  Worker_Test_Context context = {0};
  Seobeo_Worker_Pool *p_pool =
      Seobeo_Worker_Pool_Create(1, 1);
  assert(p_pool);
  assert(Seobeo_Worker_Pool_Submit(
      p_pool,
      blocking_task,
      &context,
      cleanup_task) == SEOBEO_WORKER_OK);
  wait_for_value(&context.entered, 1);
  assert(Seobeo_Worker_Pool_Submit(
      p_pool,
      count_task,
      &context,
      cleanup_task) == SEOBEO_WORKER_OK);
  assert(Seobeo_Worker_Pool_Submit(
      p_pool,
      count_task,
      &context,
      cleanup_task) == SEOBEO_WORKER_QUEUE_FULL);
  atomic_store(&context.release, 1);
  assert(Seobeo_Worker_Pool_Wait(p_pool) == SEOBEO_WORKER_OK);
  assert(atomic_load(&context.executed) == 2);
  assert(atomic_load(&context.cleaned) == 2);
  Seobeo_Worker_Pool_Destroy(p_pool);
}

static void test_pool_cancel_queue(void)
{
  Worker_Test_Context context = {0};
  Seobeo_Worker_Pool *p_pool =
      Seobeo_Worker_Pool_Create(1, 4);
  assert(p_pool);
  assert(Seobeo_Worker_Pool_Submit(
      p_pool,
      blocking_task,
      &context,
      cleanup_task) == SEOBEO_WORKER_OK);
  wait_for_value(&context.entered, 1);
  assert(Seobeo_Worker_Pool_Submit(
      p_pool,
      count_task,
      &context,
      cleanup_task) == SEOBEO_WORKER_OK);
  Recursive_Cleanup_Context recursive = {
    .p_pool = p_pool,
    .wait_result = SEOBEO_WORKER_OK,
    .shutdown_result = SEOBEO_WORKER_OK,
  };
  assert(Seobeo_Worker_Pool_Submit(
      p_pool,
      no_op_task,
      &recursive,
      recursive_pool_cleanup) == SEOBEO_WORKER_OK);

  atomic_store(&context.hold_cleanup, 1);
  Pool_Shutdown_Context shutdown = {
    .p_pool = p_pool,
    .result = SEOBEO_WORKER_THREAD_ERROR,
  };
  Seobeo_Thread *p_shutdown = Seobeo_Thread_Start(
      shutdown_pool_task,
      &shutdown,
      NULL);
  assert(p_shutdown);
  wait_for_value(&context.cleanup_entered, 1);

  Pool_Wait_Context wait = {
    .p_pool = p_pool,
  };
  Seobeo_Thread *p_wait = Seobeo_Thread_Start(
      wait_pool_task,
      &wait,
      NULL);
  assert(p_wait);
  usleep(20000);
  assert(atomic_load(&wait.returned) == 0);

  atomic_store(&context.release, 1);
  assert(Seobeo_Thread_Join(p_shutdown) == SEOBEO_WORKER_OK);
  assert(Seobeo_Thread_Join(p_wait) == SEOBEO_WORKER_OK);
  assert(shutdown.result == SEOBEO_WORKER_OK);
  assert(atomic_load(&context.executed) == 1);
  assert(atomic_load(&context.cleaned) == 2);
  assert(atomic_load(&recursive.cleaned) == 1);
  assert(recursive.wait_result == SEOBEO_WORKER_INVALID_ARGUMENT);
  assert(recursive.shutdown_result == SEOBEO_WORKER_INVALID_ARGUMENT);
  Seobeo_Worker_Pool_Destroy(p_pool);
}

int main(void)
{
  assert(Seobeo_Thread_Current_Id() != 0);
  assert(Seobeo_Thread_Start(NULL, NULL, NULL) == NULL);
  assert(Seobeo_Thread_Start_Detached(
      NULL,
      NULL,
      NULL) == SEOBEO_WORKER_INVALID_ARGUMENT);
  assert(Seobeo_Worker_Pool_Create(0, 1) == NULL);
  assert(Seobeo_Worker_Pool_Create(1, 0) == NULL);

  test_joinable_thread();
  test_detached_thread();
  test_pool_drain();
  test_pool_queue_limit();
  test_pool_cancel_queue();
  return 0;
}