#include <algorithm>
#include <array>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <iterator>
#include <optional>
#include <ratio>
#include <vector>

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

constexpr int DEFAULT_TIME = 7;
constexpr int Q_DEPTH_LIMIT = 4;
constexpr int BOOK_DEPTH = 100;

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;

double timeToThingMS = -1;
std::chrono::time_point<std::chrono::steady_clock> searchStartTime;
bool searchStopped = false;

static int ScoreMove(const Game *board, const uint16_t &move,
                     const uint16_t *bestMove) {
  uint8_t from = getFromValueFromMove(move);
  uint8_t to = getToValueFromMove(move);
  PieceType promotion = getPromotionTypeFromMove(move);

  // Indexed by PieceType (NONE, PAWN, KNIGHT, BISHOP, ROOK, QUEEN, KING).
  static constexpr std::array<int, 7> 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[from];
  const Piece captured = board->pieces[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];
  }

  // promotion bonus
  if (IndexToPosition(to).rank == (board->turn ? 7 : 0)) {
    if (promotion == QUEEN) {
      score += 8000;
    }
    score += 3000;
  }

  if (IndexToPosition(to).rank == (board->turn ? 6 : 1) &&
      moving.type == PAWN) {
    score += 8000;
  }

  return score;
}

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

static int quiescenceSearch(Game *b, int qdepth, int alpha, int beta, int ply) {
  if (timeToThingMS == -1) {
    assert(false && "Expected set time: internal error");
    exit(1);
  }
  Nodes++;
  if ((Nodes & 2047) == 0) {
    double elapsedMiliseconds =
        std::chrono::duration<double, std::milli>(
            std::chrono::steady_clock::now() - searchStartTime)
            .count();
    if (elapsedMiliseconds >= timeToThingMS) {
      searchStopped = true;
      return 0;
    }
  }

  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, CAPTUARES_ONLY, nullptr);

  for (uint16_t move : moves) {
    Undo undo = MakeMove(move, b);

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

    UndoMove(undo, b);
    if (searchStopped) {
      return 0;
    }

    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) {
  if (timeToThingMS == -1) {
    assert(false && "Expected set time: internal error");
    exit(1);
  }
  Nodes++;
  if ((Nodes & 2047) == 0) {
    double elapsedMiliseconds =
        std::chrono::duration<double, std::milli>(
            std::chrono::steady_clock::now() - searchStartTime)
            .count();
    if (elapsedMiliseconds >= timeToThingMS) {
      searchStopped = true;
      return 0;
    }
  }
  const uint64_t gameHash = GenerateZobristKey(b);

  if (isRepetionDraw(gameHash, b)) {
    return 0;
  };

  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);
  }

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

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

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

  for (uint16_t move : moves) {
    Undo undo = MakeMove(move, b);

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

    UndoMove(undo, b);
    if (searchStopped) {
      break;
    }
    if (score > bestScore) {
      bestScore = score;
      bestMove = move;
    }
    alpha = std::max(alpha, score);
    if (alpha >= beta) {
      break;
    }
  }
  if (!searchStopped) {
    Flag flag = EXACT;
    if (bestScore >= beta) {
      flag = LOWERBOUND;
    } else if (bestScore <= alphaOrig) {
      flag = UPPERBOUND;
    }

    (*b->Transpositions)[gameHash] = {
        .depth = depth, .Eval = bestScore, .flag = flag, .bestMove = bestMove};
  }
  return bestScore;
}

struct SearchResult {
  uint16_t bestMove;
  int score;
};

/*
 * Search nodes until depth X.
 * Cutting everything bad using alpha-beta pruning.
 */
static SearchResult SearchDepth(Game *b, int depth,
                                const uint16_t *previousBest) {
  auto moves = GetSortedLegalMoves(b, ALL, previousBest);

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

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

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

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

    UndoMove(undo, b);

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

  return {.bestMove = bestMove, .score = bestEval};
}

struct Info {
  const int depth;
  const int score;
  const std::optional<int> nodes;
  const std::optional<int> nps;
};
static void PrintInfo(const Info &info) {

  if (std::abs(info.score) >= MATE_THRESHOLD) {
    int movesToMate = (MATE - std::abs(info.score) + 1) / 2;
    movesToMate = std::max(1, movesToMate);
    std::cout << "info depth " << info.depth << " score mate "
              << (info.score > 0 ? movesToMate : -movesToMate) << "\n"
              << std::flush;
  } else {
    std::cout << "info depth " << info.depth << " score cp " << info.score;
    if (info.nodes.has_value()) {
      std::cout << " nodes " << info.nodes.value();
    }
    if (info.nps.has_value()) {
      std::cout << " nps " << info.nps.value();
    }
    std::cout << "\n" << std::flush;
  }
}

uint16_t GetBestMove(Game *b, int maxDepth, move_options options) {
  searchStopped = false;
  Nodes = 0;

  uint64_t hash = GenerateZobristKey(b);

  if (b->MoveClock <= BOOK_DEPTH) {
    std::optional<uint16_t> bookMove = ProbeBook(hash);
    if (bookMove.has_value()) {
      return bookMove.value();
    }
  }

  // use depth INF if caller hasnt provided depth
  int actualDepth = maxDepth > 0 ? maxDepth : INF;

  bool hasSetSpecialTimeLimit = false;

  auto legalMoves = GetSortedLegalMoves(b, ALL, nullptr);
  if (legalMoves.empty()) {
    assert(false && "GetBestMove called with no legal moves");
    return {};
  }

  uint16_t bestMove = legalMoves[0];

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

  timeToThingMS = DEFAULT_TIME * 1000; // to big number to ever achiave

  if (options.wtime != -1 && options.btime != -1) {
    hasSetSpecialTimeLimit = true;
    int increment = b->turn ? options.wncr : options.bncr;
    double time_remaning = b->turn ? options.wtime : options.btime;

    timeToThingMS = ((time_remaning / 15) + (increment * 0.8));
  }

  bool continueSearching = true;
  int depth = 1;

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

    // If the clock expired inside this iteration, the returned score/move may
    // come from a partially searched tree. Keep the last fully completed
    // iteration instead of reporting INF as a fake mate or playing a random
    // first move.
    if (searchStopped) {
      break;
    }

    bestMove = result.bestMove;

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

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

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

    // 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;
    }
    double elapsedSeconds =
        std::chrono::duration<double>(std::chrono::steady_clock::now() -
                                      searchStartTime)
            .count();
    if (elapsedSeconds >= DEFAULT_TIME && !hasSetSpecialTimeLimit) {
      continueSearching = false;
    }
  }

  return bestMove;
}