/*
  This file is part of ctak.

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with ctak. If not, see <https://www.gnu.org/licenses/>.
*/

#include "tak.h"
#include <stdio.h>
#include <stdlib.h>

// ===================================================================
// Helpers
// ===================================================================

#define NUM_SQUARES(board_size) (board_size * board_size)

tak_state_p new_tak_state(const uint8_t board_size) {
  tak_state_p state = malloc(sizeof(struct tak_state_s));
  reset_state(state, board_size);
  return state;
}

void free_tak_state(tak_state_p state) { free(state); }

// ===================================================================
// General state stuff
// ===================================================================

void reset_state(tak_state_p state, const uint8_t new_board_size) {
  if (new_board_size == 6) {
    state->board_size = 6;
    state->white_count = 128 | 30;
    state->black_count = 128 | 30;
  } else {
    state->board_size = 5;
    state->white_count = 128 | 21;
    state->black_count = 128 | 21;
  }

  state->ply = 0;
  state->won = 0xFF; // i may live to regret this hack
  state->current_colour = C_BLACK;
  state->hash = 0; // empty board, ply 0

  for (uint8_t k = 0; k < NUM_SQUARES(new_board_size); k++) {
    state->celldat[k] = 0;
  }
}

void next_ply(tak_state_p state) {
  state->ply++;
  if (state->ply == 2) {
    state->current_colour = C_WHITE;
  } else {
    if (state->current_colour == C_BLACK)
      state->current_colour = C_WHITE;
    else
      state->current_colour = C_BLACK;
  }
}

// ===================================================================
// Placing stones
// ===================================================================

enum ACT_RESULT try_place(tak_state_p state, const int8_t location,
                          const enum COLOUR colour,
                          const enum STONE_VARIANT stone) {
  // Game is over?
  if (state->won < 0xFF)
    return GAME_END;
  // Can't place on an occupied square
  if (COUNT_AT(state, location)) {
    return ACT_ILLEGAL;
  } else {
    switch (stone) {
    case STONE_STANDING:
      if (state->ply < 2)
        return ACT_ILLEGAL;
    // behold the magic GCC comment which defeates
    // -Wimplicit-fallthrough:
    // fall through
    case STONE_FLAT: {
      if (colour == C_BLACK) {
        if (state->black_count & 127)
          state->black_count--;
        else
          return ACT_ILLEGAL;
      } else {
        if (state->white_count & 127)
          state->white_count--;
        else
          return ACT_ILLEGAL;
      }
      break;
    }
    case STONE_CAPSTONE: {
      if (state->ply < 2)
        return ACT_ILLEGAL;
      if (colour == C_BLACK) {
        if (state->black_count & 128)
          state->black_count &= 127;
        else
          return ACT_ILLEGAL;
      } else {
        if (state->white_count & 128)
          state->white_count &= 127;
        else
          return ACT_ILLEGAL;
      }
      break;
    }
    }

    state->colours[location] = colour;
    state->celldat[location] = NUM_INC | stone;
    return ACT_OK;
  }
}

// ===================================================================
// Moving stacks
// ===================================================================

static inline void push_stones(tak_state_p state, const int8_t location,
                               const uint8_t count, const uint8_t new_colours,
                               const enum STONE_VARIANT top_stone) {
  state->colours[location] = (state->colours[location] << count) | new_colours;
  state->celldat[location] =
      top_stone |
      ((state->celldat[location] + ((count << NUM_SHIFT))) & NUM_MASK);
}

