/*
  This file is part of ct.

  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 ct. If not, see <https://www.gnu.org/licenses/>.
*/

#include "tak.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;

  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 ((state->colours[location] & 1) != 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
// ===================================================================

// Check for the presence of a road connecting opposite sides
static enum WIN_TYPE check_road_colour(tak_state_p state,
                                       const enum COLOUR colour) {
  int component[NUM_SQUARES(state->board_size)],
      touching[NUM_SQUARES(state->board_size)];

  /*
    We're doing a poor version of a disjoint set data structure to
    track and merge connected components. The array `component' stores
    indices to the representative cells of each connected component. A
    cell is representative if component[cell] = cell. We track only
    whether representatives are touching sides, and do a 2-dimensional
    DP approach to forming these from the board.

    Note: we don't track the rank/size of each tree, because we're not
    interested in good asymptotic complexity in the size of the board,
    merely good performance for a single board size in practice. For
    the same reason we also don't do path flattening/halving or
    anything.
  */
  for (int k = 0; k < NUM_SQUARES(state->board_size); k++) {
    component[k] = k; // every square is in its own connected
    // component initially
    touching[k] = 0; // and not connected to any sides
  }

  // touching is the bit mask for connectivity,
  // bottom | top | left | right
  //    1      2     4       8

  int touch = 5;
  for (int row = 0; row < state->board_size; row++) {
    for (int col = 0; col < state->board_size; col++) {
      const int cur = THE_COORDS(state->board_size, col, row);
      if (COUNT_AT(state, cur) && (state->colours[cur] & 1) == colour &&
          STONE_AT(state, cur) != STONE_STANDING) {

        // do we have any neighbours to the left and below?
        const int left_neighbour =
            ((cur % state->board_size > 0) &&
             (COUNT_AT(state, cur - 1)) // wont ever be out of bounds
             && ((state->colours[cur - 1] & 1) == colour) &&
             (STONE_AT(state, cur - 1) != STONE_STANDING));

        const int lowr_neighbour =
            ((cur >= state->board_size) &&
             (COUNT_AT(state, cur - state->board_size)) &&
             ((state->colours[cur - state->board_size] & 1) == colour) &&
             (STONE_AT(state, cur - state->board_size) != STONE_STANDING));

        // always take the component of the lower neighbour if
        // possible, failing that take the left neighbour, otherwise
        // we're not yet connected, so update our own component.
        if (lowr_neighbour) {
          // look up the root of the lower neighbour
          int root = cur - state->board_size;
          while (root != component[root])
            root = component[root];
          // join the set
          component[cur] = root;
          if (touch) {
            // something new
            touching[root] |= touch;
            // are we done?
            if ((touching[root] & 0x3) == 0x3 || (touching[root] & 0xC) == 0xC)
              return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
          }
          // if we also have a left neighbour then we should `merge'
          // sets, and here we assume that the left neighbour set is
          // always smaller (may not be) for the direction of merge
          if (left_neighbour) {
            int left_root = cur - 1;
            while (left_root != component[left_root])
              left_root = component[left_root];
            // merge
            const int left_touch = touching[left_root];
            if (left_touch) {
              touching[root] |= left_touch;
              if ((touching[root] & 0x3) == 0x3 ||
                  (touching[root] & 0xC) == 0xC)
                return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
            }
            component[left_root] = root;
          }
        } else if (left_neighbour) {
          int root = cur - 1;
          while (root != component[root])
            root = component[root];
          component[cur] = root;
          if (touch) {
            touching[root] |= touch;
            if ((touching[root] & 0x3) == 0x3 || (touching[root] & 0xC) == 0xC)
              return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
          }
        } else if (touch) {
          // we had no left or lower neighbour, so we're on our own
          touching[cur] = touch;
        }
      }
      if (col + 2 == state->board_size)
        touch |= 8;
      else
        touch &= 0x3;
    }
    if (row + 2 == state->board_size)
      touch = 6;
    else
      touch = 4;
  }
  return 0xFF;
}

enum WIN_TYPE check_win(tak_state_p state) {
  // Road?
  enum WIN_TYPE rb, rw;
  rb = check_road_colour(state, C_BLACK);
  rw = check_road_colour(state, C_WHITE);

  if (rb == WIN_ROAD_BLACK && rw == WIN_ROAD_WHITE) {
    return (state->ply & 1) ? rb : rw; // Dragons
  } else if (rw == WIN_ROAD_WHITE) {
    return rw;
  } else if (rb == WIN_ROAD_BLACK) {
    return rb;
  }

  // Do we do a flat count?
  int8_t total = 0, board_full = 1;
  for (uint8_t k = 0; k < NUM_SQUARES(state->board_size); k++) {
    if (COUNT_AT(state, k) == 0) {
      board_full = 0;
    } else if (STONE_AT(state, k) == STONE_FLAT) {
      total += ((state->colours[k] & 1) == C_BLACK) ? +1 : -1;
    }
  }
  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 < 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;
}