// ============================================================
// Chess960 — Open Chess
//
// Fischer Random. The board picks one of the 960 legal starting arrays and
// walks you through setup piece type by piece type so you never have to
// decode the position yourself. Both sides get the same back rank. Castling
// still exists: the king ends on g/c and the rook on f/d, same as usual.
//
// This mode needs the Connect page. The sensors only see occupied or empty,
// so the page is what names each square — king, queen, bishop — while the
// board lights where that piece belongs.
//
//   Button A  shuffle a new position (setup) or restart (during a game)
//   Button B  pause: the board goes dark until A is pressed
//
// To castle, lift the king and then the rook you want to castle with. The
// board lights the two destination squares and waits for you to put them
// there. That is how it tells a castle from a king that simply walked.
//
// Everything hardware-related — sensors, LEDs, battery, the USB link to the
// Connect page — is handled by OpenChessCore. This file is only the game.
// ============================================================

#include <OpenChessCore.h>
#include <oc_chess.h>
#include <string.h>

using namespace oc;
using chess::Move;
using chess::Position;

// ── Things you might want to change ─────────────────────────────────────────

const Color HINT_COLOR    = GREEN;
const Color CAPTURE_COLOR = ORANGE;
const Color ILLEGAL_COLOR = AMBER;
const Color CHECK_COLOR   = RED;
const Color PROMO_COLOR   = GOLD;
const Color CASTLE_COLOR  = CYAN;
const Color WHITE_WINS    = WHITE;
const Color BLACK_WINS    = PURPLE;
const Color TURN_COLOR    = GREEN.scaled(70);

// Setup walkthrough — one colour per piece so Connect and the board agree.
const Color SETUP_PAWN   = WHITE;
const Color SETUP_KING   = GOLD;
const Color SETUP_QUEEN  = PURPLE;
const Color SETUP_BISHOP = BLUE;
const Color SETUP_KNIGHT = GREEN;
const Color SETUP_ROOK   = CYAN;
const Color SETUP_PLACED = GREEN.scaled(45);

// After castling, king on g/c, rook on f/d — the usual squares.
const int KING_KS_FILE = 6, ROOK_KS_FILE = 5;
const int KING_QS_FILE = 2, ROOK_QS_FILE = 3;

const char STEP_KIND[] = { 'P', 'K', 'Q', 'B', 'N', 'R' };
const char *STEP_MSG[] = {
  "Place all 16 pawns on ranks 2 and 7",
  "Place both kings on the gold squares",
  "Place both queens on the purple squares",
  "Place both bishops on the blue squares",
  "Place both knights on the green squares",
  "Place both rooks on the cyan squares",
};

// ── Game state ──────────────────────────────────────────────────────────────

Position pos;
char     layout[8][8];
int      positionId   = 518;
int      kingFile     = 4;
int      rookQSFile   = 0;
int      rookKSFile   = 7;
bool     whiteKingMoved = false, blackKingMoved = false;
bool     whiteQSMoved   = false, whiteKSMoved   = false;
bool     blackQSMoved   = false, blackKSMoved   = false;
bool     awaitingSetup  = true;
bool     paused         = false;
bool     seeded         = false;

struct Castle {
  bool kingside;
  int  rank, kFrom, rFrom, kTo, rTo;
};

void sendPosition();
void showTurn();
void restart();
void pauseGame();
void dealPosition();
void waitForSetup();
bool occupancyMatches(const char grid[8][8]);
bool castleLegal(bool kingside, Castle &out);
bool completeCastle(const Castle &castle);
void applyCastle(const Castle &castle);
void noteRights(const Move &m, char moved);
void gameOverScreen(Color color);
bool waitForMove(int fromRow, int fromCol, Move &out, bool &applied);
void commitMove(const Move &m);
void afterMove();
void waitForPromotion(int row, int col);

void setup() {
  begin("chess960", "Chess960", "2.0.0");
  app.onIdentify = sendPosition;
}

