#include "bot.hpp"
#include "board/board.hpp"
#include "moves.hpp"
#include "zobrist/zobrist.hpp"
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <vector>

constexpr int DEFAULT_DEPTH = 4;
constexpr int Q_DEPTH_LIMIT = 3;

constexpr int PAWN_VALUE = 100;
constexpr int KNIGHT_VALUE = 320;
constexpr int BISHOP_VALUE = 400;
constexpr int ROOK_VALUE = 500;
constexpr int QUEEN_VALUE = 900;
constexpr int MATE = 10000;

static int PAWN_TABLE[64] = {
    0,  0,  0,   0,   0,   0,   0,  0,  // last rank promotes to a quuen
    50, 50, 50,  50,  50,  50,  50, 50, // comments to stop formating to ruin it
    10, 10, 20,  35,  35,  20,  10, 10, //
    5,  5,  10,  30,  30,  10,  5,  5,  //
    0,  0,  0,   25,  25,  0,   0,  0,  //
    5,  -5, -10, 0,   0,   -10, -5, 5,  //
    5,  10, 10,  -20, -20, 10,  10, 5,  //
    0,  0,  0,   0,   0,   0,   0,  0   //
};

static int KNIGHT_TABLE[64] = {
    -50, -40, -30, -30, -30, -30, -40, -50, //
    -40, -20, 0,   0,   0,   0,   -20, -40, //
    -30, 0,   10,  15,  15,  10,  0,   -30, //
    -30, 5,   15,  20,  20,  15,  5,   -30, //
    -30, 0,   15,  20,  20,  15,  0,   -30, //
    -30, 5,   10,  15,  15,  10,  5,   -30, //
    -40, -20, 0,   5,   5,   0,   -20, -40, //
    -50, -40, -30, -30, -30, -30, -40, -50, //
};

static int BISHOP_TABLE[64] = {
    -20, -10, -10, -10, -10, -10, -10, -20, //
    -10, 5,   0,   0,   0,   0,   5,   -10, //
    -10, 10,  10,  10,  10,  10,  10,  -10, //
    -10, 0,   10,  15,  15,  10,  0,   -10, //
    -10, 5,   5,   10,  10,  5,   5,   -10, //
    -10, 0,   5,   10,  10,  5,   0,   -10, //
    -10, 0,   0,   0,   0,   0,   0,   -10, //
    -20, -10, -10, -10, -10, -10, -10, -20, //
};

static int ROOK_TABLE[64] = {
    0,  0,  5,  10, 10, 5,  0,  0,  //
    5,  10, 10, 10, 10, 10, 10, 5,  //
    -5, 0,  0,  0,  0,  0,  0,  -5, //
    -5, 0,  0,  5,  5,  0,  0,  -5, //
    -5, 0,  0,  5,  5,  0,  0,  -5, //
    -5, 0,  0,  0,  0,  0,  0,  -5, //
    5,  10, 10, 10, 10, 10, 10, 5,  //
    0,  0,  5,  10, 10, 5,  0,  0,  //
};

static int QUEEN_TABLE[64] = {
    -20, -10, -10, -5, -5, -10, -10, -20, //
    -10, 0,   0,   0,  0,  0,   0,   -10, //
    -10, 0,   5,   5,  5,  5,   0,   -10, //
    -5,  0,   5,   5,  5,  5,   0,   -5,  //
    0,   0,   5,   5,  5,  5,   0,   -5,  //
    -10, 5,   5,   5,  5,  5,   0,   -10, //
    -10, 0,   5,   0,  0,  0,   0,   -10, //
    -20, -10, -10, -5, -5, -10, -10, -20, //
};

static int KING_TABLE_EARLY[64] = {
    -30, -40, -40, -50, -50, -40, -40, -30, //
    -30, -40, -40, -50, -50, -40, -40, -30, //
    -30, -40, -40, -50, -50, -40, -40, -30, //
    -30, -40, -40, -50, -50, -40, -40, -30, //
    -20, -30, -30, -40, -40, -30, -30, -20, //
    -10, -20, -20, -20, -20, -20, -20, -10, //
    20,  20,  0,   0,   0,   0,   20,  20,  //
    20,  30,  10,  0,   0,   10,  30,  20,  //
};

int ScoreMove(const Game *board, const Move &move) {
  int score = 0;

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

  // Captures (MVV-LVA)
  if (captured.type != NONEPIECE) {
    static const int pieceValue[] = {
        0,     // NONE
        20000, // KING (should never happen)
        900,   // QUEEN
        500,   // ROOK
        330,   // BISHOP
        320,   // KNIGHT
        100    // PAWN
    };

    score += 10000;
    score += pieceValue[captured.type] * 10;
    score -= pieceValue[moving.type];
  }

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

  return score;
}

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

