#include "book.hpp"

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <optional>
#include <random>
#include <vector>

#include "board/board.hpp"

struct BookEntry {
  uint64_t key;
  uint16_t move;
  uint16_t weight;
  uint32_t learn;
};

static std::vector<BookEntry> book;

static std::mt19937 rng{std::random_device{}()};

void InitBook(const std::string &path) {
  std::ifstream file(path, std::ios::binary);
  if (!file.is_open()) {
    std::cerr << "info string Failed to open: " << path << "\n";
    return;
  }

  file.seekg(0, std::ios::end);
  size_t fileSize = static_cast<size_t>(file.tellg());
  file.seekg(0, std::ios::beg);

  if (fileSize == 0 || fileSize % sizeof(BookEntry) != 0) {
    std::cerr << "info string Invalid book file size\n";
    return;
  }

  size_t numEntries = fileSize / sizeof(BookEntry);
  book.resize(numEntries);
  file.read(reinterpret_cast<char *>(book.data()),
            static_cast<std::streamsize>(fileSize));

  if (!file) {
    std::cerr << "info string Failed to read book file\n";
    book.clear();
    return;
  }

  // Byteswap if needed (Polyglot books are big-endian)
  for (auto &entry : book) {
    uint64_t k = entry.key;
    entry.key = ((k & 0x00000000000000FFULL) << 56) |
                ((k & 0x000000000000FF00ULL) << 40) |
                ((k & 0x0000000000FF0000ULL) << 24) |
                ((k & 0x00000000FF000000ULL) << 8) |
                ((k & 0x000000FF00000000ULL) >> 8) |
                ((k & 0x0000FF0000000000ULL) >> 24) |
                ((k & 0x00FF000000000000ULL) >> 40) |
                ((k & 0xFF00000000000000ULL) >> 56);

    uint16_t m = entry.move;
    entry.move = static_cast<uint16_t>((m >> 8) | (m << 8));

    uint16_t w = entry.weight;
    entry.weight = static_cast<uint16_t>((w >> 8) | (w << 8));

    uint32_t l = entry.learn;
    entry.learn = ((l & 0x000000FFU) << 24) | ((l & 0x0000FF00U) << 8) |
                  ((l & 0x00FF0000U) >> 8) | ((l & 0xFF000000U) >> 24);
  }

  std::cerr << "info string " << path << ": " << book.size() << " entries\n";
}

static uint16_t PolyglotToMono(uint16_t polyglotMove, const Game &g) {
  uint8_t to = polyglotMove & 63;
  uint8_t from = (polyglotMove >> 6) & 63;
  uint8_t promo = (polyglotMove >> 12) & 7;

  if (g.pieces[from].type == KING) {
    if (from == 4 && to == 7) {
      to = 6; // white kingside: e1h1 -> e1g1
    } else if (from == 4 && to == 0) {
      to = 2; // white queenside: e1a1 -> e1c1
    } else if (from == 60 && to == 63) {
      to = 62; // black kingside: e8h8 -> e8g8
    } else if (from == 60 && to == 56) {
      to = 58; // black queenside: e8a8 -> e8c8
    }
  }

  PieceType promotion = NONEPIECE;
  switch (promo) {
  case 0:
    promotion = KNIGHT;
    break;
  case 1:
    promotion = BISHOP;
    break;
  case 2:
    promotion = ROOK;
    break;
  case 3:
    promotion = QUEEN;
    break;
  default:
    break;
  }

  return CreateMove(from, to, promotion);
}

std::optional<uint16_t> ProbeBook(uint64_t key, const Game &g) {
  if (book.empty()) {
    return {};
  }

  // Binary search for first matching entry
  auto it = std::lower_bound(
      book.begin(), book.end(), key,
      [](const BookEntry &a, uint64_t k) { return a.key < k; });

  // Collect all matching entries
  std::vector<const BookEntry *> matches;
  while (it != book.end() && it->key == key) {
    matches.push_back(&(*it));
    ++it;
  }

  if (matches.empty()) {
    return {};
  }

  // Weighted random selection
  int totalWeight = 0;
  for (const auto *e : matches) {
    totalWeight += e->weight;
  }

  std::uniform_int_distribution<int> dist(0, totalWeight - 1);
  int r = dist(rng);

  for (const auto *e : matches) {
    r -= e->weight;
    if (r < 0) {
      return PolyglotToMono(e->move, g);
    }
  }

  return PolyglotToMono(matches.back()->move, g);
}