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

#define ROOT_MAX_DEPTH 100
#define MATE_WINDOW 1000

// convert a mate score into a "mate in N plies" distance (negative when we
// are getting mated), 0 when the score isn't near a mate.
static int __score_to_mate(i32 score) {
  int abs_score = score < 0 ? -score : score;
  if (abs_score < SEARCH_MATE - MATE_WINDOW) return 0;
  int plies = SEARCH_MATE - abs_score + 1;
  int mate_in = (plies + 1) / 2;
  return score > 0 ? mate_in : -mate_in;
}

void threads_go(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_create(
      handler->threads + i,
      NULL,
      i == 0 ? root_thread : child_thread,
      handler->args + i
    );
  }
}

void *child_thread(void *_args) {
  struct child_thread_args *args = (struct child_thread_args*)_args;
  args->stop = 0;

  atomic_bool stop = false;
  u64 nodes = 0;
  search_ctx_t ctx = {
    .ttable = args->ttable,
    .go_args = args->go_args,
    .stop = &stop,
    .nodes = &nodes,
  };

  engine_message_t *message = args->message->count > 0
    ? args->message->newest_valid
    : engine_message_handler_push(args->message);
  message->multipv = args->id;

  i32 best_score = -SEARCH_INF;
  move_t best_move = (move_t) { 0 };

  // DTS worker (§2.3–2.8): HELP peers by claiming moves from published
  // split points, searching them, and Sharing the result back.
  while (!atomic_load(&args->stop) && !atomic_load(&stop)) {
    split_point_t *sp = __pool_find_work(
      args->split_points,
      args->sp_pool_capacity
    );

    // nothing active anywhere: the pool is drained, so we are done helping
    // (the root re-publishes per iteration once root splitting lands)
    if (sp == NULL) break;

    // register as a worker at this block so Merge knows the headcount
    atomic_fetch_add(&sp->workers, 1);

    int alpha, beta;
    move_t move;
    while (split_point_claim_move(sp, &move)) {
      alpha = -atomic_load(&sp->beta);
      beta = -atomic_load(&sp->alpha);

      position_t child = sp->position_snapshot;
      position_make_move(&child, move);

      search_result_t res = search(child, sp->depth - 1, alpha, beta, &ctx);

      // Share (§2.8): negamax — our value flips sign for the parent side
      split_point_share(sp, -res.score, move);
      if (-res.score > best_score) {
        best_score = -res.score;
        best_move = move;
      }

      message->depth = sp->depth;
      message->nodes = nodes > INT_MAX ? INT_MAX : (int)nodes;
      message->score_cp = best_score;
      message->pv.count = 0;
      add_ipc_move(&message->pv, (ipc_move_t) { best_move.from, best_move.to });
      engine_message_mark_ready(message);
    }

    // Merge (§2.6): when the last worker leaves, Unsplit the block
    if (split_point_merge(sp)) {
      atomic_store(&sp->state, SPLITPOINT_STATE_FREE);
    }
  }

  message->best_move = (ipc_move_t) { best_move.from, best_move.to };
  message->next = NULL;
  engine_message_mark_ready(message);

  atomic_store(&args->finished, 1);
  return NULL;
}