enum ACT_RESULT try_move(tak_state_p state, const int8_t location,
                         const enum MOVE_DIRECTION direction,
                         const uint8_t steps, const uint8_t drops[5]) {
  // Game is over?
  if (state->won < 0xFF)
    return GAME_END;
  // Can't do this
  if (steps == 0 || steps > state->board_size)
    return ACT_ILLEGAL;
  // Check for stones at all
  const uint8_t avail = COUNT_AT(state, location);
  if (avail == 0)
    return ACT_ILLEGAL;
  // Does the current player own the pile?
  if (TOP_COLOUR_AT(state, location) != state->current_colour)
    return ACT_ILLEGAL;
  // Is the desired direction and count on the board?
  int8_t delta = 0;
  switch (direction) {
  case M_UP: {
    delta = +state->board_size;
    if (location + delta * steps > NUM_SQUARES(state->board_size))
      return ACT_ILLEGAL;
    break;
  };
  case M_DOWN: {
    delta = -state->board_size;
    if (location + delta * steps < 0)
      return ACT_ILLEGAL;
    break;
  };
  case M_RIGHT: {
    delta = +1;
    if ((location + steps * delta) / state->board_size >
        location / state->board_size)
      return ACT_ILLEGAL;
    break;
  };
  case M_LEFT: {
    delta = -1;
    // We need the extra check for zero here because, irritatingly,
    // -1 / board_size == 1 / board_size
    if ((location + steps * delta < 0) ||
        ((location + steps * delta) / state->board_size <
         location / state->board_size))
      return ACT_ILLEGAL;
    break;
  };
  };

  // For every square in the direction
  uint8_t total = 0;
  for (uint8_t k = 0; k < steps; k++) {
    // Can't drop 0 anywhere because we're past the first square
    if (drops[k] == 0)
      return ACT_ILLEGAL;
    // Can't drop more than BOARD_SIZE stones in a square
    if (drops[k] > state->board_size)
      return ACT_ILLEGAL;
    // Check for overflows
    if (COUNT_AT(state, location + (k + 1) * delta) + drops[k] > 0x0F)
      return ACT_OVERFLOW;
    // Check for capstone
    if (STONE_AT(state, location + (k + 1) * delta) == STONE_CAPSTONE)
      return ACT_ILLEGAL;
    // Check for wall
    if ((STONE_AT(state, location + (k + 1) * delta) == STONE_STANDING)
        // If not last drop, or not dropping just one, or not a cap
        && ((k + 1 < steps) || (drops[k] != 1) ||
            (STONE_AT(state, location) != STONE_CAPSTONE)))
      return ACT_ILLEGAL;
    total += drops[k];
  }

  // Can't ask to move 0, more than board_size, or stones available
  if ((total == 0) || (total > state->board_size) || (total > avail))
    return ACT_ILLEGAL;

  // Nothing illegal, do it. First we add the stones to the
  // destination squares
  uint8_t j = total;
  for (uint8_t k = 0; k < steps; k++) {
    j -= drops[k];
    push_stones(state, location + (k + 1) * delta, drops[k],
                (state->colours[location] >> j) & (0xFFFF >> (0x10 - drops[k])),
                (k == steps - 1) ? STONE_AT(state, location) : STONE_FLAT);
  }
  // Then we drop them from the source
  state->colours[location] >>= total;
  const uint8_t dec_count = state->celldat[location] - (total << NUM_SHIFT);
  state->celldat[location] = dec_count & NUM_MASK;

  return ACT_OK;
}

// ===================================================================
// Checking win
// ===================================================================

// One bit per colour, set once a road for that colour has been found
#define ROAD_WHITE (1u << C_WHITE)
#define ROAD_BLACK (1u << C_BLACK)
#define ROAD_BOTH (ROAD_WHITE | ROAD_BLACK)

// One bit per square, marking the cells on a given side of the board.
// Indexed by board_size - 5 (0 for 5x5, 1 for 6x6):
//   5x5: bottom bits 0-4, top bits 20-24,
//        left bits 0,5,10,15,20, right bits 4,9,14,19,24
//   6x6: bottom bits 0-5, top bits 30-35,
//        left bits 0,6,12,18,24,30, right bits 5,11,17,23,29,35
static const uint64_t side_bottom_row[2] = {0x000000000000001FULL,
                                            0x000000000000003FULL};
