// ============================================================
// Atomic Chess — Open Chess
//
// Chess, except captures detonate. Take a piece and the blast removes the
// piece you took, the piece you took it with, and everything standing on the
// eight squares around it — your own pieces included. Pawns are the one
// exception: a pawn only dies if it was the piece captured, so a pawn wall
// walks out of a blast that levels everything else around it.
//
// You do not need checkmate. Catch the enemy king in an explosion and the
// game is over. The flip side is that the board will not let you play a move
// that destroys your own king, which is also why a king can never capture.
//
//   Button A  restart (back to the setup screen)
//   Button B  pause: the board goes dark until A is pressed
//
// Captures are signalled the same way as in Classic Chess: take the piece you
// are capturing off the board first. In Atomic that is all the board needs,
// because your capturing piece is destroyed too and never lands — both pieces
// go straight in the box.
//
// After a blast the board holds the game and pulses red under every piece the
// explosion killed. Take them off and play carries on. Nobody has to agree on
// what the blast hit; you just clear whatever is still glowing.
//
// 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>

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

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

const Color HINT_COLOR    = GREEN;              // empty squares the piece can reach
const Color CAPTURE_COLOR = ORANGE;             // a capture, and so an explosion
const Color ILLEGAL_COLOR = AMBER;              // where you may not put it
const Color CHECK_COLOR   = RED;                // the board flashes this on check
const Color PROMO_COLOR   = GOLD;               // pawn reached the last rank
const Color TURN_COLOR    = GREEN.scaled(70);   // soft glow under the side to move
const Color WHITE_WINS    = WHITE;
const Color BLACK_WINS    = PURPLE;

// The blast, from the inside out: a white-hot core, orange behind the shock
// front, and dull red where it has already burned through.
const Color FUSE_COLOR  = Color(255, 40, 0);
const Color BLAST_CORE  = Color(255, 210, 60, 180);
const Color BLAST_MID   = Color(255, 80, 0);
const Color BLAST_EDGE  = Color(120, 16, 0);
const Color DEBRIS      = RED;                  // a dead piece still on the board

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

Position pos;
bool     awaitingSetup = true;
bool     paused        = false;


// Every square whose piece an explosion destroys.
struct Blast {
  bool doomed[8][8];
  int  count;        // pieces killed
};

void setup() {
  begin("atomic-chess", "Atomic Chess", "2.0.0");
  app.onIdentify = sendPosition;
  randomSeed(micros());
}

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

  // Look for a piece of the side to move that has just been lifted.
  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;
      if (!waitForMove(row, col, move)) return;
      commitMove(move);
      afterMove();
      return;
    }
  }
  board.sync();
}

// ============================================================
// THE ATOMIC RULES
//
// Layered on top of the core's ordinary move generator rather than replacing
// it. Pieces still move the way they always did; what changes is which of
// those moves you are allowed to play, and what is left standing afterwards.
// ============================================================

bool isPawn(char piece) { return chess::kindOf(piece) == 'P'; }

// Marks every square a capture would clear. Reads the board as it stands
// before the move, which is what makes pawn immunity work: a pawn next to the
// blast is still a pawn when the blast arrives.
void blastFor(const Position &p, const Move &m, Blast &out) {
  memset(out.doomed, 0, sizeof(out.doomed));
  out.count = 0;

  auto kill = [&](int row, int col) {
    if (!onBoard(row, col) || out.doomed[row][col]) return;
    out.doomed[row][col] = true;
    out.count++;
  };

  kill(m.fr, m.fc);   // the capturing piece goes with its target
  kill(m.tr, m.tc);   // the piece that was taken

  for (int dr = -1; dr <= 1; dr++) {
    for (int dc = -1; dc <= 1; dc++) {
      if (dr == 0 && dc == 0) continue;
      const int row = m.tr + dr, col = m.tc + dc;
      if (!onBoard(row, col)) continue;
      const char neighbour = p.at(row, col);
      if (neighbour == ' ' || isPawn(neighbour)) continue;   // pawns ride it out
      kill(row, col);
    }
  }
}

