#include "../search.h"

bitboard_t get_white_pawn_forwardmoves(moves_t* moves, position_t position) {
  bitboard_t enemy_bitboard = blacks(position);
  bitboard_t friendly_pieces =
    position.bitboards[WHITE_KING] |
    position.bitboards[WHITE_QUEEN] |
    position.bitboards[WHITE_ROOK] |
    position.bitboards[WHITE_BISHOP] |
    position.bitboards[WHITE_KNIGHT] |
    position.bitboards[WHITE_PAWN];
  bitboard_t friendly_pawns = position.bitboards[WHITE_PAWN];
  bitboard_t y = friendly_pieces | enemy_bitboard;
  bitboard_t x = ~((friendly_pawns << 8) & y) & (friendly_pawns << 8);
  bitboard_t first_move_x = x & BITMASK_RANK_3;

  first_move_x |= (first_move_x << 8) & (~y);
  x |= first_move_x;

  return x;
}
bitboard_t get_black_pawn_forwardmoves(moves_t* moves, position_t position) {
  bitboard_t enemy_bitboard = whites(position);
  bitboard_t friendly_pieces =
    position.bitboards[BLACK_KING] |
    position.bitboards[BLACK_QUEEN] |
    position.bitboards[BLACK_ROOK] |
    position.bitboards[BLACK_BISHOP] |
    position.bitboards[BLACK_KNIGHT] |
    position.bitboards[BLACK_PAWN];
  bitboard_t friendly_pawns = position.bitboards[BLACK_PAWN];
  bitboard_t y = friendly_pieces | enemy_bitboard;
  bitboard_t x = ~((friendly_pawns >> 8) & y) & (friendly_pawns >> 8);
  bitboard_t first_move_x = x & BITMASK_RANK_6;

  first_move_x |= (first_move_x >> 8) & (~y);
  x |= first_move_x;

  return x;
}

void get_pawn_moves(moves_t* moves, position_t position) {
  assert_valid_position(position);

  bitboard_t forward_moves;
  bitboard_t friendly_pawns;
  bitboard_t enemy_bitboard;
  if (position.turn == WHITE_TURN) {
    forward_moves = get_white_pawn_forwardmoves(moves, position);
    friendly_pawns = position.bitboards[WHITE_PAWN];
    enemy_bitboard = blacks(position);
  } else {
    forward_moves = get_black_pawn_forwardmoves(moves, position);
    friendly_pawns = position.bitboards[BLACK_PAWN];
    enemy_bitboard = whites(position);
  }

  bitboard_t file_masks[8] = {
    BITMASK_FILE_A,
    BITMASK_FILE_B,
    BITMASK_FILE_C,
    BITMASK_FILE_D,
    BITMASK_FILE_E,
    BITMASK_FILE_F,
    BITMASK_FILE_G,
    BITMASK_FILE_H,
  };

  for (int i = 0; i < 8; i++) {
    bitboard_t item = file_masks[i] & forward_moves;
    bitboard_t pawns = file_masks[i] & friendly_pawns;

    int from = __builtin_ctzll(pawns);
    pawns &= (pawns - 1);
    int next_from = __builtin_ctzll(pawns);
    if (pawns == 0) next_from = 64; // to work around UB
    
    while (item) {
      int to = __builtin_ctzll(item);
      if (to > next_from) {
        from = next_from;
        pawns &= (pawns - 1);
        next_from = __builtin_ctzll(pawns);
        if (pawns == 0) next_from = 64;
      }
      item &= (item - 1);
      add_move(moves, from, to);
    }
  }

  if (position.passantable_column != 0) {
    bitboard_t PASSANT_RANK;
    int rank_shift;
    int from_shift;
    if (position.turn == WHITE_TURN) {
      PASSANT_RANK = BITMASK_RANK_5;
      rank_shift = 40;
      from_shift = -9;
    } else {
      PASSANT_RANK = BITMASK_RANK_4;
      from_shift = 7;
      rank_shift = 16;
    }

    int left = 0, right = 0;
    if (position.passantable_column > 1) left = position.passantable_column - 1;
    if (position.passantable_column < 8) right = position.passantable_column + 1;
    if (left != 0) {
      bitboard_t MASK = file_masks[left - 1] & friendly_pawns & PASSANT_RANK;
      if (MASK != 0) {
        add_move(
          moves,
          left + rank_shift + from_shift,
          position.passantable_column + rank_shift - 1
        );
      }
    }
    if (right != 0) {
      bitboard_t MASK = file_masks[right - 1] & friendly_pawns & PASSANT_RANK;
      if (MASK != 0) {
        add_move(
          moves,
          right + rank_shift + from_shift,
          position.passantable_column + rank_shift - 1
        );
      }
    }
  }


  bitboard_t ATTACK_MASK;
  int shift;
  if (position.turn == WHITE_TURN) {
    shift = 0;
  } else {
    shift = -16;
  }
  for (int i = 0; i < 8; i++) {
    if (i == 0) ATTACK_MASK = 512;
    else if (i == 7) ATTACK_MASK = 128;
    else {
      ATTACK_MASK = 640;
    }
    bitboard_t pawns = file_masks[i] & friendly_pawns;
    while (pawns) {
      int from = __builtin_ctzll(pawns);
      pawns &= pawns - 1;

      bitboard_t MASK; 
      if (from + shift > 0) {
        MASK = ATTACK_MASK << (from + shift);
      } else {
        MASK = ATTACK_MASK >> -(from + shift);
      }

      MASK &= enemy_bitboard;
      while (MASK) {
        int to = __builtin_ctzll(MASK);
        MASK &= MASK - 1;
        add_move(moves, from, to);
      }
    }
  }
}