void loop() {
  service();

  if (paused) {
    if (buttons.pressed(BUTTON_A)) { paused = false; awaitingSetup = true; }
    return;
  }

  if (awaitingSetup) {
    waitForSetup();
    return;
  }

  if (buttons.pressed(BUTTON_A)) { restart(); return; }
  if (buttons.pressed(BUTTON_B)) { pauseGame(); return; }

  for (int row = 0; row < 8; row++) {
    for (int col = 0; col < 8; col++) {
      if (!board.lifted(row, col)) continue;
      const char piece = pos.at(row, col);
      if (piece == ' ' || chess::colorOf(piece) != pos.turn) continue;

      Move move;
      bool applied = false;
      if (!waitForMove(row, col, move, applied)) return;
      if (!applied) commitMove(move);
      afterMove();
      return;
    }
  }
  board.sync();
}

// ── Chess960 numbering (Scharnagl). 518 is the ordinary starting array. ─────

void generateBackRank(int n, char out[8]) {
  memset(out, 0, 8);
  n = ((n % 960) + 960) % 960;

  out[1 + (n % 4) * 2] = 'B';           // light-squared bishop: b d f h
  n /= 4;
  out[(n % 4) * 2] = 'B';               // dark-squared bishop: a c e g
  n /= 4;

  int empty[8], count = 0;
  for (int i = 0; i < 8; i++) if (!out[i]) empty[count++] = i;
  out[empty[n % 6]] = 'Q';
  n /= 6;

  count = 0;
  for (int i = 0; i < 8; i++) if (!out[i]) empty[count++] = i;
  static const int N1[10] = { 0, 0, 0, 0, 1, 1, 1, 2, 2, 3 };
  static const int N2[10] = { 1, 2, 3, 4, 2, 3, 4, 3, 4, 4 };
  out[empty[N1[n]]] = 'N';
  out[empty[N2[n]]] = 'N';

  count = 0;
  for (int i = 0; i < 8; i++) if (!out[i]) empty[count++] = i;
  out[empty[0]] = 'R';
  out[empty[1]] = 'K';
  out[empty[2]] = 'R';
}

void dealPosition() {
  if (!seeded) {
    randomSeed(micros());
    seeded = true;
  }
  positionId = (int)random(960);

  char back[8];
  generateBackRank(positionId, back);

  memset(layout, ' ', sizeof(layout));
  kingFile = rookQSFile = rookKSFile = -1;
  for (int col = 0; col < 8; col++) {
    layout[0][col] = back[col];
    layout[1][col] = 'P';
    layout[6][col] = 'p';
    layout[7][col] = (char)(back[col] + 32);
    if (back[col] == 'K') kingFile = col;
    if (back[col] == 'R') {
      if (rookQSFile < 0) rookQSFile = col;
      else                rookKSFile = col;
    }
  }

  memcpy(pos.squares, layout, sizeof(layout));
  pos.turn = 'w';
  whiteKingMoved = blackKingMoved = false;
  whiteQSMoved = whiteKSMoved = blackQSMoved = blackKSMoved = false;

  char shown[9];
  memcpy(shown, back, 8);
  shown[8] = 0;
  app.log("Chess960 #%d: %s", positionId, shown);
  sendPosition();
}

Color colorForKind(char kind) {
  switch (kind) {
    case 'P': return SETUP_PAWN;
    case 'K': return SETUP_KING;
    case 'Q': return SETUP_QUEEN;
    case 'B': return SETUP_BISHOP;
    case 'N': return SETUP_KNIGHT;
    case 'R': return SETUP_ROOK;
  }
  return WHITE;
}

int currentSetupStep() {
  for (int step = 0; step < 6; step++) {
    const char kind = STEP_KIND[step];
    for (int row = 0; row < 8; row++)
      for (int col = 0; col < 8; col++)
        if (chess::kindOf(layout[row][col]) == kind && !board.occupied(row, col))
          return step;
  }
  return 6;
}