static const uint64_t side_top_row[2] = {0x0000000001F00000ULL,
                                         0x00000FC0000000ULL};
static const uint64_t side_left_col[2] = {0x0000000000108421ULL,
                                          0x000000041041041ULL};
static const uint64_t side_right_col[2] = {0x000000001084210ULL,
                                           0x000820820820ULL};
/*
  Flood fill, restricted to the squares set in `mask', starting from the
  squares set in `start'. Each iteration grows the fill by one square in all
  four directions (the left/right edge masks stop the horizontal shifts from
  wrapping to the next row), so the fill reaches a fixed point within the
  board's diameter, 2*(board_size - 1), iterations.
*/
static inline uint64_t side_flood(const uint64_t mask, const uint64_t start,
                                  const uint8_t board_size,
                                  const uint64_t left,
                                  const uint64_t right) {
  uint64_t fill = start & mask;
  for (uint8_t i = 0; i < 2 * (board_size - 1); i++) {
    // << 1 moves toward the right column (and would wrap to the next row's
    // left edge, hence the ~right mask); >> 1 is the mirror image
    const uint64_t grow =
        ((fill & ~right) << 1) | ((fill & ~left) >> 1) |
        (fill << board_size) | (fill >> board_size);
    const uint64_t nf = fill | (grow & mask);
    if (nf == fill)
      break;
    fill = nf;
  }
  return fill;
}

enum WIN_TYPE check_win(tak_state_p state) {
  const uint8_t board_size = state->board_size;
  const uint8_t si = board_size - 5;
  const uint64_t bottom = side_bottom_row[si], top = side_top_row[si],
      left = side_left_col[si], right = side_right_col[si];

  /*
    A cell is road-eligible only for the colour of its top stone, and only
    if that stone is not a standing stone. We build one bitset of
    road-eligible squares per colour, then flood fill from each pair of
    opposite sides: a road exists for a colour if the fill from the left
    column (or the bottom row) reaches the opposite side. Both colours are
    checked in a single pass over the board, which also counts top-flat
    stones. The dragon rule needs the road status of both colours.
  */
  uint64_t m[2] = {0, 0};

  // flat count: +1 per top-flat black stone, -1 per top-flat white
  int8_t total = 0;
  uint8_t board_full = 1;

  for (uint32_t k = 0; k < NUM_SQUARES(board_size); k++) {
    const uint8_t cd = state->celldat[k];
    const uint8_t count = cd >> NUM_SHIFT;
    if (count == 0) {
      board_full = 0;
    } else {
      const uint8_t stone = cd & STONE_MASK;
      if (stone == STONE_FLAT)
        total += (state->colours[k] & 1) ? +1 : -1;
      if (stone != STONE_STANDING)
        m[state->colours[k] & 1] |= 1ULL << k;
    }
  }

  uint8_t road = 0;
  for (uint8_t c = 0; c < 2; c++) {
    if (m[c] == 0)
      continue;
    if ((side_flood(m[c], left, board_size, left, right) & right) != 0 ||
        (side_flood(m[c], bottom, board_size, left, right) & top) != 0)
      road |= 1u << c;
  }

  const enum WIN_TYPE rw = (road & ROAD_WHITE) ? WIN_ROAD_WHITE : 0xFF,
                      rb = (road & ROAD_BLACK) ? WIN_ROAD_BLACK : 0xFF;

  if (rb == WIN_ROAD_BLACK && rw == WIN_ROAD_WHITE) {
    return WIN_DRAGON;
  } else if (rw == WIN_ROAD_WHITE) {
    return rw;
  } else if (rb == WIN_ROAD_BLACK) {
    return rb;
  }

  if (state->black_count == 0 || state->white_count == 0 || board_full) {
    // Decide based on count
    if (total > 0)
      return WIN_FLAT_BLACK;
    else if (total < 0)
      return WIN_FLAT_WHITE;
    else
      return WIN_DRAW;
  }

  return 0xFF;
}
// ===================================================================
// PTN place parser
// ===================================================================