// DTS requires a root thread which is initialize the search tree, and the
// others who use that search tree. it's as if root_thread is the owner of the
// tree
void *root_thread(void *_args) {
  struct root_thread_args *args = (struct root_thread_args*)_args;
  args->stop = 0;

  atomic_bool stop = false;
  u64 nodes = 0;
  search_ctx_t ctx = {
    .ttable = args->ttable,
    .go_args = args->go_args,
    .stop = &stop,
    .nodes = &nodes,
  };

  int max_depth = args->go_args != NULL && args->go_args->depth > 0
    ? args->go_args->depth
    : ROOT_MAX_DEPTH;

  // reuse the handler's pre-created message when present, otherwise push one.
  engine_message_t *message = args->message->count > 0
    ? args->message->newest_valid
    : engine_message_handler_push(args->message);
  search_result_t result = { .score = 0, .best_move = (move_t) { 0 } };

  // iterative deepening: search at increasing depths, publishing an info
  // message per completed depth so the UCI thread can emit `info ... pv`.
  for (int depth = 1; depth <= max_depth; depth++) {
    result = search(*args->position, depth, -SEARCH_INF, SEARCH_INF, &ctx);

    if (atomic_load(&stop)) break;

    message->depth = depth;
    message->nodes = nodes > INT_MAX ? INT_MAX : (int)nodes;
    message->score_cp = result.score;
    message->mate = __score_to_mate(result.score);
    // pv is already initialized by engine_message_init / handler_push.
    message->pv.count = 0;
    add_ipc_move(&message->pv, (ipc_move_t) {
      result.best_move.from,
      result.best_move.to,
    });

    if (!should_continue(message, args->go_args)) {
      break;
    }

    engine_message_mark_ready(message);
    message = engine_message_handler_push(args->message);
  }

  // final message carries the best move of the last completed iteration.
  message->best_move = (ipc_move_t) { result.best_move.from, result.best_move.to };
  message->ponder = (ipc_move_t) { 0, 0 };
  message->next = NULL;
  engine_message_mark_ready(message);

  // signal children to stop their loops and mark this thread as finished.
  atomic_store(&args->stop, 1);
  atomic_store(&args->finished, 1);

  return NULL;
}

#ifdef TEST_MOD
#include <stdbool.h>
#include <stdlib.h>

#include "fen.h"

// the root_thread driver + tests need the IPC implementation (messages,
// moves, should_continue) which lives outside of engine.c's aggregation.
#include "../src/ipc/go_args.c"
#include "../src/ipc/moves.c"
#include "../src/ipc/engine_message.c"

static bool __root_thread_run(
  struct root_thread_args *args,
  go_args_t *go_args,
  position_t position,
  int ttable_log
) {
  args->message = malloc(sizeof(*args->message));
  engine_message_handler_init(args->message);
  args->go_args = go_args;
  args->position = &position;

  ttable_t tt;
  ttable_init(&tt, ttable_log);

  root_thread(args);

  bool ok = true;
  ok &= atomic_load(&args->finished) == 1;
  ok &= atomic_load(&args->stop) == 1;

  ttable_deinit(tt);
  return ok;
}

static engine_message_t *__root_thread_pop_all(
  engine_message_handler_t *handler,
  int *count_out
) {
  engine_message_t *head = NULL;
  engine_message_t *tail = NULL;
  int count = 0;
  while (engine_message_handler_can_pop(handler)) {
    engine_message_t *m = engine_message_handler_pop(handler);
    m->next = NULL;
    if (head == NULL) head = m;
    else tail->next = m;
    tail = m;
    count++;
  }
  if (count_out != NULL) *count_out = count;
  return head;
}

static void __root_thread_free_messages(engine_message_t *head) {
  while (head != NULL) {
    engine_message_t *next = head->next;
    engine_message_deinit(head);
    head = next;
  }
}

bool test_root_thread_emits_info_and_bestmove() {
  struct root_thread_args *args = calloc(sizeof(*args), 1);
  go_args_t go_args = { .depth = 3 };
  position_t position = position_starting();

  bool ok = __root_thread_run(
    args,
    &go_args,
    position,
    2
  );

  engine_message_handler_t *handler = args->message;
  int count = 0;
  engine_message_t *messages = __root_thread_pop_all(handler, &count);

  ok &= count >= 1;

  // every published message must have a valid depth and on-board pv move.
  for (engine_message_t *m = messages; m != NULL; m = m->next) {
    ok &= m->depth >= 1 && m->depth <= 3;
    ok &= m->pv.count == 1;
    ok &= m->pv.data[0].from >= 0 && m->pv.data[0].from < 64;
    ok &= m->pv.data[0].to >= 0 && m->pv.data[0].to < 64;
  }

  // the final message must carry the best move.
  engine_message_t *tail = messages;
  while (tail->next != NULL) tail = tail->next;
  ok &= tail->best_move.from >= 0 && tail->best_move.from < 64;
  ok &= tail->best_move.to >= 0 && tail->best_move.to < 64;
  ok &= tail->depth == 3;

  __root_thread_free_messages(messages);
  engine_message_handler_deinit(handler);
  free(handler);
  free(args);
  return ok;
}

