#define ENGINE_THREADS_INTERNAL
#include "engine/threads.h"
#include "engine/search.h"
#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include <stdio.h>

void threads_init(struct threads *handler, size_t capacity) {
  // maybe this will be a speed up?
  // why am i micro optimizing here??????
  // idk i just watched the eskil steenberg video about UB
  memset(handler, 0, sizeof(*handler));

  handler->threads = malloc(sizeof(*handler->threads) * capacity);
  handler->args = malloc(sizeof(*handler->args) * capacity);

  // split point pool (DTS §2.4): zeroed so every block starts SPLITPOINT_STATE_FREE
  handler->split_point_pool = calloc(sizeof(split_point_t), capacity);

  handler->capacity = capacity;
  handler->count = 0;
  handler->running_count = 0;
  handler->engine_messages = NULL;
  handler->go_args = NULL;
  handler->position = NULL;
}

void threads_deinit(struct threads *handler) {
  assert(handler->running_count == 0);
  free(handler->threads);
  free(handler->args);
  free(handler->split_point_pool);
}

void threads_reserve(struct threads *handler, size_t capacity) {
  if (handler->capacity >= capacity) return;

  handler->threads = realloc(
    handler->threads,
    sizeof(*handler->threads) * capacity
  );
  handler->args = realloc(
    handler->args,
    sizeof(*handler->args) * capacity
  );
  handler->split_point_pool = realloc(
    handler->split_point_pool,
    sizeof(split_point_t) * capacity
  );
  // new blocks must start FREE
  memset(
    (char*)handler->split_point_pool + handler->capacity * sizeof(split_point_t),
    0,
    (capacity - handler->capacity) * sizeof(split_point_t)
  );
  handler->capacity = capacity;
}

bool threads_go_loop(struct threads *handler) {
  return threads_all_done(handler);
}

bool threads_all_done(struct threads *handler) {
  int running_count = handler->running_count;
  assert(running_count > 0);
  for (int i = 0; i < running_count; i++) {
    // `finished` is part of the common initial sequence of both union
    // members (stop, finished), so reading it via either member is well-defined.
    if (atomic_load(&handler->args[i].root_thread_args.finished) == 1) continue;
    return false;
  }
  return true;
}

void threads_cleanup(struct threads *handler) {
  int running_count = handler->running_count;
  for (int i = 0; i < running_count; i++) {
    __threads_args_init(handler->args + i, handler, i);
    pthread_cancel(handler->threads[i]);
  }
  handler->running_count = 0;
}

void __threads_args_init(
  thread_args_t* _args,
  struct threads *handler,
  int i
) {
  if (i > 0) {
    struct child_thread_args *args = (struct child_thread_args*)_args;
    atomic_init(&args->stop, 0);
    atomic_init(&args->finished, 0);
    args->id = i;
    args->position = handler->position;
    args->go_args = handler->go_args;
    args->ttable = &handler->ttable;
    args->split_points = handler->split_point_pool;
    args->sp_pool_capacity = handler->capacity;
    args->message = handler->engine_messages->data[i];
    return;
  }

  struct root_thread_args *args = (struct root_thread_args*)_args;
  atomic_init(&args->stop, 0);
  atomic_init(&args->finished, 0);
  args->id = i;
  args->position = handler->position;
  args->go_args = handler->go_args;
  args->message = handler->engine_messages->data[i];
  args->ttable = &handler->ttable;
}

split_point_t *split_point_alloc(struct threads *handler) {
  for (size_t i = 0; i < handler->capacity; i++) {
    split_point_t *sp = (split_point_t*)handler->split_point_pool + i;
    int expected = SPLITPOINT_STATE_FREE;
    if (!atomic_compare_exchange_weak(&sp->state, &expected, SPLITPOINT_STATE_ACTIVE)) continue;

    atomic_store(&sp->next_move_idx, 0);
    atomic_store(&sp->alpha, 0);
    atomic_store(&sp->beta, 0);
    atomic_store(&sp->best_score, -SEARCH_INF);
    // workers counts only registered helpers; the creator registers itself
    atomic_store(&sp->workers, 0);
    sp->best_move = (move_t) { 0 };
    return sp;
  }
  return NULL;
}

void split_point_free(struct threads *handler, split_point_t *split_point) {
  // Unsplit (§2.7): drop the block back into the pool
  atomic_store(&split_point->state, SPLITPOINT_STATE_FREE);
}

