#include "seobeo/seobeo_worker.h"

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

typedef struct {
  Seobeo_Work_Function function;
  void *p_context;
  Seobeo_Work_Cleanup cleanup;
} Seobeo_Work_Item;

struct Seobeo_Thread {
  pthread_t thread;
  Seobeo_Work_Item item;
};

struct Seobeo_Worker_Pool {
  pthread_t *p_threads;
  Seobeo_Work_Item *p_queue;
  uint32 worker_count;
  uint32 queue_capacity;
  uint32 queue_head;
  uint32 queue_count;
  uint32 active_count;
  uint32 cleanup_count;
  boolean accepting;
  boolean stopping;
  boolean drain;
  boolean shutdown_started;
  boolean joined;
  pthread_mutex_t mutex;
  pthread_cond_t work_available;
  pthread_cond_t idle;
};

static _Thread_local Seobeo_Worker_Pool *g_current_worker_pool = NULL;

static void Seobeo_Work_Item_Run(Seobeo_Work_Item *p_item)
{
  p_item->function(p_item->p_context);
  if (p_item->cleanup)
    p_item->cleanup(p_item->p_context);
}

static void *Seobeo_Thread_Run(void *p_argument)
{
  Seobeo_Thread *p_thread = p_argument;
  Seobeo_Work_Item_Run(&p_thread->item);
  return NULL;
}

Seobeo_Thread *Seobeo_Thread_Start(
    Seobeo_Work_Function function,
    void *p_context,
    Seobeo_Work_Cleanup cleanup)
{
  if (!function)
    return NULL;

  Seobeo_Thread *p_thread = calloc(1, sizeof(*p_thread));
  if (!p_thread)
    return NULL;
  p_thread->item = (Seobeo_Work_Item){
    .function = function,
    .p_context = p_context,
    .cleanup = cleanup,
  };
  if (pthread_create(
          &p_thread->thread,
          NULL,
          Seobeo_Thread_Run,
          p_thread) != 0)
  {
    free(p_thread);
    return NULL;
  }
  return p_thread;
}

Seobeo_Worker_Result Seobeo_Thread_Join(Seobeo_Thread *p_thread)
{
  if (!p_thread)
    return SEOBEO_WORKER_INVALID_ARGUMENT;
  if (pthread_equal(pthread_self(), p_thread->thread))
    return SEOBEO_WORKER_INVALID_ARGUMENT;

  int result = pthread_join(p_thread->thread, NULL);
  free(p_thread);
  return result == 0
      ? SEOBEO_WORKER_OK
      : SEOBEO_WORKER_THREAD_ERROR;
}

static void *Seobeo_Thread_Run_Detached(void *p_argument)
{
  Seobeo_Work_Item *p_item = p_argument;
  Seobeo_Work_Item_Run(p_item);
  free(p_item);
  return NULL;
}

Seobeo_Worker_Result Seobeo_Thread_Start_Detached(
    Seobeo_Work_Function function,
    void *p_context,
    Seobeo_Work_Cleanup cleanup)
{
  if (!function)
    return SEOBEO_WORKER_INVALID_ARGUMENT;

  Seobeo_Work_Item *p_item = malloc(sizeof(*p_item));
  if (!p_item)
    return SEOBEO_WORKER_OUT_OF_MEMORY;
  *p_item = (Seobeo_Work_Item){
    .function = function,
    .p_context = p_context,
    .cleanup = cleanup,
  };

  pthread_attr_t attributes;
  if (pthread_attr_init(&attributes) != 0)
  {
    free(p_item);
    return SEOBEO_WORKER_THREAD_ERROR;
  }
  int result = pthread_attr_setdetachstate(
      &attributes,
      PTHREAD_CREATE_DETACHED);
  pthread_t thread;
  if (result == 0)
  {
    result = pthread_create(
        &thread,
        &attributes,
        Seobeo_Thread_Run_Detached,
        p_item);
  }
  pthread_attr_destroy(&attributes);
  if (result != 0)
  {
    free(p_item);
    return SEOBEO_WORKER_THREAD_ERROR;
  }
  return SEOBEO_WORKER_OK;
}

uint64 Seobeo_Thread_Current_Id(void)
{
  pthread_t thread = pthread_self();
  const uint8 *p_bytes = (const uint8 *)&thread;
  uint64 hash = 1469598103934665603ULL;
  for (size_t i = 0; i < sizeof(thread); i++)
  {
    hash ^= p_bytes[i];
    hash *= 1099511628211ULL;
  }
  return hash;
}

static boolean Seobeo_Worker_Pool_Is_Current_Thread(
    Seobeo_Worker_Pool *p_pool)
{
  pthread_t current = pthread_self();
  for (uint32 i = 0; i < p_pool->worker_count; i++)
  {
    if (pthread_equal(current, p_pool->p_threads[i]))
      return TRUE;
  }
  return FALSE;
}