bool test_root_thread_respects_depth_limit() {
  struct root_thread_args *args = calloc(sizeof(*args), 1);
  go_args_t go_args = { .depth = 2 };
  position_t position = position_starting();

  bool ok = __root_thread_run(args, &go_args, position, 2);

  int count = 0;
  engine_message_t *messages = __root_thread_pop_all(args->message, &count);

  // depth 2 -> exactly two info messages, depths 1 then 2.
  ok &= count == 2;
  ok &= messages->depth == 1;
  ok &= messages->next->depth == 2;

  __root_thread_free_messages(messages);
  engine_message_handler_deinit(args->message);
  free(args->message);
  free(args);
  return ok;
}

bool test_root_thread_stops_on_node_limit() {
  struct root_thread_args *args = calloc(sizeof(*args), 1);
  go_args_t go_args = { .nodes = 1 };
  position_t position = position_starting();

  bool ok = __root_thread_run(args, &go_args, position, 2);

  int count = 0;
  engine_message_t *messages = __root_thread_pop_all(args->message, &count);

  // node budget of 1 exhausts should_continue after the first iteration.
  ok &= count == 1;
  ok &= messages->node_limit;

  __root_thread_free_messages(messages);
  engine_message_handler_deinit(args->message);
  free(args->message);
  free(args);
  return ok;
}

bool test_root_thread_reports_mate() {
  struct root_thread_args *args = calloc(sizeof(*args), 1);
  go_args_t go_args = { .depth = 4 };
  struct fen_load r = load_fen("k7/8/KQ6/8/8/8/8/8 w - - 0 1");
  if (r.failed) return false;

  bool ok = __root_thread_run(args, &go_args, r.position, 2);

  int count = 0;
  engine_message_t *messages = __root_thread_pop_all(args->message, &count);

  // mate in 1 must be reported via the `mate` field.
  engine_message_t *tail = messages;
  while (tail->next != NULL) tail = tail->next;
  ok &= tail->mate == 1;
  ok &= tail->score_cp == SEARCH_MATE;

  __root_thread_free_messages(messages);
  engine_message_handler_deinit(args->message);
  free(args->message);
  free(args);
  return ok;
}

// ---- Phase E tests: DTS child worker loop ----

#include <pthread.h>

typedef struct {
  struct threads *handler;
  struct child_thread_args args[4];
  pthread_t threads[4];
} __child_test_env_t;

static void __child_test_setup(
  __child_test_env_t *env,
  split_point_t *sp,
  int worker_count
) {
  threads_init(env->handler, (size_t)worker_count);
  ttable_init(&env->handler->ttable, 2);

  for (int i = 0; i < worker_count; i++) {
    env->args[i].id = i + 1;
    atomic_init(&env->args[i].stop, 0);
    atomic_init(&env->args[i].finished, 0);
    env->args[i].go_args = NULL;
    env->args[i].ttable = &env->handler->ttable;
    env->args[i].split_points = env->handler->split_point_pool;
    env->args[i].sp_pool_capacity = env->handler->capacity;
    env->args[i].message = malloc(sizeof(engine_message_handler_t));
    engine_message_handler_init(env->args[i].message);
  }

  if (sp != NULL) {
    sp->moves.length = 0; // filled by caller
  }
}

static void __child_test_teardown(__child_test_env_t *env, int worker_count) {
  ttable_deinit(env->handler->ttable);
  for (int i = 0; i < worker_count; i++) {
    while (engine_message_handler_can_pop(env->args[i].message)) {
      engine_message_deinit(engine_message_handler_pop(env->args[i].message));
    }
    engine_message_handler_deinit(env->args[i].message);
    free(env->args[i].message);
  }
  threads_deinit(env->handler);
}