// HELP (§2.3): a worker publishes its block so idle peers can pick up moves
void split_point_publish(split_point_t *split_point) {
  atomic_store(&split_point->state, SPLITPOINT_STATE_ACTIVE);
}

// Select (§2.5): lock-free claim of the next move via the shared counter
bool split_point_claim_move(split_point_t *split_point, move_t *move) {
  int idx = atomic_fetch_add(&split_point->next_move_idx, 1);
  if (idx >= (int)split_point->moves.length) return false;
  *move = split_point->moves.moves[idx];
  return true;
}

// Merge (§2.6): decrement workers; returns true when the caller was the last
// worker and now owns the node
bool split_point_merge(split_point_t *split_point) {
  return atomic_fetch_sub(&split_point->workers, 1) == 1;
}

// Share (§2.8): back up an improved best score; peers see it on their next Select
void split_point_share(split_point_t *split_point, i32 score, move_t move) {
  int current = atomic_load(&split_point->best_score);
  while (score > current) {
    if (atomic_compare_exchange_weak(&split_point->best_score, &current, score)) {
      split_point->best_move = move;
      break;
    }
  }
}

split_point_t *__pool_find_work(split_point_t *pool, size_t capacity) {
  for (size_t i = 0; i < capacity; i++) {
    split_point_t *sp = pool + i;
    if (atomic_load(&sp->state) != SPLITPOINT_STATE_ACTIVE) continue;
    if (atomic_load(&sp->next_move_idx) >= (int)sp->moves.length) continue;
    return sp;
  }
  return NULL;
}

split_point_t *threads_find_work(struct threads *handler) {
  return __pool_find_work(handler->split_point_pool, handler->capacity);
}

#ifdef TEST_MOD
bool test_threads_args_init_root_propagates() {
  engine_message_handler_t dummy[2];
  engine_message_handler_t *dummy_ptrs[2] = { &dummy[0], &dummy[1] };
  engine_messages_t messages = {
    .data = dummy_ptrs,
    .count = 2,
    .capacity = 2,
  };

  go_args_t go_args = {0};
  position_t pos = {0};

  struct threads handler = {0};
  handler.engine_messages = &messages;
  handler.go_args = &go_args;
  handler.position = &pos;
  ttable_init(&handler.ttable, 2);

  thread_args_t args;
  __threads_args_init(&args, &handler, 0);

  bool ok = true;
  ok &= args.root_thread_args.id == 0;
  ok &= atomic_load(&args.root_thread_args.stop) == 0;
  ok &= atomic_load(&args.root_thread_args.finished) == 0;
  ok &= args.root_thread_args.ttable == &handler.ttable;
  ok &= args.root_thread_args.go_args == &go_args;
  ok &= args.root_thread_args.position == &pos;
  ok &= args.root_thread_args.message == &dummy[0];

  ttable_deinit(handler.ttable);
  return ok;
}

bool test_threads_args_init_child_propagates() {
  engine_message_handler_t dummy[2];
  engine_message_handler_t *dummy_ptrs[2] = { &dummy[0], &dummy[1] };
  engine_messages_t messages = {
    .data = dummy_ptrs,
    .count = 2,
    .capacity = 2,
  };

  go_args_t go_args = {0};
  position_t pos = {0};

  struct threads handler = {0};
  handler.engine_messages = &messages;
  handler.go_args = &go_args;
  handler.position = &pos;
  ttable_init(&handler.ttable, 2);

  thread_args_t args;
  __threads_args_init(&args, &handler, 1);

  bool ok = true;
  ok &= args.child_thread_args.id == 1;
  ok &= args.child_thread_args.ttable == &handler.ttable;
  ok &= args.child_thread_args.go_args == &go_args;
  ok &= args.child_thread_args.position == &pos;
  ok &= args.child_thread_args.message == &dummy[1];
  ok &= args.child_thread_args.split_points == handler.split_point_pool;
  ok &= args.child_thread_args.sp_pool_capacity == handler.capacity;

  ttable_deinit(handler.ttable);
  return ok;
}

// ---- Phase D tests: split point pool + parallel primitives ----

static void __test_fill_moves(split_point_t *sp, int count) {
  moves_init_wcapacity(&sp->moves, count);
  for (int i = 0; i < count; i++) {
    add_move(&sp->moves, (square_t)i, (square_t)(i + 8));
  }
}