bool occupancyMatches(const char grid[8][8]) {
  for (int row = 0; row < 8; row++)
    for (int col = 0; col < 8; col++)
      if ((grid[row][col] != ' ') != board.occupied(row, col)) return false;
  return true;
}

void waitForSetup() {
  dealPosition();

  char msg[96];
  snprintf(msg, sizeof(msg), "Chess960 #%d — open Connect, then place the pawns", positionId);
  app.status("setup", msg);

  int lastStep = -1;
  while (!occupancyMatches(layout)) {
    service();
    if (buttons.pressed(BUTTON_A)) { restart(); return; }
    if (buttons.pressed(BUTTON_B)) { pauseGame(); return; }

    const int step = currentSetupStep();
    if (step != lastStep && step < 6) {
      lastStep = step;
      snprintf(msg, sizeof(msg), "Chess960 #%d — %s", positionId, STEP_MSG[step]);
      app.status("setup", msg);
    }

    const bool pulse = (millis() / 280) % 2 == 0;
    lights.clear();
    for (int row = 0; row < 8; row++) {
      for (int col = 0; col < 8; col++) {
        const char piece = layout[row][col];
        if (piece == ' ') {
          if (board.occupied(row, col)) lights.set(row, col, ILLEGAL_COLOR);
          continue;
        }
        if (board.occupied(row, col)) {
          lights.set(row, col, SETUP_PLACED);
          continue;
        }
        const bool current = step < 6 && chess::kindOf(piece) == STEP_KIND[step];
        if (current && pulse) lights.set(row, col, colorForKind(STEP_KIND[step]));
        else if (current)     lights.set(row, col, colorForKind(STEP_KIND[step]).scaled(70));
      }
    }
    lights.show();
    wait(80);
  }

  memcpy(pos.squares, layout, sizeof(layout));
  pos.turn = 'w';
  anim::ripple(3, 3, GOLD, 45);
  board.sync();
  awaitingSetup = false;

  sound.gameStart();
  snprintf(msg, sizeof(msg), "Chess960 #%d — White to move", positionId);
  app.status("playing", msg);
  sendPosition();
  showTurn();
}

void showTurn() {
  lights.clear();
  for (int row = 0; row < 8; row++)
    for (int col = 0; col < 8; col++) {
      const char p = pos.at(row, col);
      if (p != ' ' && chess::colorOf(p) == pos.turn) lights.set(row, col, TURN_COLOR);
    }
  lights.show();
}

void restart() {
  anim::flashAll(GREEN, 3);
  awaitingSetup = true;
  app.log("New game");
}

void pauseGame() {
  for (int col = 7; col >= 0; col--) {
    lights.set(0, col, RED);
    lights.show();
    wait(60);
  }
  lights.clear();
  lights.show();
  paused = true;
  app.status("idle", "Paused - press button A to start a new game");
}

void sendPosition() {
  char fen[96];
  pos.fen(fen);
  app.state(fen, pos.turn);
}

// ── Castling ────────────────────────────────────────────────────────────────

bool sideRights(char color, bool kingside) {
  if (color == 'w') {
    if (whiteKingMoved) return false;
    return kingside ? !whiteKSMoved : !whiteQSMoved;
  }
  if (blackKingMoved) return false;
  return kingside ? !blackKSMoved : !blackQSMoved;
}

bool vacantExcept(int rank, int from, int to, int keepA, int keepB) {
  const int lo = min(from, to), hi = max(from, to);
  for (int file = lo; file <= hi; file++) {
    if (file == keepA || file == keepB) continue;
    if (pos.at(rank, file) != ' ') return false;
  }
  return true;
}

bool kingPathSafe(int rank, int from, int to, char enemy) {
  const int step = from < to ? 1 : from > to ? -1 : 0;
  int file = from;
  for (;;) {
    if (pos.attacked(rank, file, enemy)) return false;
    if (file == to) break;
    file += step;
  }
  return true;
}