Move EngineGetBestMove(Game *b, const int depth) {
  const auto moves = GetSortedLegalMoves(b);
  if (moves.empty()) {
    std::cout << "Expected a position with legal moves";
    assert(false && "Error zero legal moves for bot");
    exit(1);
  }

  Move bestMove = moves[0];
  int BestEval = b->turn ? std::numeric_limits<int>::lowest()
                         : std::numeric_limits<int>::max();
  for (const Move move : moves) {
    const UndoMove undo = MakeMove(move, b);
    const float alpha = -INFINITY;
    const float beta = INFINITY;

    const int eval =
        minimax(depth <= 0 ? DEFAULT_DEPTH : depth, b, alpha, beta);

    UnMakeMove(undo, b);
    if (b->turn) {
      if (eval > BestEval) {
        BestEval = eval;
        bestMove = move;
      }
    } else {
      if (eval < BestEval) {
        BestEval = eval;
        bestMove = move;
      }
    }
  }
  return bestMove;
}
int quiescenceSearch(Game *b, int Qdepth) {
  int standPat = EvaluateBoardForWhite(b);
  if (Qdepth >= Q_DEPTH_LIMIT) {
    return standPat;
  }
  std::vector<Move> attackMoves = GetSortedLegalMoves(b, false);
  if (attackMoves.empty()) {
    return standPat;
  }

  int bestScore = b->turn ? std::numeric_limits<int>::lowest()
                          : std::numeric_limits<int>::max();

  if (b->turn) {
    for (Move move : attackMoves) {
      UndoMove undo = MakeMove(move, b);

      int eval = quiescenceSearch(b, Qdepth + 1);

      UnMakeMove(undo, b);
      bestScore = std::max(bestScore, eval);
    }
  } else {
    for (Move move : attackMoves) {
      UndoMove undo = MakeMove(move, b);

      int eval = quiescenceSearch(b, Qdepth + 1);

      UnMakeMove(undo, b);
      bestScore = std::min(bestScore, eval);
    }
  }

  return bestScore;
}

int minimax(int depth, Game *b, float alpha, float beta) {
  uint64_t gameHash = GenerateZobristKey(b);

  if (b->ThreeFoldMap[gameHash] >= 2) {
    return 0;
  }
  if (b->Transpositions->contains(gameHash)) {
    TranspositionsEntry data = b->Transpositions->at(gameHash);
    if (data.depth >= depth) {
      return data.Eval;
    }
  }

  if (depth <= 0) {
    return quiescenceSearch(b, 0);
  }
  auto moves = GetSortedLegalMoves(b);
  if (moves.empty()) {
    return EvaluateBoardForWhite(b);
  }

  bool shouldStore = true;

  int bestScore = b->turn ? std::numeric_limits<int>::lowest()
                          : std::numeric_limits<int>::max();

  if (b->turn) {
    for (Move move : moves) {
      UndoMove undo = MakeMove(move, b);

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

      UnMakeMove(undo, b);
      bestScore = std::max(bestScore, eval);

      alpha = std::max(alpha, static_cast<float>(bestScore));

      if (alpha >= beta) {
        shouldStore = false;
        break; // *snips*
      }
    }
  } else {
    for (Move move : moves) {
      UndoMove undo = MakeMove(move, b);

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

      UnMakeMove(undo, b);
      bestScore = std::min(bestScore, eval);

      beta = std::min(beta, static_cast<float>(bestScore));
      if (alpha >= beta) {
        shouldStore = false;
        break; // *snips*
      }
    }
  }

  if (shouldStore) {
    b->Transpositions->operator[](gameHash) = {.depth = depth,
                                               .Eval = bestScore};
  }
  return bestScore;
}

static int PSTIndex(const int square, const bool white) {
  return white ? square : (56 ^ square);
}

static int ForceKingToEdgeBonus(const Position enemyKing,
                                const Position myKing) {
  int bonus = 0;

  // Push enemy king toward edge
  const int distToCenter =
      std::abs(enemyKing.file - 3) + std::abs(enemyKing.rank - 3);

  bonus += distToCenter * 10;

  // Bring own king closer
  const int kingDistance = std::abs(myKing.file - enemyKing.file) +
                           std::abs(myKing.rank - enemyKing.rank);

  bonus += (14 - kingDistance) * 5;

  return bonus;
}
static bool IsEndgame(Game *g) {
  int queens = 0;
  int rooks = 0;

  for (const Piece &piece : g->pieces) {
    switch (piece.type) {
    case QUEEN:
      queens++;
      break;
    case ROOK:
      rooks++;
      break;
    default:
      break;
    }
  }

  return queens == 0 || (queens == 2 && rooks <= 1);
}

int EvaluateBoardForWhite(Game *g) {
  int score = 0;

  switch (GetNewGameState(g)) {
  case WHITE_WON:
    return MATE;
    break;
  case BLACK_WON:
    return -MATE;
    break;
  case TURN:
    break;
  case DRAW:
    return 0;
  }
  for (int i = 0; i < 64; i++) {
    const Piece piece = g->pieces[i];
    if (piece.type == NONEPIECE)
      continue;

    int value = 0;

    switch (piece.type) {
    case PAWN:
      value = PAWN_VALUE;
      value += PAWN_TABLE[PSTIndex(i, piece.color)];
      break;
    case KNIGHT:
      value = KNIGHT_VALUE;
      value += KNIGHT_TABLE[PSTIndex(i, piece.color)];
      break;
    case BISHOP:
      value = BISHOP_VALUE;
      value += BISHOP_TABLE[PSTIndex(i, piece.color)];
      break;
    case ROOK:
      value = ROOK_VALUE;
      value += ROOK_TABLE[PSTIndex(i, piece.color)];
      break;
    case QUEEN:
      value = QUEEN_VALUE;
      value += QUEEN_TABLE[PSTIndex(i, piece.color)];
      break;
    case KING:
      value = KING_TABLE_EARLY[PSTIndex(i, piece.color)];
      break;
    default:
      break;
    }

    if (piece.color) {
      score += value;
    } else {
      score -= value;
    }
  }

  Position whiteKing = FindKing(g, true);
  Position blackKing = FindKing(g, false);

  if (IsEndgame(g)) {
    score +=
        ForceKingToEdgeBonus(blackKing, whiteKing); // White attacking black
    score -=
        ForceKingToEdgeBonus(whiteKing, blackKing); // Black attacking white
  }
  return score;
}