#include "engine/search.h"

#include <assert.h>

#include "engine/evaluate.h"
#include "bitboard.h"

// stop the search early? (time limit reached, node budget exhausted by the UCI
// thread, or an external stop signal).
static bool search_should_stop(search_ctx_t *ctx) {
  if (ctx == NULL) return false;
  if (ctx->stop != NULL && atomic_load(ctx->stop)) return true;
  if (ctx->go_args != NULL && ctx->go_args->nodes != 0 &&
      ctx->nodes != NULL && *ctx->nodes >= (u64)ctx->go_args->nodes) {
    return true;
  }
  return false;
}

// material value of a captured piece, used for delta pruning in qsearch.
static i32 piece_capture_value(position_t position, square_t to_sq) {
  for (int p = 0; p < PIECE_TYPE_COUNT; p++) {
    if (occupied_by(position.bitboards[p], to_sq)) {
      switch (p) {
        case WHITE_PAWN: case BLACK_PAWN:   return 100;
        case WHITE_KNIGHT: case BLACK_KNIGHT: return 320;
        case WHITE_BISHOP: case BLACK_BISHOP: return 320;
        case WHITE_ROOK: case BLACK_ROOK:   return 500;
        case WHITE_QUEEN: case BLACK_QUEEN:  return 900;
        default: return 0;
      }
    }
  }
  return 0;
}

// quiescence search: only captures and promotions, until stand-pat or no
// capture improves alpha.  Avoids the horizon effect at depth 0.
static search_result_t qsearch(
  position_t position,
  int alpha,
  int beta,
  search_ctx_t *ctx
) {
  if (ctx != NULL && ctx->nodes != NULL) (*ctx->nodes)++;
  if (search_should_stop(ctx)) {
    return (search_result_t) { .score = 0, .best_move = (move_t){0} };
  }

  i32 stand_pat = evaluate(position);
  if (stand_pat >= beta) {
    return (search_result_t) { .score = beta, .best_move = (move_t){0} };
  }
  if (stand_pat > alpha) alpha = stand_pat;

  // generate only captures (and promotions)
  moves_t moves;
  moves_init(&moves);
  get_legal_moves(&moves, position);

  // filter to captures/promotions only
  move_t cap_moves[64];
  int cap_count = 0;
  for (u32 i = 0; i < moves.length; i++) {
    move_t m = moves.moves[i];
    bool is_capture = (m.flags & MOVE_CAPTURE) || (m.flags & MOVE_EN_PASSANT);
    bool is_promo = m.flags & (MOVE_PROMOTE_Q | MOVE_PROMOTE_R | MOVE_PROMOTE_B | MOVE_PROMOTE_N);
    if (is_capture || is_promo) {
      cap_moves[cap_count++] = m;
    }
  }
  moves_deinit(moves);

  if (cap_count == 0) {
    return (search_result_t) { .score = stand_pat, .best_move = (move_t){0} };
  }

  search_result_t result = { .score = stand_pat, .best_move = cap_moves[0] };

  for (int i = 0; i < cap_count; i++) {
    move_t m = cap_moves[i];

    // delta pruning: if the capture can't possibly raise alpha, skip it.
    // (largest possible gain is queen ~900; add a safety margin.)
    i32 gain = piece_capture_value(position, m.to);
    if (stand_pat + gain + 200 < alpha) continue;

    position_t child = position;
    position_make_move(&child, m);
    search_result_t child_res = qsearch(child, -beta, -alpha, ctx);
    i32 score = -child_res.score;

    if (score > result.score) {
      result.score = score;
      result.best_move = m;
    }
    if (score > alpha) alpha = score;
    if (alpha >= beta) break;
  }

  return result;
}