bool castleLegal(bool kingside, Castle &out) {
  const char color = pos.turn;
  if (!sideRights(color, kingside)) return false;
  if (pos.inCheck(color)) return false;

  const int rank = color == 'w' ? 0 : 7;
  const int rFrom = kingside ? rookKSFile : rookQSFile;
  const int kTo   = kingside ? KING_KS_FILE : KING_QS_FILE;
  const int rTo   = kingside ? ROOK_KS_FILE : ROOK_QS_FILE;
  const char rook = color == 'w' ? 'R' : 'r';
  if (pos.at(rank, kingFile) != (color == 'w' ? 'K' : 'k')) return false;
  if (pos.at(rank, rFrom) != rook) return false;

  if (!vacantExcept(rank, kingFile, kTo, kingFile, rFrom)) return false;
  if (!vacantExcept(rank, rFrom, rTo, kingFile, rFrom)) return false;
  if (!kingPathSafe(rank, kingFile, kTo, chess::other(color))) return false;

  out.kingside = kingside;
  out.rank = rank;
  out.kFrom = kingFile;
  out.rFrom = rFrom;
  out.kTo = kTo;
  out.rTo = rTo;
  return true;
}

void applyCastle(const Castle &c) {
  const char king = pos.turn == 'w' ? 'K' : 'k';
  const char rook = pos.turn == 'w' ? 'R' : 'r';
  pos.squares[c.rank][c.kFrom] = ' ';
  pos.squares[c.rank][c.rFrom] = ' ';
  pos.squares[c.rank][c.kTo] = king;
  pos.squares[c.rank][c.rTo] = rook;
  pos.turn = chess::other(pos.turn);
  if (c.rank == 0) { whiteKingMoved = whiteQSMoved = whiteKSMoved = true; }
  else             { blackKingMoved = blackQSMoved = blackKSMoved = true; }
}

bool completeCastle(const Castle &c) {
  char expected[8][8];
  memcpy(expected, pos.squares, sizeof(expected));
  const char king = pos.turn == 'w' ? 'K' : 'k';
  const char rook = pos.turn == 'w' ? 'R' : 'r';
  expected[c.rank][c.kFrom] = ' ';
  expected[c.rank][c.rFrom] = ' ';
  expected[c.rank][c.kTo] = king;
  expected[c.rank][c.rTo] = rook;

  char msg[80];
  snprintf(msg, sizeof(msg), "Castle %s — king to %c%c, rook to %c%c",
           c.kingside ? "kingside" : "queenside",
           (char)('a' + c.kTo), (char)('1' + c.rank),
           (char)('a' + c.rTo), (char)('1' + c.rank));
  app.status("playing", msg);

  while (!occupancyMatches(expected)) {
    service();
    if (buttons.pressed(BUTTON_A)) { restart(); return false; }
    if (buttons.pressed(BUTTON_B)) { pauseGame(); return false; }

    const bool pulse = (millis() / 240) % 2 == 0;
    lights.clear();
    lights.set(c.rank, c.kTo, pulse ? CASTLE_COLOR : GOLD);
    lights.set(c.rank, c.rTo, pulse ? GOLD : CASTLE_COLOR);
    for (int row = 0; row < 8; row++)
      for (int col = 0; col < 8; col++)
        if (expected[row][col] == ' ' && board.occupied(row, col) &&
            !(row == c.rank && (col == c.kTo || col == c.rTo)))
          lights.set(row, col, ILLEGAL_COLOR);
    lights.show();
    wait(40);
  }

  applyCastle(c);
  sound.castle();

  char uci[6];
  uci[0] = (char)('a' + c.kFrom);
  uci[1] = (char)('1' + c.rank);
  const int toFile = c.kFrom == c.kTo ? c.rFrom : c.kTo;
  uci[2] = (char)('a' + toFile);
  uci[3] = (char)('1' + c.rank);
  uci[4] = 0;
  app.move(uci);
  sendPosition();
  app.log("%s: O-O%s", chess::other(pos.turn) == 'w' ? "White" : "Black",
          c.kingside ? "" : "-O");

  lights.clear();
  lights.show();
  board.sync();
  return true;
}