bool hasKing(const Position &p, char color) {
  const char king = color == 'w' ? 'K' : 'k';
  for (int row = 0; row < 8; row++)
    for (int col = 0; col < 8; col++)
      if (p.squares[row][col] == king) return true;
  return false;
}

// Applies a capture: everything in the blast comes off, and the turn passes.
void applyBlast(Position &p, const Move &m) {
  Blast blast;
  blastFor(p, m, blast);
  for (int row = 0; row < 8; row++)
    for (int col = 0; col < 8; col++)
      if (blast.doomed[row][col]) p.squares[row][col] = ' ';
  p.turn = chess::other(p.turn);
}

// The position a move leaves behind.
Position after(const Position &p, const Move &m) {
  Position next = p;
  if (m.isCapture()) applyBlast(next, m);
  else               next.play(m);
  return next;
}

// Atomic check is not ordinary check.
//
// A piece only threatens your king if taking it would leave its own king
// standing. That one condition covers the two places atomic chess parts ways
// with the rules people expect: the kings may stand next to each other,
// because neither can take the other without dying with it — and a piece is
// just as dangerous next to your king as it is attacking it, because a capture
// anywhere alongside the king takes the king too.
bool inAtomicCheck(const Position &p, char color) {
  static Move replies[chess::MAX_MOVES];

  const char enemy = chess::other(color);
  const int  count = p.pseudoMoves(enemy, replies);

  for (int i = 0; i < count; i++) {
    const Position next = after(p, replies[i]);
    if (!hasKing(next, color) && hasKing(next, enemy)) return true;
  }
  return false;
}

bool isAtomicLegal(const Position &p, const Move &m) {
  const char mover = chess::colorOf(p.at(m.fr, m.fc));
  const Position next = after(p, m);

  // Your own king caught in your own blast. Never allowed, and the reason a
  // king cannot capture: its own explosion would always take it.
  if (!hasKing(next, mover)) return false;

  // Taking the enemy king wins on the spot, so it outranks any check you
  // happen to be in while doing it.
  if (!hasKing(next, chess::other(mover))) return true;

  return !inAtomicCheck(next, mover);
}

// The legal moves for one piece — what the LED hints show.
int atomicMovesFrom(const Position &p, int row, int col, Move *out) {
  static Move all[chess::MAX_MOVES];

  const char piece = p.at(row, col);
  if (piece == ' ') return 0;

  const int total = p.pseudoMoves(chess::colorOf(piece), all);
  int count = 0;
  for (int i = 0; i < total; i++) {
    if (all[i].fr != row || all[i].fc != col) continue;
    if (!isAtomicLegal(p, all[i])) continue;
    out[count++] = all[i];
  }
  return count;
}

bool hasAtomicMove(const Position &p, char color) {
  static Move all[chess::MAX_MOVES];

  const int total = p.pseudoMoves(color, all);
  for (int i = 0; i < total; i++)
    if (isAtomicLegal(p, all[i])) return true;
  return false;
}

// ============================================================
// THE EXPLOSION
//
// A fuse, a blast, and then the board asking for the wreckage to be cleared.
// ============================================================

// The buzzer can only make tones, so the bang is a fast tumble down the
// register with the pitch knocked about a little. It will never be a real
// explosion, but it falls fast enough to read as one — much better than any
// single note, and better than anything the shared sound library should be
// carrying around for one game's sake.
//
// sound.beep() hands the note to the buzzer's own task and returns straight
// away, so each frame of the animation fires the next note and the two stay in
// step without either waiting on the other.

// Shows what is about to die while the fuse burns down, accelerating into the
// blast. This is the beat before you commit: if the explosion is going to cost
// you your own queen, you can see her lit up before it goes off.
void fuse(const Blast &blast, int row, int col) {
  static const uint16_t STEP_MS[] = {240, 190, 150, 115, 85, 62, 44, 30};
  const int steps = sizeof(STEP_MS) / sizeof(STEP_MS[0]);

  for (int i = 0; i < steps; i++) {
    sound.beep(1100 + i * 130, 22);

    for (int lit = 1; lit >= 0; lit--) {
      lights.clear();
      if (lit) {
        for (int r = 0; r < 8; r++)
          for (int c = 0; c < 8; c++)
            if (blast.doomed[r][c]) lights.set(r, c, BLAST_EDGE);
        lights.set(row, col, FUSE_COLOR);
      }
      lights.show();
      wait(STEP_MS[i] / 2);
    }
  }
}

