#include "tak.h"

// -------------------------------------------------------------------
// Helpers

#define NUM_INC       (0x1<<NUM_SHIFT) // 0b00010000
#define NUM_MASK      (0xF<<NUM_SHIFT) // 0b11110000
#define DFS_MASK      0b00001100
#define NOT_DFS_MASK  0b11110011

#define NUM_SQUARES   (board_size * board_size)

// -------------------------------------------------------------------
// General state stuff

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

  ply = 0;
  won = -1;
  current_colour = C_BLACK;

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

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

// -------------------------------------------------------------------
// Place stone

enum E_RESULT
try_place(const int8_t location, const enum COLOUR colour,
          const enum STONE_VARIANT stone)
{
  // Game is over?
  if (won < -1) return GAME_END;
  // Can't place on an occupied square
  if (COUNT_AT(location)) {
    return ACT_ILLEGAL;
  } else {
    switch (stone) {
      case STONE_STANDING: if (ply < 2) return ACT_ILLEGAL;
      case STONE_FLAT: {
        if (colour == C_BLACK) {
          if (black_count & 127) black_count--;
          else return ACT_ILLEGAL;
        } else {
          if (white_count & 127) white_count--;
          else return ACT_ILLEGAL;
        }
        break;
      }
      case STONE_CAPSTONE: {
        if (ply < 2) return ACT_ILLEGAL;
        if (colour == C_BLACK) {
          if (black_count & 128) black_count &= 127;
          else return ACT_ILLEGAL;
        } else {
          if (white_count & 128) white_count &= 127;
          else return ACT_ILLEGAL;
        }
        break;
      }
    }

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

// -------------------------------------------------------------------
// Move stack

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

static void
drop_stones(const int8_t location, const uint8_t count) {
  // Calling this with count = 0 is destructive
  colours[location] >>= count;
  const uint8_t dec_count = celldat[location] - (count << NUM_SHIFT);
  celldat[location] = dec_count & NUM_MASK;
}

enum E_RESULT
try_move(const int8_t location, const enum MOVE_DIRECTION direction,
         const uint8_t steps, const uint8_t drops[5]) {
  // Game is over?
  if (won < -1) return GAME_END;
  // Can't do this
  if (steps == 0 || steps > 5) return ACT_ILLEGAL;
  // Is the desired direction and count on the board?
  int8_t delta = 0;
  switch (direction) {
    case M_UP:    {
      delta = +board_size;
      if (location + delta * steps > NUM_SQUARES) return ACT_ILLEGAL;
      break;
    };
    case M_DOWN:  {
      delta = -board_size;
      if (location + delta * steps < 0) return ACT_ILLEGAL;
      break;
    };
    case M_RIGHT: {
      delta = +1;
      if (location + delta * steps > NUM_SQUARES) return ACT_ILLEGAL;
      break;
    };
    case M_LEFT:  {
      delta = -1;
      if (location + delta * steps < 0) 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] > board_size)
      return ACT_ILLEGAL;
    // Check for overflows
    if (COUNT_AT(location+(k+1)*delta) + drops[k] > 0x0F)
      return ACT_OVERFLOW;
    // Check for capstone
    if (STONE_AT(location+(k+1)*delta) == STONE_CAPSTONE)
      return ACT_ILLEGAL;
    // Check for wall
    if ( (STONE_AT(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(location) != STONE_CAPSTONE) )
         )
      return ACT_ILLEGAL;
    total += drops[k];
  }

  // Can't ask to move more than BOARD_SIZE or stones available
  if ( (total > board_size) || (total > COUNT_AT(location)) )
    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(location+(k+1)*delta,
                drops[k],
                (colours[location] >> j) & (0xFFFF >> (0x10 - drops[k])),
                (k == steps - 1) ? STONE_AT(location) : STONE_FLAT);
  }
  // Then we drop them from the source
  drop_stones(location, total);

  return ACT_OK;
}

// -------------------------------------------------------------------
// Win conditions

static uint8_t
board_full(void) {
  for (uint8_t k = 0; k < NUM_SQUARES; k++) {
    if (COUNT_AT(k) == 0) return 0;
  }
  return 1;
}

// Depth-first search of the board for a road with a given
// directionality:
// direction = 0 --> left-to-right, direction = 1 --> bottom-to-top
static uint8_t
dfs_road(uint8_t dfs_stack[NUM_SQUARES], uint8_t dfs_pntr,
         const enum COLOUR colour, const uint8_t direction) {
  while (dfs_pntr > 0) {
    const uint8_t cur = dfs_stack[--dfs_pntr];
    // Made it to the other side?
    if ( (direction == 1 && cur >= board_size * (board_size - 1))
         || (direction == 0 && cur % board_size == 1) )
        return 1;

    // Check the four neighbours of this cell, provided they exist,
    // are inhabited, and are of the appropriate colour

    // Direction: > (same row)
    if ( ((cur % board_size) + 1 < board_size)
         && (COUNT_AT(cur + 1)) // wont ever be out of bounds
         && ((colours[cur+1] & 1) == colour)
         && ((celldat[cur+1] & DFS_MASK) == 0) ) {
      dfs_stack[dfs_pntr++] = cur + 1;
      celldat[cur+1] |= DFS_MASK;
    }

    // Direction: <  (same row)
    if ( (cur % board_size > 0)
         && (COUNT_AT(cur - 1)) // wont ever be out of bounds
         && ((colours[cur-1] & 1) == colour)
         && ((celldat[cur-1] & DFS_MASK) == 0) ) {
      dfs_stack[dfs_pntr++] = cur - 1;
      celldat[cur-1] |= DFS_MASK;
    }

    // Direction: +
    if ( (cur + board_size < NUM_SQUARES)
         && (COUNT_AT(cur + board_size))
         && ((colours[cur+board_size] & 1) == colour)
         && ((celldat[cur+board_size] & DFS_MASK) == 0) ) {
      dfs_stack[dfs_pntr++] = cur + board_size;
      celldat[cur+board_size] |= DFS_MASK;
    }

    // Direction: -
    if ( (cur >= board_size)
         && (COUNT_AT(cur - board_size))
         && ((colours[cur-board_size] & 1) == colour)
         && ((celldat[cur-board_size] & DFS_MASK) == 0) ) {
      dfs_stack[dfs_pntr++] = cur - board_size;
      celldat[cur-board_size] |= DFS_MASK;
    }
  }
  return 0;
}

