#include "bot.hpp"
#include "board/board.hpp"
#include "evaluate.hpp"
#include "moves.hpp"
#include "zobrist.hpp"

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <ctime>
#include <iostream>
#include <iterator>
#include <vector>

constexpr int MAXIMUM_DEPTH = 6;
constexpr int MAXIMUM_TIME_PER_MOVE = 7;
constexpr int Q_DEPTH_LIMIT = 4;

constexpr int MATE = 10000;

// Material + PST bonuses can never reach this, so anything at/above it is a
// forced mate. Used to stop iterative deepening once a mate is found.
constexpr int MATE_THRESHOLD = MATE - 1000;
constexpr int INF = 100000000;

uint64_t Nodes = 0;

static int ScoreMove(const Game *board, const Move &move,
                     const Move *bestMove) {
  // Indexed by PieceType (NONE, PAWN, KNIGHT, BISHOP, ROOK, QUEEN, KING).
  static constexpr int PIECE_VALUES[] = {
      0,     // NONEPIECE
      100,   // PAWN
      320,   // KNIGHT
      330,   // BISHOP
      500,   // ROOK
      900,   // QUEEN
      20000, // KING (never captured in legal play)
  };
  static_assert(std::size(PIECE_VALUES) == KING + 1);

  int score = 0;

  // Search the previous iteration's best move first (iterative deepening).
  if (bestMove != nullptr && move == *bestMove) {
    score += 1000000;
  }

  const Piece moving = board->pieces[PositionToIndex(move.From)];
  const Piece captured = board->pieces[PositionToIndex(move.To)];

  // MVV-LVA: value the capture by what we win, penalise by what we spend.
  if (captured.type != NONEPIECE) {
    score += 10000;
    score += PIECE_VALUES[captured.type] * 10;
    score -= PIECE_VALUES[moving.type];
  }

  if (move.promotion != NONEPIECE) {
    score += 8000;
  }

  return score;
}

static std::vector<Move> GetSortedLegalMoves(Game *g, bool generateQuietMoves,
                                             const Move *bestMove) {
  auto moves = GetLegalMoves(g, generateQuietMoves);
  std::ranges::sort(moves, [&](const Move &a, const Move &c) {
    return ScoreMove(g, a, bestMove) > ScoreMove(g, c, bestMove);
  });
  return moves;
}

// Quiescence search: only captures, so the eval isn't blind to hanging pieces.
// `standPat` lets the side to move decline every capture. Terminal positions
// are scored with the same `MATE - ply` convention as the main search so that
// mate distances stay consistent (a mate found inside quiescence must not look
// faster than a real mate-in-1). White perspective throughout.
static int quiescenceSearch(Game *b, int qdepth, int alpha, int beta, int ply) {
  Nodes++;
  // Runs GetNewGameState internally, so b->state is up to date afterwards.
  const int standPat = EvaluateBoard(b);

  switch (b->state) {
  case WHITE_WON:
  case BLACK_WON:
    return -(MATE - ply);
    break;
  case DRAW:
    return 0;
    break;
  case TURN:
    break;
  }

  if (qdepth >= Q_DEPTH_LIMIT) {
    return standPat;
  }

  if (standPat >= beta) {
    return beta;
  }
  alpha = std::max(alpha, standPat);

  auto moves = GetSortedLegalMoves(b, false, nullptr);

  for (Move move : moves) {
    UndoMove undo = MakeMove(move, b);

    int score = -quiescenceSearch(b, qdepth + 1, -beta, -alpha, ply + 1);

    UnMakeMove(undo, b);

    if (score >= beta) {
      return beta;
    }
    alpha = std::max(alpha, score);
  };
  return alpha;
};