// ── Making a move ───────────────────────────────────────────────────────────

void showHints(int fromRow, int fromCol, Move *options, int count, const Castle *ks, const Castle *qs) {
  lights.clear();
  lights.set(fromRow, fromCol, WHITE);
  for (int i = 0; i < count; i++) {
    lights.set(options[i].tr, options[i].tc,
               options[i].promotion ? PROMO_COLOR : options[i].isCapture() ? CAPTURE_COLOR : HINT_COLOR);
  }
  if (ks) lights.set(ks->rank, ks->rFrom, CASTLE_COLOR);
  if (qs) lights.set(qs->rank, qs->rFrom, CASTLE_COLOR);
  lights.show();
}

bool waitForMove(int fromRow, int fromCol, Move &out, bool &applied) {
  static Move options[chess::MAX_MOVES];
  const int count = pos.legalMovesFrom(fromRow, fromCol, options);
  applied = false;

  Castle ks, qs;
  const bool isKing = chess::kindOf(pos.at(fromRow, fromCol)) == 'K';
  const bool haveKS = isKing && castleLegal(true,  ks);
  const bool haveQS = isKing && castleLegal(false, qs);
  showHints(fromRow, fromCol, options, count, haveKS ? &ks : nullptr, haveQS ? &qs : nullptr);

  for (;;) {
    service();
    if (buttons.pressed(BUTTON_A)) { restart(); return false; }
    if (buttons.pressed(BUTTON_B)) { pauseGame(); return false; }

    if (board.occupied(fromRow, fromCol)) {
      showTurn();
      return false;
    }

    if (haveKS && board.lifted(ks.rank, ks.rFrom)) {
      if (!completeCastle(ks)) return false;
      applied = true;
      return true;
    }
    if (haveQS && board.lifted(qs.rank, qs.rFrom)) {
      if (!completeCastle(qs)) return false;
      applied = true;
      return true;
    }

    for (int i = 0; i < count; i++) {
      const int tr = options[i].tr, tc = options[i].tc;
      if (options[i].isCapture()) {
        if (board.lifted(tr, tc)) {
          lights.clear();
          lights.set(tr, tc, CAPTURE_COLOR);
          lights.show();
          while (!board.occupied(tr, tc)) {
            service();
            if (buttons.pressed(BUTTON_A)) { restart(); return false; }
          }
          out = options[i];
          return true;
        }
      } else if (board.placed(tr, tc)) {
        out = options[i];
        return true;
      }
    }

    // King dropped on the castle destination (more than one square away).
    if (haveKS && board.placed(ks.rank, ks.kTo) && abs(ks.kTo - ks.kFrom) > 1) {
      if (!completeCastle(ks)) return false;
      applied = true;
      return true;
    }
    if (haveQS && board.placed(qs.rank, qs.kTo) && abs(qs.kTo - qs.kFrom) > 1) {
      if (!completeCastle(qs)) return false;
      applied = true;
      return true;
    }

    for (int row = 0; row < 8; row++) {
      for (int col = 0; col < 8; col++) {
        if (!board.placed(row, col) || pos.at(row, col) != ' ') continue;
        if (haveKS && row == ks.rank && col == ks.kTo) continue;
        if (haveQS && row == qs.rank && col == qs.kTo) continue;
        sound.illegal();
        while (board.occupied(row, col)) {
          anim::flashSquare(row, col, ILLEGAL_COLOR, 1, 120);
          if (buttons.pressed(BUTTON_A)) { restart(); return false; }
        }
        showHints(fromRow, fromCol, options, count, haveKS ? &ks : nullptr, haveQS ? &qs : nullptr);
      }
    }
    wait(20);
  }
}