static void *Seobeo_Worker_Pool_Run(void *p_argument)
{
  Seobeo_Worker_Pool *p_pool = p_argument;
  while (TRUE)
  {
    pthread_mutex_lock(&p_pool->mutex);
    while (p_pool->queue_count == 0 && !p_pool->stopping)
      pthread_cond_wait(&p_pool->work_available, &p_pool->mutex);

    if (p_pool->stopping &&
        (!p_pool->drain || p_pool->queue_count == 0))
    {
      pthread_mutex_unlock(&p_pool->mutex);
      break;
    }

    Seobeo_Work_Item item = p_pool->p_queue[p_pool->queue_head];
    p_pool->queue_head =
        (p_pool->queue_head + 1) % p_pool->queue_capacity;
    p_pool->queue_count--;
    p_pool->active_count++;
    pthread_mutex_unlock(&p_pool->mutex);

    Seobeo_Worker_Pool *p_previous_pool = g_current_worker_pool;
    g_current_worker_pool = p_pool;
    Seobeo_Work_Item_Run(&item);
    g_current_worker_pool = p_previous_pool;

    pthread_mutex_lock(&p_pool->mutex);
    p_pool->active_count--;
    if (p_pool->queue_count == 0 &&
        p_pool->active_count == 0 &&
        p_pool->cleanup_count == 0)
      pthread_cond_broadcast(&p_pool->idle);
    pthread_mutex_unlock(&p_pool->mutex);
  }
  return NULL;
}

Seobeo_Worker_Pool *Seobeo_Worker_Pool_Create(
    uint32 worker_count,
    uint32 queue_capacity)
{
  if (worker_count == 0 || queue_capacity == 0)
    return NULL;

  Seobeo_Worker_Pool *p_pool = calloc(1, sizeof(*p_pool));
  if (!p_pool)
    return NULL;
  p_pool->p_threads = calloc(worker_count, sizeof(*p_pool->p_threads));
  p_pool->p_queue = calloc(queue_capacity, sizeof(*p_pool->p_queue));
  if (!p_pool->p_threads || !p_pool->p_queue)
  {
    free(p_pool->p_threads);
    free(p_pool->p_queue);
    free(p_pool);
    return NULL;
  }

  p_pool->worker_count = worker_count;
  p_pool->queue_capacity = queue_capacity;
  p_pool->accepting = TRUE;
  boolean mutex_initialized = FALSE;
  boolean work_condition_initialized = FALSE;
  boolean idle_condition_initialized = FALSE;
  if (pthread_mutex_init(&p_pool->mutex, NULL) == 0)
    mutex_initialized = TRUE;
  if (mutex_initialized &&
      pthread_cond_init(&p_pool->work_available, NULL) == 0)
    work_condition_initialized = TRUE;
  if (work_condition_initialized &&
      pthread_cond_init(&p_pool->idle, NULL) == 0)
    idle_condition_initialized = TRUE;
  if (!idle_condition_initialized)
  {
    if (work_condition_initialized)
      pthread_cond_destroy(&p_pool->work_available);
    if (mutex_initialized)
      pthread_mutex_destroy(&p_pool->mutex);
    free(p_pool->p_threads);
    free(p_pool->p_queue);
    free(p_pool);
    return NULL;
  }

  uint32 created = 0;
  for (; created < worker_count; created++)
  {
    if (pthread_create(
            &p_pool->p_threads[created],
            NULL,
            Seobeo_Worker_Pool_Run,
            p_pool) != 0)
      break;
  }
  if (created != worker_count)
  {
    pthread_mutex_lock(&p_pool->mutex);
    p_pool->stopping = TRUE;
    pthread_cond_broadcast(&p_pool->work_available);
    pthread_mutex_unlock(&p_pool->mutex);
    for (uint32 i = 0; i < created; i++)
      pthread_join(p_pool->p_threads[i], NULL);
    pthread_cond_destroy(&p_pool->idle);
    pthread_cond_destroy(&p_pool->work_available);
    pthread_mutex_destroy(&p_pool->mutex);
    free(p_pool->p_threads);
    free(p_pool->p_queue);
    free(p_pool);
    return NULL;
  }
  return p_pool;
}

Seobeo_Worker_Result Seobeo_Worker_Pool_Submit(
    Seobeo_Worker_Pool *p_pool,
    Seobeo_Work_Function function,
    void *p_context,
    Seobeo_Work_Cleanup cleanup)
{
  if (!p_pool || !function)
    return SEOBEO_WORKER_INVALID_ARGUMENT;

  pthread_mutex_lock(&p_pool->mutex);
  if (!p_pool->accepting)
  {
    pthread_mutex_unlock(&p_pool->mutex);
    return SEOBEO_WORKER_STOPPED;
  }
  if (p_pool->queue_count == p_pool->queue_capacity)
  {
    pthread_mutex_unlock(&p_pool->mutex);
    return SEOBEO_WORKER_QUEUE_FULL;
  }

  uint32 tail =
      (p_pool->queue_head + p_pool->queue_count) %
      p_pool->queue_capacity;
  p_pool->p_queue[tail] = (Seobeo_Work_Item){
    .function = function,
    .p_context = p_context,
    .cleanup = cleanup,
  };
  p_pool->queue_count++;
  pthread_cond_signal(&p_pool->work_available);
  pthread_mutex_unlock(&p_pool->mutex);
  return SEOBEO_WORKER_OK;
}