static enum WIN_TYPE
check_road_colour(const enum COLOUR colour) {
  uint8_t dfs_stack[NUM_SQUARES];
  uint8_t dfs_pntr;

  // Prime the depth-first-search stack with all boundary cells of
  // colour COLOUR.

  const enum WIN_TYPE winner =
    (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;

  // We're using the two left-over bits in data_t to track whether
  // we've seen it. Reset those before anything.

  for (uint8_t k=0; k<NUM_SQUARES; k++)
    celldat[k] &= NOT_DFS_MASK;

  // First left-to-right
  dfs_pntr = 0;
  for (uint8_t y=0; y<board_size; y++) {
    if ( COUNT_AT(THE_COORDS(0, y)) &&
        (colours[THE_COORDS(0, y)] & 1) == colour) {
      dfs_stack[dfs_pntr++] = THE_COORDS(0, y);
      celldat[THE_COORDS(0, y)] |= DFS_MASK;
    }
  }
  if (dfs_road(dfs_stack, dfs_pntr, colour, 0)) return winner;

  // Clear visited squares
  for (uint8_t k=0; k<NUM_SQUARES; k++)
    celldat[k] &= NOT_DFS_MASK;

  // Then bottom-to-top
  dfs_pntr = 0;
  for (uint8_t x=0; x<board_size; x++) {
    if ( COUNT_AT(THE_COORDS(x, 0)) &&
         (colours[THE_COORDS(x, 0)] & 1) == colour) {
      dfs_stack[dfs_pntr++] = THE_COORDS(x, 0);
      celldat[THE_COORDS(x, 0)] |= DFS_MASK;
    }
  }
  if (dfs_road(dfs_stack, dfs_pntr, colour, 1)) return winner;

  return -1;
}

enum WIN_TYPE
check_win(void) {
  // Do we do a flat count?
  if ((black_count & 127) == 0 || (white_count & 127) == 0 || board_full()) {
      int8_t total = 0;
      for (uint8_t k = 0; k < NUM_SQUARES; k++) {
        if (STONE_AT(k) == STONE_FLAT) {
          total += ((colours[k] & 1) == C_BLACK) ? +1 : -1 ;
        }
      }
      if (total > 0) {
        return WIN_FLAT_BLACK;
      } else if (total < 0) {
        return WIN_FLAT_WHITE;
      } else {
        return WIN_DRAW;
      }
  }

  // Road?

  enum WIN_TYPE rb, rw;
  rb = check_road_colour(C_BLACK);
  rw = check_road_colour(C_WHITE);

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

// -------------------------------------------------------------------
// PTN place parser

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

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

  if (board_size < 5 || board_size > 6) return PTN_INVALID;

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

// -------------------------------------------------------------------
// PTN move parser

enum E_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]) {

  if (board_size < 5 || board_size > 6) return PTN_INVALID;

  ASSERT_NONEMPTY;

  uint8_t picked_up = 1;

  if ( (*ptn >= '1') && (*ptn <= '0' + board_size)) {
    picked_up = *ptn - '0';
    ptn++; 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');

  ptn++; ASSERT_MORE;

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

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

  *out_steps = 0;
  uint8_t total = 0;
  while (*ptn) {
      if ( (*ptn < '1') || (*ptn > '0' + board_size) ) {
        return PTN_INVALID;
      }

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

  if ( total != picked_up ) return PTN_INVALID;

  return PTN_VALID;
}

// -------------------------------------------------------------------
// 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 E_RESULT
do_ptn(char *ptn) {

  if (won < -1) return GAME_END;

  enum E_RESULT res;
  uint8_t location;

  if (is_not_placement(ptn)) {
    uint8_t steps, drops[5];
    enum MOVE_DIRECTION direction;
    res = parse_move(board_size, ptn, &location,
                     &direction, &steps, drops);
    if (res == PTN_VALID) {
      if (ply < 2) return ACT_ILLEGAL;
      res = try_move(location, direction, steps, drops);
    }
  } else {
    enum STONE_VARIANT stone;
    res = parse_place(board_size, ptn, &location, &stone);
    if (res == PTN_VALID)
      res = try_place(location, current_colour, stone);
  }

  if (res == ACT_OK) {
    next_ply();
    // Could be less conservative here :)
    if (ply > board_size) {
      won = check_win();
      if (won < -1) return GAME_END;
    }
  }

   return res;
}