// An expanding shock front. Each frame the front moves out one square; the
// squares it has already passed cool from white-hot to orange to dull red.
void detonation(int row, int col) {
  const int FRAMES = 8;

  for (int t = 0; t < FRAMES; t++) {
    const int pitch = 640 - t * 80 + (int)random(-50, 50);
    sound.beep((uint16_t)max(70, pitch), 62);

    lights.clear();
    for (int r = 0; r < 8; r++) {
      for (int c = 0; c < 8; c++) {
        // Square distance, so the front travels as a box and every one of the
        // eight neighbours is caught on the same frame.
        const int distance = max(abs(r - row), abs(c - col));
        if (distance > t) continue;

        switch (t - distance) {      // frames since the front went through
          case 0:  lights.set(r, c, BLAST_CORE); break;
          case 1:  lights.set(r, c, BLAST_MID);  break;
          case 2:  lights.set(r, c, BLAST_EDGE); break;
          default: break;              // burned out
        }
      }
    }
    lights.show();
    wait(62);
  }

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


// The board cannot take pieces off itself, so it asks. Every square still
// holding a destroyed piece pulses red until that piece is gone, and a
// survivor picked up by mistake pulses amber until it is put back — which
// keeps the board's idea of the position and the actual pieces from drifting
// apart in the middle of the mess.
void clearWreckage(const Blast &blast) {
  app.status("playing", "Take off the pieces the blast destroyed");

  for (;;) {
    service();

    // Pulse, so a lit square reads as a request rather than as decoration.
    const bool lit = (millis() / 220) % 2 == 0;
    bool waiting = false;

    lights.clear();
    for (int row = 0; row < 8; row++) {
      for (int col = 0; col < 8; col++) {
        if (blast.doomed[row][col]) {
          if (!board.occupied(row, col)) continue;
          waiting = true;
          if (lit) lights.set(row, col, DEBRIS);
        } else if (pos.at(row, col) != ' ' && !board.occupied(row, col)) {
          waiting = true;
          if (lit) lights.set(row, col, ILLEGAL_COLOR);
        }
      }
    }
    lights.show();

    if (!waiting) break;
    if (buttons.pressed(BUTTON_A)) { restart(); return; }
    wait(40);
  }

  board.sync();
}

// ============================================================
// PLAYING A MOVE
// ============================================================

// Shows where the lifted piece may go and waits for it to land. Returns false
// if the piece went back to its square or a button was pressed.
bool waitForMove(int fromRow, int fromCol, Move &out) {
  static Move options[chess::MAX_MOVES];
  const int count = atomicMovesFrom(pos, fromRow, fromCol, options);
  showHints(fromRow, fromCol, options, count);

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

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

    for (int i = 0; i < count; i++) {
      const int toRow = options[i].tr, toCol = options[i].tc;

      if (options[i].isCapture()) {
        // Taking the target off the board is the whole signal. Your capturing
        // piece is already in your hand and is about to be destroyed, so
        // unlike Classic Chess there is nothing left to wait for it to land
        // on — both pieces are out of the game from here.
        if (board.lifted(toRow, toCol)) {
          out = options[i];
          return true;
        }
      } else if (board.placed(toRow, toCol)) {
        out = options[i];
        return true;
      }
    }

    // Landed somewhere illegal: flash it until the piece is lifted again.
    for (int row = 0; row < 8; row++) {
      for (int col = 0; col < 8; col++) {
        if (!board.placed(row, col) || pos.at(row, col) != ' ') 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);
      }
    }
    wait(20);
  }
}

void showHints(int fromRow, int fromCol, Move *options, int count) {
  lights.clear();
  lights.set(fromRow, fromCol, WHITE);
  for (int i = 0; i < count; i++) {
    lights.set(options[i].tr, options[i].tc,
               options[i].isCapture() ? CAPTURE_COLOR
               : options[i].promotion ? PROMO_COLOR
               : HINT_COLOR);
  }
  lights.show();
}