search_result_t search(
  position_t position,
  int depth,
  int alpha,
  int beta,
  search_ctx_t *ctx
) {
  if (ctx != NULL && ctx->nodes != NULL) (*ctx->nodes)++;
  if (search_should_stop(ctx)) {
    // TODO: this might be a problem, if all the paths are losing,
    // and it reaches this condition. then it will suddenly think
    // WOW, this move will equialise the position!!
    // Which, obviously it didn't, maybe returning a score of relative negative
    // infinite will help
    return (search_result_t) { .score = 0, .best_move = (move_t){0} };
  }

  if (depth <= 0) {
    return qsearch(position, alpha, beta, ctx);
  }

  moves_t moves;
  moves_init(&moves);
  get_legal_moves(&moves, position);

  // terminal node: no legal moves -> checkmate or stalemate.
  if (moves.length == 0) {
    moves_deinit(moves);
    u8 king = position.turn == WHITE_TURN ? WHITE_KING : BLACK_KING;
    int king_sq = __builtin_ctzll(position.bitboards[king]);
    u8 opponent = position.turn == WHITE_TURN ? BLACK_TURN : WHITE_TURN;
    if (square_attacked(position, king_sq, opponent)) {
      return (search_result_t) {
        .score = -SEARCH_MATE,
        .best_move = (move_t){0},
      };
    }
    return (search_result_t) { .score = 0, .best_move = (move_t){0} };
  }

  // transposition-table probe.
  u64 key = 0;
  tentry_t *tt_entry = NULL;
  bool have_tt = ctx != NULL && ctx->ttable != NULL;
  if (have_tt) {
    key = zobrist_hash(position);
    tt_entry = ttable_find(*ctx->ttable, key);
  }

  search_result_t result = { .score = -SEARCH_INF, .best_move = moves.moves[0] };
  int orig_alpha = alpha;

  if (tt_entry != NULL && tt_entry->depth >= (u16)depth) {
    i16 tt_score = tt_entry->score;
    if (tt_entry->flag == TT_EXACT) {
      result.score = tt_score;
      result.best_move = tt_entry->best_move;
      moves_deinit(moves);
      return result;
    }
    if (tt_entry->flag == TT_LOWERBOUND && tt_score > alpha) alpha = tt_score;
    else if (tt_entry->flag == TT_UPPERBOUND && tt_score < beta) beta = tt_score;
    if (alpha >= beta) {
      result.score = tt_score;
      result.best_move = tt_entry->best_move;
      moves_deinit(moves);
      return result;
    }
  }

  for (u32 i = 0; i < moves.length; i++) {
    position_t child = position;
    position_make_move(&child, moves.moves[i]);

    search_result_t child_res = search(child, depth - 1, -beta, -alpha, ctx);
    i32 score = -child_res.score;

    if (score > result.score) {
      result.score = score;
      result.best_move = moves.moves[i];
    }
    if (score > alpha) alpha = score;
    if (alpha >= beta) break;
  }

  // transposition-table store.
  if (have_tt) {
    u8 flag = (result.score >= beta) ? TT_LOWERBOUND
              : (result.score <= orig_alpha) ? TT_UPPERBOUND
              : TT_EXACT;
    ttable_insert(*ctx->ttable,
      .key = key,
      .move = result.best_move,
      .score = (i16)result.score,
      .eval = 0,
      .depth = (u16)depth,
      .flag = flag,
      .generation = 0
    );
  }

  moves_deinit(moves);
  return result;
}

// Reconstruct the principal variation by following TT best-move entries.
int reconstruct_pv(
  position_t position,
  ttable_t tt,
  move_t *pv_out,
  int max_len
) {
  int count = 0;
  position_t pos = position;
  // guard against cycles: at most 128 plies
  for (int i = 0; i < max_len && i < 128; i++) {
    u64 key = zobrist_hash(pos);
    tentry_t *entry = ttable_find(tt, key);
    if (entry == NULL) break;

    move_t m = entry->best_move;
    if (m.from == 0 && m.to == 0) break;

    pv_out[count++] = m;
    position_make_move(&pos, m);
  }
  return count;
}

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

#include "fen.h"

bool test_search_finds_mate_in_one() {
  // Ka6, Qb6 vs Ka8 -> Qa7# or Qb7#
  struct fen_load r = load_fen("k7/8/KQ6/8/8/8/8/8 w - - 0 1");
  if (r.failed) return false;

  ttable_t tt;
  ttable_init(&tt, 2);

  _Atomic bool stop = false;
  go_args_t go = {0};
  u64 nodes = 0;
  search_ctx_t ctx = { .ttable = &tt, .go_args = &go, .stop = &stop, .nodes = &nodes };

  search_result_t res = search(r.position, 2, -SEARCH_INF, SEARCH_INF, &ctx);
  ttable_deinit(tt);

  // the queen delivers mate, so the value at the root must be a (positive) mate.
  return res.score == SEARCH_MATE;
}

