#include "engine/moves.h"
#include "bitboard.h"

#include "moves/vec.c"
#include "moves/magic.c"
#include "moves/knight.c"
#include "moves/pawn.c"
#include "moves/rook.c"
#include "moves/bishop.c"
#include "moves/queen.c"
#include "moves/attack.c"
#ifndef TEST_MOD
#include "moves/king.c"
#endif

#include <assert.h>

void get_moves(moves_t* moves, position_t position) {
  get_pawn_moves(moves, position);
  get_knight_moves(moves, position);
  get_king_moves(moves, position);
  get_rook_moves(moves, position);
  get_bishop_moves(moves, position);
  get_queen_moves(moves, position);
}

void get_legal_moves(moves_t* moves, position_t position) {
  get_moves(moves, position);

  u8 opponent = position.turn == WHITE_TURN ? BLACK_TURN : WHITE_TURN;
  int write = 0;
  for (int read = 0; read < moves->length; read++) {
    move_t m = moves->moves[read];

    if (m.flags & MOVE_SHORT_CASTLE) {
      int king_sq = position.turn == WHITE_TURN ? 4 : 60;
      int pass_sq = position.turn == WHITE_TURN ? 5 : 61;
      if (square_attacked(position, king_sq, opponent)) continue;
      if (square_attacked(position, pass_sq, opponent)) continue;
    }
    if (m.flags & MOVE_LONG_CASTLE) {
      int king_sq = position.turn == WHITE_TURN ? 4 : 60;
      int pass_sq = position.turn == WHITE_TURN ? 3 : 59;
      if (square_attacked(position, king_sq, opponent)) continue;
      if (square_attacked(position, pass_sq, opponent)) continue;
    }

    position_t copy = position;
    position_make_move(&copy, m);

    bitboard_t friendly_king;
    if (position.turn == WHITE_TURN) {
      friendly_king = copy.bitboards[WHITE_KING];
    } else {
      friendly_king = copy.bitboards[BLACK_KING];
    }
    int king_square = __builtin_ctzll(friendly_king);

    if (!square_attacked(copy, king_square, opponent)) {
      moves->moves[write++] = m;
    }
  }
  moves->length = write;
}