bool test_child_workers_claim_and_share() {
  enum { WORKERS = 4 };
  __child_test_env_t env;
  struct threads handler;

  env.handler = &handler;
  __child_test_setup(&env, NULL, WORKERS);

  split_point_t *sp = split_point_alloc(&handler);
  sp->depth = 3;
  sp->ply = 0;
  atomic_store(&sp->alpha, -SEARCH_INF);
  atomic_store(&sp->beta, SEARCH_INF);
  sp->position_snapshot = position_starting();
  moves_init(&sp->moves);
  get_legal_moves(&sp->moves, sp->position_snapshot);
  split_point_publish(sp);

  for (int i = 0; i < WORKERS; i++) {
    pthread_create(env.threads + i, NULL, child_thread, env.args + i);
  }
  for (int i = 0; i < WORKERS; i++) {
    pthread_join(env.threads[i], NULL);
  }

  bool ok = true;

  // every worker finished cleanly
  for (int i = 0; i < WORKERS; i++) {
    ok &= atomic_load(&env.args[i].finished) == 1;
    ok &= atomic_load(&env.args[i].message->newest_valid->ready) == 1;
  }

  // the whole move list was consumed and the block was Unsplitten
  ok &= atomic_load(&sp->next_move_idx) >= (int)sp->moves.length;
  ok &= atomic_load(&sp->state) == SPLITPOINT_STATE_FREE;
  ok &= atomic_load(&sp->workers) == 0;

  // real search happened: nodes were counted and a sane best was shared
  ok &= atomic_load(&sp->best_score) > -SEARCH_INF;
  ok &= atomic_load(&sp->best_score) <= SEARCH_MATE;

  moves_deinit(sp->moves);
  __child_test_teardown(&env, WORKERS);
  return ok;
}

bool test_child_thread_stops_on_signal() {
  __child_test_env_t env;
  struct threads handler;
  enum { WORKERS = 1 };

  env.handler = &handler;
  __child_test_setup(&env, NULL, WORKERS);

  split_point_t *sp = split_point_alloc(&handler);
  sp->depth = 5;
  atomic_store(&sp->alpha, -SEARCH_INF);
  atomic_store(&sp->beta, SEARCH_INF);
  sp->position_snapshot = position_starting();
  moves_init(&sp->moves);
  get_legal_moves(&sp->moves, sp->position_snapshot);
  split_point_publish(sp);

  pthread_create(env.threads, NULL, child_thread, env.args);

  // wait until the worker has started claiming, then signal stop
  while (atomic_load(&sp->next_move_idx) == 0) sched_yield();
  atomic_store(&env.args[0].stop, 1);
  pthread_join(env.threads[0], NULL);

  bool ok = true;
  ok &= atomic_load(&env.args[0].finished) == 1;
  ok &= atomic_load(&env.args[0].message->newest_valid->ready) == 1;

  moves_deinit(sp->moves);
  __child_test_teardown(&env, WORKERS);
  return ok;
}

bool test_child_workers_find_mate_score() {
  // Ka6/Qb6 vs Ka8: white mates in 1 with Qa7#/Qb7#. Children searching the
  // mate move must Share a mate score back to the split point.
  enum { WORKERS = 2 };
  __child_test_env_t env;
  struct threads handler;

  env.handler = &handler;
  __child_test_setup(&env, NULL, WORKERS);

  struct fen_load r = load_fen("k7/8/KQ6/8/8/8/8/8 w - - 0 1");
  if (r.failed) return false;

  split_point_t *sp = split_point_alloc(&handler);
  sp->depth = 2;
  atomic_store(&sp->alpha, -SEARCH_INF);
  atomic_store(&sp->beta, SEARCH_INF);
  sp->position_snapshot = r.position;
  moves_init(&sp->moves);
  get_legal_moves(&sp->moves, r.position);
  split_point_publish(sp);

  for (int i = 0; i < WORKERS; i++) {
    pthread_create(env.threads + i, NULL, child_thread, env.args + i);
  }
  for (int i = 0; i < WORKERS; i++) {
    pthread_join(env.threads[i], NULL);
  }

  bool ok = true;
  ok &= atomic_load(&sp->best_score) == SEARCH_MATE;
  ok &= sp->best_move.from >= 0 && sp->best_move.from < 64;
  ok &= sp->best_move.to >= 0 && sp->best_move.to < 64;

  moves_deinit(sp->moves);
  __child_test_teardown(&env, WORKERS);
  return ok;
}
#endif