Seobeo_Worker_Result Seobeo_Worker_Pool_Wait(
    Seobeo_Worker_Pool *p_pool)
{
  if (!p_pool ||
      Seobeo_Worker_Pool_Is_Current_Thread(p_pool) ||
      g_current_worker_pool == p_pool)
    return SEOBEO_WORKER_INVALID_ARGUMENT;

  pthread_mutex_lock(&p_pool->mutex);
  while (p_pool->queue_count > 0 ||
         p_pool->active_count > 0 ||
         p_pool->cleanup_count > 0)
    pthread_cond_wait(&p_pool->idle, &p_pool->mutex);
  pthread_mutex_unlock(&p_pool->mutex);
  return SEOBEO_WORKER_OK;
}

Seobeo_Worker_Result Seobeo_Worker_Pool_Shutdown(
    Seobeo_Worker_Pool *p_pool,
    boolean drain)
{
  if (!p_pool ||
      Seobeo_Worker_Pool_Is_Current_Thread(p_pool) ||
      g_current_worker_pool == p_pool)
    return SEOBEO_WORKER_INVALID_ARGUMENT;

  pthread_mutex_lock(&p_pool->mutex);
  if (p_pool->joined)
  {
    pthread_mutex_unlock(&p_pool->mutex);
    return SEOBEO_WORKER_OK;
  }
  if (p_pool->shutdown_started)
  {
    while (!p_pool->joined)
      pthread_cond_wait(&p_pool->idle, &p_pool->mutex);
    pthread_mutex_unlock(&p_pool->mutex);
    return SEOBEO_WORKER_OK;
  }
  p_pool->shutdown_started = TRUE;
  p_pool->accepting = FALSE;
  p_pool->stopping = TRUE;
  p_pool->drain = drain;

  if (!drain)
  {
    while (p_pool->queue_count > 0)
    {
      Seobeo_Work_Item item = p_pool->p_queue[p_pool->queue_head];
      p_pool->queue_head =
          (p_pool->queue_head + 1) % p_pool->queue_capacity;
      p_pool->queue_count--;
      p_pool->cleanup_count++;
      pthread_mutex_unlock(&p_pool->mutex);
      if (item.cleanup)
      {
        Seobeo_Worker_Pool *p_previous_pool = g_current_worker_pool;
        g_current_worker_pool = p_pool;
        item.cleanup(item.p_context);
        g_current_worker_pool = p_previous_pool;
      }
      pthread_mutex_lock(&p_pool->mutex);
      p_pool->cleanup_count--;
      if (p_pool->queue_count == 0 &&
          p_pool->active_count == 0 &&
          p_pool->cleanup_count == 0)
        pthread_cond_broadcast(&p_pool->idle);
    }
  }
  pthread_cond_broadcast(&p_pool->work_available);
  pthread_mutex_unlock(&p_pool->mutex);

  for (uint32 i = 0; i < p_pool->worker_count; i++)
  {
    if (pthread_join(p_pool->p_threads[i], NULL) != 0)
      return SEOBEO_WORKER_THREAD_ERROR;
  }

  pthread_mutex_lock(&p_pool->mutex);
  p_pool->joined = TRUE;
  if (p_pool->active_count == 0 &&
      p_pool->cleanup_count == 0)
    pthread_cond_broadcast(&p_pool->idle);
  pthread_mutex_unlock(&p_pool->mutex);
  return SEOBEO_WORKER_OK;
}

void Seobeo_Worker_Pool_Destroy(Seobeo_Worker_Pool *p_pool)
{
  if (!p_pool)
    return;
  if (Seobeo_Worker_Pool_Is_Current_Thread(p_pool) ||
      g_current_worker_pool == p_pool)
    return;
  if (!p_pool->joined)
    (void)Seobeo_Worker_Pool_Shutdown(p_pool, TRUE);
  pthread_cond_destroy(&p_pool->idle);
  pthread_cond_destroy(&p_pool->work_available);
  pthread_mutex_destroy(&p_pool->mutex);
  free(p_pool->p_threads);
  free(p_pool->p_queue);
  free(p_pool);
}

uint32 Seobeo_Worker_Pool_Pending(Seobeo_Worker_Pool *p_pool)
{
  if (!p_pool)
    return 0;
  pthread_mutex_lock(&p_pool->mutex);
  uint32 count = p_pool->queue_count;
  pthread_mutex_unlock(&p_pool->mutex);
  return count;
}

uint32 Seobeo_Worker_Pool_Active(Seobeo_Worker_Pool *p_pool)
{
  if (!p_pool)
    return 0;
  pthread_mutex_lock(&p_pool->mutex);
  uint32 count = p_pool->active_count;
  pthread_mutex_unlock(&p_pool->mutex);
  return count;
}