void noteRights(const Move &m, char moved) {
  if (moved == 'K') whiteKingMoved = true;
  if (moved == 'k') blackKingMoved = true;
  if (moved == 'R' && m.fr == 0) {
    if (m.fc == rookQSFile) whiteQSMoved = true;
    if (m.fc == rookKSFile) whiteKSMoved = true;
  }
  if (moved == 'r' && m.fr == 7) {
    if (m.fc == rookQSFile) blackQSMoved = true;
    if (m.fc == rookKSFile) blackKSMoved = true;
  }
  if (m.captured == 'R' && m.tr == 0) {
    if (m.tc == rookQSFile) whiteQSMoved = true;
    if (m.tc == rookKSFile) whiteKSMoved = true;
  }
  if (m.captured == 'r' && m.tr == 7) {
    if (m.tc == rookQSFile) blackQSMoved = true;
    if (m.tc == rookKSFile) blackKSMoved = true;
  }
}

void commitMove(const Move &m) {
  const char moved = pos.at(m.fr, m.fc);
  if (m.isCapture()) {
    sound.capture();
    anim::ripple(m.tr, m.tc, CAPTURE_COLOR, 50);
  } else {
    sound.move();
  }
  noteRights(m, moved);
  pos.play(m);

  char uci[6];
  m.uci(uci);
  app.move(uci);
  sendPosition();
  app.log("%s: %s", pos.turn == 'b' ? "White" : "Black", uci);

  if (m.promotion) waitForPromotion(m.tr, m.tc);

  lights.clear();
  lights.show();
  board.sync();
}

void afterMove() {
  const char side = pos.turn;

  if (!pos.hasLegalMoves(side)) {
    Castle ks, qs;
    const bool canCastle = castleLegal(true, ks) || castleLegal(false, qs);
    if (!canCastle) {
      if (pos.inCheck(side)) {
        const bool whiteWon = side == 'b';
        app.result("checkmate", whiteWon ? "w" : "b");
        app.status("over", whiteWon ? "Checkmate - White wins" : "Checkmate - Black wins");
        anim::flashAll(CHECK_COLOR, 6, 200, 180);
        gameOverScreen(whiteWon ? WHITE_WINS : BLACK_WINS);
      } else {
        app.result("stalemate", "draw");
        app.status("over", "Stalemate - draw");
        anim::flashAll(Color(200, 150, 0), 6, 300, 200);
        gameOverScreen(Color(200, 150, 0));
      }
      awaitingSetup = true;
      return;
    }
  }

  if (pos.inCheck(side)) {
    sound.check();
    anim::flashAll(CHECK_COLOR, 3, 120, 120);
    app.status("playing", side == 'w' ? "White is in check" : "Black is in check");
  } else {
    char msg[64];
    snprintf(msg, sizeof(msg), "Chess960 #%d — %s to move", positionId,
             side == 'w' ? "White" : "Black");
    app.status("playing", msg);
  }
  showTurn();
}

void gameOverScreen(Color color) {
  for (int row = 0; row < 8; row++)
    for (int col = 0; col < 8; col++)
      lights.set(row, col, pos.at(row, col) != ' ' ? color : OFF);
  lights.show();
  while (!buttons.anyPressed()) service();
  lights.clear();
  lights.show();
}

void waitForPromotion(int row, int col) {
  sound.promote();
  app.status("playing", "Pawn promoted - swap in a queen");
  anim::columnChase(col, PROMO_COLOR);

  while (board.occupied(row, col)) anim::flashSquare(row, col, PROMO_COLOR, 1, 200);
  while (!board.occupied(row, col)) anim::flashSquare(row, col, PROMO_COLOR, 1, 200);
  app.log("Promotion complete");
}