#define ASSERT_NONEMPTY                                                        \
  {                                                                            \
    if (ptn == NULL || *ptn == 0)                                              \
      return PTN_INVALID;                                                      \
  }

#define ASSERT_MORE                                                            \
  {                                                                            \
    if (*ptn == 0)                                                             \
      return PTN_INVALID;                                                      \
  }

enum PTN_RESULT parse_place(const uint8_t board_size, char *ptn,
                            uint8_t *out_location,
                            enum STONE_VARIANT *out_stone) {

  ASSERT_NONEMPTY;

  *out_stone = STONE_FLAT;
  switch (*ptn) {
  case 'C': {
    ptn++;
    *out_stone = STONE_CAPSTONE;
    break;
  };
  case 'S': {
    ptn++;
    *out_stone = STONE_STANDING;
    break;
  };
  case 'F': {
    ptn++;
    break;
  };
  }

  ASSERT_MORE;

  if ((*ptn < 'a') || (*ptn > '`' + board_size))
    return PTN_INVALID;
  *out_location = *ptn - 'a';

  ptn++;
  ASSERT_MORE;

  if ((*ptn < '1') || (*ptn > board_size + '0'))
    return PTN_INVALID;
  *out_location += board_size * (*ptn - '1');

  if (*(++ptn) > 0)
    return PTN_INVALID;

  return PTN_OK;
}

// ===================================================================
// PTN move parser
// ===================================================================

enum PTN_RESULT parse_move(const uint8_t board_size, char *ptn,
                           uint8_t *out_location,
                           enum MOVE_DIRECTION *out_direction,
                           uint8_t *out_steps, uint8_t out_drops[5]) {

  ASSERT_NONEMPTY;

  uint8_t picked_up = 1;

  // Optionally indicate how many stones picked up
  if ((*ptn >= '1') && (*ptn <= '0' + board_size)) {
    picked_up = *ptn - '0';
    ptn++;
    ASSERT_MORE;
  }

  // column must be on the board
  if ((*ptn < 'a') || (*ptn > '`' + board_size))
    return PTN_INVALID;
  *out_location = *ptn - 'a';

  ptn++;
  ASSERT_MORE;

  // row must be on the board
  if ((*ptn < '1') || (*ptn > board_size + '0'))
    return PTN_INVALID;
  *out_location += board_size * (*ptn - '1');

  ptn++;
  ASSERT_MORE;

  // valid direction
  switch (*ptn) {
  case '+': {
    *out_direction = M_UP;
    break;
  }
  case '-': {
    *out_direction = M_DOWN;
    break;
  }
  case '<': {
    *out_direction = M_LEFT;
    break;
  }
  case '>': {
    *out_direction = M_RIGHT;
    break;
  }
  default:
    return PTN_INVALID;
  }

  // Handle the case 'n<column><row><direction>' as
  // 'n<column><row><direction>n' for convenience, if n is omitted
  // assume n = 1
  ptn++;
  if (*ptn == 0) {
    *out_steps = 1;
    out_drops[0] = picked_up;
    return PTN_OK;
  }

  // Parse the drops in each subsequent square
  *out_steps = 0;
  uint8_t total = 0;
  while (*ptn) {
    // can't drop more than the carry limit, or less than 1
    if ((*ptn < '1') || (*ptn > '0' + board_size))
      return PTN_INVALID;

    // can't move more than the size of the board in any direction
    if ((*out_steps + 1 >= board_size) && *ptn)
      return PTN_INVALID;

    out_drops[*out_steps] = *ptn - '0';
    total += out_drops[*out_steps];
    *out_steps += 1;
    ptn++;
  }

  // Mismatch between number of stones picked up and total dropped
  if (total != picked_up)
    return PTN_INVALID;

  return PTN_OK;
}

// ===================================================================
// Generate PTN for place
// ===================================================================