void commitMove(const Move &m) {
  const char mover = pos.turn;
  char uci[6];
  m.uci(uci);

  if (!m.isCapture()) {
    sound.move();
    pos.play(m);
    app.move(uci);
    app.log("%s: %s", mover == 'w' ? "White" : "Black", uci);
    if (m.promotion) waitForPromotion(m.tr, m.tc);
    lights.clear();
    lights.show();
    board.sync();
    return;
  }

  Blast blast;
  blastFor(pos, m, blast);

  sound.capture();
  fuse(blast, m.tr, m.tc);
  detonation(m.tr, m.tc);

  applyBlast(pos, m);

  app.move(uci);
  app.log("%s: %s - %d pieces destroyed", mover == 'w' ? "White" : "Black", uci, blast.count);
  sendPosition();

  // A king in the blast means the game is already over; the pieces come off
  // during setup for the next one, so there is no point asking for them now.
  const bool kingDown = !hasKing(pos, 'w') || !hasKing(pos, 'b');
  if (!kingDown) clearWreckage(blast);
}

// Explosions, check, checkmate, stalemate — and the status line for the
// Connect page.
void afterMove() {
  const char side = pos.turn;

  // Atomic does not need checkmate. A king caught in a blast ends it here.
  if (!hasKing(pos, side)) {
    const bool whiteWon = side == 'b';
    sound.win();
    app.result("king-destroyed", whiteWon ? "w" : "b");
    app.status("over", whiteWon ? "King destroyed - White wins" : "King destroyed - Black wins");
    anim::flashAll(CHECK_COLOR, 6, 180, 160);
    gameOverScreen(whiteWon ? WHITE_WINS : BLACK_WINS);
    awaitingSetup = true;
    return;
  }

  if (!hasAtomicMove(pos, side)) {
    if (inAtomicCheck(pos, side)) {
      const bool whiteWon = side == 'b';
      sound.checkmate();
      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 {
      sound.draw();
      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 (inAtomicCheck(pos, side)) {
    sound.check();
    anim::flashAll(CHECK_COLOR, 3, 120, 120);
    app.status("playing", side == 'w' ? "White is in check" : "Black is in check");
  } else {
    app.status("playing", side == 'w' ? "White to move" : "Black to move");
  }
  showTurn();
}

// ============================================================
// SETUP, TURNS, AND THE END OF THE GAME
// ============================================================

void waitForSetup() {
  app.status("setup", "Place all 32 pieces on their starting squares");
  lights.clear();
  lights.show();

  while (!board.covers(chess::START_LAYOUT)) {
    service();
    if (buttons.pressed(BUTTON_B)) { pauseGame(); return; }
    for (int row = 0; row < 8; row++)
      for (int col = 0; col < 8; col++)
        lights.set(row, col, chess::START_LAYOUT[row][col] != ' ' && !board.occupied(row, col) ? WHITE : OFF);
    lights.show();
    wait(80);
  }

  pos.reset();
  anim::ripple(3, 3, WHITE, 45);
  board.sync();
  awaitingSetup = false;

  sound.gameStart();
  app.status("playing", "White to move");
  sendPosition();
  showTurn();
}

// Lights every piece of the side to move. This is the resting picture of the
// board between moves.
void showTurn() {
  lights.clear();
  for (int row = 0; row < 8; row++)
    for (int col = 0; col < 8; col++) {
      const char piece = pos.at(row, col);
      if (piece != ' ' && chess::colorOf(piece) == pos.turn) lights.set(row, col, TURN_COLOR);
    }
  lights.show();
}

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

void pauseGame() {
  // Count down along rank 1 so an accidental press is obvious, then go dark.
  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);
}

// Holds a colour on every occupied square until a button is pressed.
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();
}

// Only reachable on a quiet move: a capturing pawn is destroyed by its own
// blast long before it gets to promote.
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);   // pawn off
  while (!board.occupied(row, col)) anim::flashSquare(row, col, PROMO_COLOR, 1, 200);  // queen on
  app.log("Promotion complete");
}