bool test_split_pool_alloc_free() {
  struct threads handler;
  threads_init(&handler, 4);

  bool ok = true;

  // pool starts all FREE; we can alloc exactly capacity blocks
  split_point_t *sps[4];
  for (int i = 0; i < 4; i++) {
    sps[i] = split_point_alloc(&handler);
    if (sps[i] == NULL) ok = false;
  }
  ok &= split_point_alloc(&handler) == NULL; // exhausted

  // freed block becomes reusable
  split_point_free(&handler, sps[2]);
  split_point_t *reused = split_point_alloc(&handler);
  ok &= reused == sps[2];

  threads_deinit(&handler);
  return ok;
}

typedef struct {
  split_point_t *sp;
  int claims;
} __claim_thread_args_t;

static void *__claim_worker(void *_args) {
  __claim_thread_args_t *args = (__claim_thread_args_t*)_args;
  args->claims = 0;
  move_t move;
  while (split_point_claim_move(args->sp, &move)) args->claims++;
  return NULL;
}

bool test_split_claim_moves_exactly_once() {
  enum { MOVE_COUNT = 100, THREAD_COUNT = 4 };
  struct threads handler;
  threads_init(&handler, THREAD_COUNT);

  split_point_t *sp = split_point_alloc(&handler);
  atomic_store(&sp->alpha, -SEARCH_INF);
  atomic_store(&sp->beta, SEARCH_INF);
  __test_fill_moves(sp, MOVE_COUNT);

  pthread_t threads[THREAD_COUNT];
  __claim_thread_args_t thread_args[THREAD_COUNT];
  for (int i = 0; i < THREAD_COUNT; i++) {
    thread_args[i].sp = sp;
    pthread_create(threads + i, NULL, __claim_worker, thread_args + i);
  }
  int total = 0;
  for (int i = 0; i < THREAD_COUNT; i++) {
    pthread_join(threads[i], NULL);
    total += thread_args[i].claims;
  }

  bool ok = total == MOVE_COUNT; // each move claimed exactly once
  ok &= !split_point_claim_move(sp, &(move_t){0}); // exhausted

  moves_deinit(sp->moves);
  threads_deinit(&handler);
  return ok;
}

bool test_split_merge_last_owns() {
  split_point_t sp;
  atomic_init(&sp.workers, 3);

  bool ok = true;
  ok &= !split_point_merge(&sp); // 3 -> 2, not last
  ok &= !split_point_merge(&sp); // 2 -> 1, not last
  ok &= split_point_merge(&sp);  // 1 -> 0, caller owns the node
  return ok;
}

bool test_split_share_propagates_best() {
  split_point_t sp;
  atomic_init(&sp.best_score, -SEARCH_INF);
  sp.best_move = (move_t) { 0 };

  bool ok = true;

  split_point_share(&sp, 50, (move_t) { 1, 2 });
  ok &= atomic_load(&sp.best_score) == 50;
  ok &= sp.best_move.from == 1 && sp.best_move.to == 2;

  // a worse score must not overwrite the shared best
  split_point_share(&sp, 10, (move_t) { 3, 4 });
  ok &= atomic_load(&sp.best_score) == 50;
  ok &= sp.best_move.from == 1 && sp.best_move.to == 2;

  split_point_share(&sp, 90, (move_t) { 5, 6 });
  ok &= atomic_load(&sp.best_score) == 90;
  ok &= sp.best_move.from == 5 && sp.best_move.to == 6;
  return ok;
}

bool test_help_scan_finds_published_work() {
  struct threads handler;
  threads_init(&handler, 4);

  bool ok = true;

  // empty pool -> no work
  ok &= threads_find_work(&handler) == NULL;

  // a published block with moves left is found
  split_point_t *working = split_point_alloc(&handler);
  __test_fill_moves(working, 4);
  split_point_publish(working);
  ok &= threads_find_work(&handler) == working;

  // an exhausted block is skipped
  atomic_store(&working->next_move_idx, 4);
  ok &= threads_find_work(&handler) == NULL;

  // a free block is skipped even if republished with moves
  split_point_t *freed = split_point_alloc(&handler);
  __test_fill_moves(freed, 2);
  split_point_publish(freed);
  split_point_free(&handler, freed);
  ok &= threads_find_work(&handler) == NULL;

  moves_deinit(working->moves);
  moves_deinit(freed->moves);
  threads_deinit(&handler);
  return ok;
}
#endif