void position_make_move(position_t* position, move_t move) {
  position->passantable_file = 0;

  if (move.flags & MOVE_LONG_CASTLE) {
    if (position->turn == WHITE_TURN) {
      position->bitboards[WHITE_KING] = (bitboard_t)1 << 2;
      position->bitboards[WHITE_ROOK] &= ~((bitboard_t)1 << 0);
      position->bitboards[WHITE_ROOK] |= (bitboard_t)1 << 3;
    } else {
      position->bitboards[BLACK_KING] = (bitboard_t)1 << 58;
      position->bitboards[BLACK_ROOK] &= ~((bitboard_t)1 << 56);
      position->bitboards[BLACK_ROOK] |= (bitboard_t)1 << 59;
    }
    position->turn = !position->turn;
    return;
  }
  if (move.flags & MOVE_SHORT_CASTLE) {
    if (position->turn == WHITE_TURN) {
      position->bitboards[WHITE_KING] = (bitboard_t)1 << 6;
      position->bitboards[WHITE_ROOK] &= ~((bitboard_t)1 << 7);
      position->bitboards[WHITE_ROOK] |= (bitboard_t)1 << 5;
    } else {
      position->bitboards[BLACK_KING] = (bitboard_t)1 << 62;
      position->bitboards[BLACK_ROOK] &= ~((bitboard_t)1 << 63);
      position->bitboards[BLACK_ROOK] |= (bitboard_t)1 << 61;
    }
    position->turn = !position->turn;
    return;
  }

  int piece_type = find_piece_on_square(position, move.from);
  position->bitboards[piece_type] &= ~((bitboard_t)1 << move.from);

  if (position->turn == BLACK_TURN) position->fullmove_clock++;
  if (piece_type == WHITE_PAWN || piece_type == BLACK_PAWN) {
    position->halfmove_clock = 0;
  } else {
    position->halfmove_clock++;
  }

  if (position->castling > 0) {
    if (piece_type == WHITE_ROOK) {
      if (move.from == 7) {
        position->castling &= ~WHITE_SHORT_CASTLE;
      } else if (move.from == 0) {
        position->castling &= ~WHITE_LONG_CASTLE;
      }
    } else if (piece_type == BLACK_ROOK) {
      if (move.from == 63) {
        position->castling &= ~BLACK_SHORT_CASTLE;
      } else if (move.from == 56) {
        position->castling &= ~BLACK_LONG_CASTLE;
      }
    } else if (piece_type == WHITE_KING) {
      position->castling &= ~(WHITE_SHORT_CASTLE | WHITE_LONG_CASTLE);
    } else if (piece_type == BLACK_KING) {
      position->castling &= ~(BLACK_SHORT_CASTLE | BLACK_LONG_CASTLE);
    }
  }

  if (move.flags & MOVE_PROMOTE_Q) {
    if (move.flags & MOVE_CAPTURE) {
      int to_remove_piece_type = find_piece_on_square(position, move.to);
      position->bitboards[to_remove_piece_type] &= ~((bitboard_t)1 << move.to);
    }
    int q_type = position->turn == WHITE_TURN ? WHITE_QUEEN : BLACK_QUEEN;
    position->bitboards[q_type] |= (bitboard_t)1 << move.to;
    position->turn = !position->turn;
    return;
  }
  if (move.flags & MOVE_PROMOTE_R) {
    if (move.flags & MOVE_CAPTURE) {
      int to_remove_piece_type = find_piece_on_square(position, move.to);
      position->bitboards[to_remove_piece_type] &= ~((bitboard_t)1 << move.to);
    }
    int q_type = position->turn == WHITE_TURN ? WHITE_ROOK : BLACK_ROOK;
    position->bitboards[q_type] |= (bitboard_t)1 << move.to;
    position->turn = !position->turn;
    return;
  }
  if (move.flags & MOVE_PROMOTE_B) {
    if (move.flags & MOVE_CAPTURE) {
      int to_remove_piece_type = find_piece_on_square(position, move.to);
      position->bitboards[to_remove_piece_type] &= ~((bitboard_t)1 << move.to);
    }
    int q_type = position->turn == WHITE_TURN ? WHITE_BISHOP : BLACK_BISHOP;
    position->bitboards[q_type] |= (bitboard_t)1 << move.to;
    position->turn = !position->turn;
    return;
  }
  if (move.flags & MOVE_PROMOTE_N) {
    if (move.flags & MOVE_CAPTURE) {
      int to_remove_piece_type = find_piece_on_square(position, move.to);
      position->bitboards[to_remove_piece_type] &= ~((bitboard_t)1 << move.to);
    }
    int q_type = position->turn == WHITE_TURN ? WHITE_KNIGHT : BLACK_KNIGHT;
    position->bitboards[q_type] |= (bitboard_t)1 << move.to;
    position->turn = !position->turn;
    return;
  }
  if (move.flags & MOVE_EN_PASSANT) {
    int target_sqr;
    if (position->turn == WHITE_TURN) target_sqr = move.to - 8;
    else target_sqr = move.to + 8;
    
    int to_remove_piece_type = find_piece_on_square(position, target_sqr);
    position->bitboards[piece_type] |= (bitboard_t)1 << move.to;
    position->bitboards[to_remove_piece_type] &= ~((bitboard_t)1 << target_sqr);
    position->turn = !position->turn;
    return;
  }

  if (move.flags & MOVE_CAPTURE) {
    int to_remove_piece_type = find_piece_on_square(position, move.to);
    position->bitboards[to_remove_piece_type] &= ~((bitboard_t)1 << move.to);

    if (to_remove_piece_type == WHITE_ROOK) {
      if (move.to == 7) position->castling &= ~WHITE_SHORT_CASTLE;
      else if (move.to == 0) position->castling &= ~WHITE_LONG_CASTLE;
    } else if (to_remove_piece_type == BLACK_ROOK) {
      if (move.to == 63) position->castling &= ~BLACK_SHORT_CASTLE;
      else if (move.to == 56) position->castling &= ~BLACK_LONG_CASTLE;
    }
  }

  if (piece_type == WHITE_PAWN && move.from / 8 == 1 && move.to / 8 == 3) {
    position->passantable_file = (move.from % 8) + 1;
  } else if (piece_type == BLACK_PAWN && move.from / 8 == 6 && move.to / 8 == 4) {
    position->passantable_file = (move.from % 8) + 1;
  }

  position->bitboards[piece_type] |= (bitboard_t)1 << move.to;
  position->turn = !position->turn;
}