void generate_place(const uint8_t board_size, const uint8_t in_location,
                    const enum STONE_VARIANT in_stone, char out_ptn[4]) {
  switch (in_stone) {
  case STONE_FLAT: {
    break;
  }
  case STONE_STANDING: {
    *out_ptn = 'S';
    out_ptn++;
    break;
  }
  case STONE_CAPSTONE: {
    *out_ptn = 'C';
    out_ptn++;
    break;
  }
  }
  *out_ptn = 'a' + (in_location % board_size);
  out_ptn++;
  *out_ptn = '1' + (in_location / board_size);
  out_ptn++;
  *out_ptn = 0;
}

// ===================================================================
// Generate PTN for move
// ===================================================================

void generate_move(const uint8_t board_size, const uint8_t in_location,
                   const enum MOVE_DIRECTION in_direction,
                   const uint8_t in_steps, const uint8_t in_drops[5],
                   char out_ptn[10]) {
  uint8_t total = 0;
  for (uint8_t k = 0; k < in_steps; k++)
    total += in_drops[k];
  if (total > 1) {
    *out_ptn = '0' + total;
    out_ptn++;
  }

  *out_ptn = 'a' + (in_location % board_size);
  out_ptn++;
  *out_ptn = '1' + (in_location / board_size);
  out_ptn++;

  switch (in_direction) {
  case M_UP: {
    *out_ptn = '+';
    break;
  }
  case M_DOWN: {
    *out_ptn = '-';
    break;
  }
  case M_LEFT: {
    *out_ptn = '<';
    break;
  }
  case M_RIGHT: {
    *out_ptn = '>';
    break;
  }
  };
  out_ptn++;

  for (uint8_t k = 0; (total > 1) && (k < in_steps); k++) {
    *out_ptn = '0' + in_drops[k];
    out_ptn++;
  }

  *out_ptn = 0;
}

// ===================================================================
// Driver
// ===================================================================

static uint8_t is_not_placement(char *ptn) {
  if (ptn == 0)
    return 0;
  for (;; ptn++) {
    switch (*ptn) {
    case '+':
    case '-':
    case '>':
    case '<':
      return 1;
    case 0:
      return 0;
    }
  }
}

enum ACT_RESULT do_ptn(tak_state_p state, char *ptn) {
  // Game over?
  if (state->won < 0xFF)
    return GAME_END;

  enum PTN_RESULT ptn_res;
  enum ACT_RESULT act_res;
  uint8_t location;

  // Placing or moving?
  if (is_not_placement(ptn)) {
    uint8_t steps, drops[5];
    enum MOVE_DIRECTION direction;
    // Parse it as a move
    ptn_res = parse_move(state->board_size, ptn, &location, &direction, &steps,
                         drops);
    // If valid PTN, try to do it
    if (ptn_res == PTN_OK) {
      if (state->ply < 2)
        return ACT_ILLEGAL;
      act_res = try_move(state, location, direction, steps, drops);
    } else {
      return ACT_INVALID_PTN;
    }
  } else {
    // It was not a move
    enum STONE_VARIANT stone;
    // Was it a valid placement?
    ptn_res = parse_place(state->board_size, ptn, &location, &stone);
    // If so, try it
    if (ptn_res == PTN_OK)
      act_res = try_place(state, location, state->current_colour, stone);
    else
      return ACT_INVALID_PTN;
  }
  // A valid ply occured
  if (act_res == ACT_OK) {
    // Don't bother checking that the game was won early on, could be
    // more conservative here :)
    if (state->ply >= state->board_size) {
      state->won = check_win(state);
      if (state->won == WIN_DRAGON)
        state->won = (state->current_colour == C_BLACK)
                         ? WIN_ROAD_BLACK
                         : WIN_ROAD_WHITE;
      if (state->won < 0xFF) {
        // Winning move, but no need to update current colour
        state->ply++;
        return GAME_END;
      }
    }
    // Only step if the game isn't over yet
    next_ply(state);
  }
  return act_res;
}