bool test_search_handles_no_moves() {
  // "7k/5Q2/6K1/8/8/8/8/8 b - - 0 1" — black to move, stalemated, not in check.
  struct fen_load r = load_fen("7k/5Q2/6K1/8/8/8/8/8 b - - 0 1");
  if (r.failed) return false;

  ttable_t tt;
  ttable_init(&tt, 2);

  _Atomic bool stop = false;
  go_args_t go = {0};
  u64 nodes = 0;
  search_ctx_t ctx = { .ttable = &tt, .go_args = &go, .stop = &stop, .nodes = &nodes };

  search_result_t res = search(r.position, 2, -SEARCH_INF, SEARCH_INF, &ctx);
  ttable_deinit(tt);

  // no legal moves + not in check -> stalemate, value 0 (not a mate).
  return res.score == 0;
}

bool test_search_respects_depth() {
  position_t pos = position_starting();
  ttable_t tt;
  ttable_init(&tt, 2);

  _Atomic bool stop = false;
  go_args_t go = {0};
  u64 nodes = 0;
  search_ctx_t ctx = { .ttable = &tt, .go_args = &go, .stop = &stop, .nodes = &nodes };

  search_result_t res = search(pos, 3, -SEARCH_INF, SEARCH_INF, &ctx);
  ttable_deinit(tt);

  // must return a legal on-board move and a finite (non-sentinel) score.
  bool on_board = res.best_move.from < 64 && res.best_move.to < 64;
  return on_board && res.score > -SEARCH_INF && res.score < SEARCH_INF;
}

bool test_search_stop_signal_halts() {
  position_t pos = position_starting();
  ttable_t tt;
  ttable_init(&tt, 2);

  _Atomic bool stop = false;
  go_args_t go = {0};
  go.nodes = 1;   // node budget of 1 -> should stop almost immediately
  u64 nodes = 0;
  search_ctx_t ctx = { .ttable = &tt, .go_args = &go, .stop = &stop, .nodes = &nodes };

  search_result_t res = search(pos, 6, -SEARCH_INF, SEARCH_INF, &ctx);
  ttable_deinit(tt);

  // with a 1-node budget the search must bail early, not run to depth 6.
  return nodes <= 2 && res.score == 0;
}

bool test_qsearch_finds_capture() {
  // position where white can capture a queen: white pawn on d4, black queen on e5.
  struct fen_load r = load_fen("rnb1kbnr/pppp1ppp/8/4q3/3P4/8/PPP1PPPP/RNBQKBNR w KQkq - 0 1");
  if (r.failed) return false;

  ttable_t tt;
  ttable_init(&tt, 2);

  _Atomic bool stop = false;
  go_args_t go = {0};
  u64 nodes = 0;
  search_ctx_t ctx = { .ttable = &tt, .go_args = &go, .stop = &stop, .nodes = &nodes };

  // depth 1: the search should see the dxe5 capture via qsearch
  search_result_t res = search(r.position, 1, -SEARCH_INF, SEARCH_INF, &ctx);
  ttable_deinit(tt);

  // dxe5 captures the queen, so the best move must be d4->e5
  bool found_capture = (res.best_move.from == 35 && res.best_move.to == 44); // d4=35, e5=44
  // even if the move ordering differs, the score should be strongly positive
  return res.score > 500;
}

bool test_pv_reconstruction() {
  position_t pos = position_starting();
  ttable_t tt;
  ttable_init(&tt, 4);

  _Atomic bool stop = false;
  go_args_t go = {0};
  u64 nodes = 0;
  search_ctx_t ctx = { .ttable = &tt, .go_args = &go, .stop = &stop, .nodes = &nodes };

  // run a shallow search to populate the TT
  search(pos, 4, -SEARCH_INF, SEARCH_INF, &ctx);

  // reconstruct PV from the TT
  move_t pv[32];
  int pv_len = reconstruct_pv(pos, tt, pv, 32);
  ttable_deinit(tt);

  // PV should have at least 1 move (the root best move)
  if (pv_len < 1) return false;
  // every PV move should be on-board
  for (int i = 0; i < pv_len; i++) {
    if (pv[i].from >= 64 || pv[i].to >= 64) return false;
  }
  return true;
}
#endif