#ifdef TEST_MOD
#include "fen.h"

int count_positions(position_t position, int depth) {
  if (depth == 0) return 1;
  moves_t moves;
  moves_init(&moves);
  get_legal_moves(&moves, position);

  int count = 0;
  for (int i = 0; i < moves.length; i++) {
    position_t copy = position;
    position_make_move(&copy, moves.moves[i]);
    count += count_positions(copy, depth - 1);
  }
  moves_deinit(moves);
  return count;
}

void perft_divide(position_t position, int depth) {
  if (depth == 0) return;
  moves_t moves;
  moves_init(&moves);
  get_legal_moves(&moves, position);

  int nodes = 0;
  for (int i = 0; i < moves.length; i++) {
    position_t copy = position;
    position_make_move(&copy, moves.moves[i]);
    int c = count_positions(copy, depth - 1);
    nodes += c;
    printf("  %c%d%c%d: %d\n",
      'a' + (moves.moves[i].from % 8), 1 + (moves.moves[i].from / 8),
      'a' + (moves.moves[i].to % 8), 1 + (moves.moves[i].to / 8), c);
  }
  printf("\nNodes: %d\n", nodes);
  moves_deinit(moves);
}

// https://www.chessprogramming.org/Perft_Results#Initial_Position
bool test_perft_starting_position() {
  position_t position = load_fen(
    "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
  ).position;
  int d;
  d = count_positions(position, 0); if (d != 1) return false;
  d = count_positions(position, 1); if (d != 20) return false;
  d = count_positions(position, 2); if (d != 400) return false;
  d = count_positions(position, 3); if (d != 8902) return false;
  d = count_positions(position, 4); if (d != 197281) return false;

  position_t castle_pos = load_fen(
    "r3k2r/8/8/8/8/8/8/R3K2R w KQkq - 0 1"
  ).position;
  d = count_positions(castle_pos, 1); if (d != 26) return false;
  d = count_positions(position, 5); if (d != 4865609) return false;
  return true;
}

bool test_perft_kiwipete_position() {
  position_t position = load_fen(
    "r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq - 0 1"
  ).position;
  int d;
  d = count_positions(position, 1); if (d != 48) return false;
  d = count_positions(position, 2); if (d != 2039) return false;
  d = count_positions(position, 3); if (d != 97862) return false;
  d = count_positions(position, 4); if (d != 4085603) return false;

  // my move gen too slow to verify this
  // d = count_positions(position, 5); if (d != 193690690) return false;

  return true;
}

bool test_perft_position3() {
  position_t position = load_fen(
    "8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - - 0 1"
  ).position;
  int d;
  d = count_positions(position, 1); if (d != 14) return false;
  d = count_positions(position, 2); if (d != 191) return false;
  d = count_positions(position, 3); if (d != 2812) return false;
  d = count_positions(position, 4); if (d != 43238) return false;
  d = count_positions(position, 5); if (d != 674624) return false;
  d = count_positions(position, 6); if (d != 11030083) return false;

  // my move gen too slow to verify this
  // d = count_positions(position, 7); if (d != 178633661) return false;
  // d = count_positions(position, 8); if (d != 3009794393) return false;

  return true;
}

#endif // TEST_MOD