static int search(int depth, Game *b, int alpha, int beta, int ply) {
  Nodes++;
  const uint64_t gameHash = GenerateZobristKey(b);

  auto repIt = b->ThreeFoldMap.find(gameHash);
  if (repIt != b->ThreeFoldMap.end() && repIt->second >= 2) {
    return 0; // repetition -> draw
  }

  TranspositionsEntry *entry = nullptr;
  if (auto it = b->Transpositions->find(gameHash);
      it != b->Transpositions->end()) {
    entry = &it->second;
    if (entry->depth >= depth) {
      if (entry->flag == EXACT) {
        return entry->Eval;
      }
      if (entry->flag == LOWERBOUND) {
        alpha = std::max(alpha, entry->Eval);
      }
      if (entry->flag == UPPERBOUND) {
        beta = std::min(beta, entry->Eval);
      }
      if (alpha >= beta) {
        return entry->Eval;
      }
    }
  }

  if (depth <= 0) {
    return quiescenceSearch(b, 0, alpha, beta, ply);
  }

  Move ttBestMove = entry != nullptr ? entry->bestMove : Move{};
  std::vector<Move> moves = GetSortedLegalMoves(b, true, &ttBestMove);

  if (moves.empty()) {
    const Position king = FindKing(b, b->turn);
    if (IsSquareAttacked(b, king, !b->turn)) {
      return -(MATE - ply); // mated
    }
    return 0; // stalemate
  }

  Move bestMove = moves[0];
  const int alphaOrig = alpha;
  int bestScore = -INF;

  for (Move move : moves) {
    UndoMove undo = MakeMove(move, b);

    int score = -search(depth - 1, b, -beta, -alpha, ply + 1);

    UnMakeMove(undo, b);

    if (score > bestScore) {
      bestScore = score;
      bestMove = move;
    }
    alpha = std::max(alpha, score);
    if (alpha >= beta) {
      break;
    }
  }

  const Flag flag = bestScore <= alphaOrig ? UPPERBOUND
                    : bestScore >= beta    ? LOWERBOUND
                                           : EXACT;
  (*b->Transpositions)[gameHash] = {
      .depth = depth, .Eval = bestScore, .flag = flag, .bestMove = bestMove};
  return bestScore;
}

struct SearchResult {
  Move bestMove;
  int score;
};

// Searches every root move to `depth` plies and returns the best one.
// `previousBest` is the best move from the previous iteration (used for move
// ordering, the core win of iterative deepening).
static SearchResult SearchDepth(Game *b, int depth, const Move *previousBest) {
  auto moves = GetSortedLegalMoves(b, true, previousBest);

  if (moves.empty()) {
    return {{}, EvaluateBoardForWhite(b)};
  }

  Move bestMove = moves[0];
  int bestEval = -INF;

  int alpha = -INF;
  int beta = INF;
  for (Move move : moves) {
    UndoMove undo = MakeMove(move, b);

    int eval = -search(depth - 1, b, -beta, -alpha, 1);

    UnMakeMove(undo, b);

    if (eval > bestEval) {
      bestEval = eval;
      bestMove = move;
    }
    alpha = std::max(alpha, eval);
  }

  return {bestMove, bestEval};
}

static void PrintInfo(const int depth, const int engineScore,
                      const uint64_t nps) {
  if (std::abs(engineScore) >= MATE_THRESHOLD) {
    int movesToMate = (MATE - std::abs(engineScore) + 1) / 2;
    if (movesToMate < 1) {
      movesToMate = 1;
    }
    std::cout << "info depth " << depth << " score mate "
              << (engineScore > 0 ? movesToMate : -movesToMate) << "\n"
              << std::flush;
  } else {
    std::cout << "info depth " << depth << " score cp " << engineScore
              << " nodes " << Nodes << " nps " << nps << "\n"
              << std::flush;
  }
}

Move GetBestMove(Game *b, const int maxDepth) {
  Nodes = 0;
  int actualDepth = maxDepth > 0 ? maxDepth : MAXIMUM_DEPTH;
  bool usingDefaultDepth = actualDepth == MAXIMUM_DEPTH;
  auto legalMoves = GetSortedLegalMoves(b, true, nullptr);
  if (legalMoves.empty()) {
    assert(false && "GetBestMove called with no legal moves");
    return {};
  }

  Move bestMove = legalMoves[0];

  time_t start = time(0);
  auto startMili = std::chrono::steady_clock::now();

  bool continueSearching = true;
  int depth = 1;

  while (continueSearching) {
    SearchResult result = SearchDepth(b, depth, &bestMove);
    bestMove = result.bestMove;

    auto now = std::chrono::steady_clock::now();

    double seconds = std::chrono::duration<double>(now - startMili).count();

    uint64_t nps =
        seconds > 0
            ? static_cast<uint64_t>(static_cast<double>(Nodes) / seconds)
            : Nodes;
    PrintInfo(depth, result.score, nps);

    // A mate was found; deeper searches can only find a faster one.
    if (std::abs(result.score) >= MATE_THRESHOLD) {
      break;
    }
    depth++;
    if (depth >= actualDepth) {
      continueSearching = false;
    }
    int since_start = static_cast<int>(difftime(time(0), start));
    if (since_start >= MAXIMUM_TIME_PER_MOVE && usingDefaultDepth) {
      continueSearching = false;
    }
  }

  return bestMove;
}