game.py

import copy
from board import create_board
from moves import pawn_moves, rook_moves, bishop_moves, knight_moves, queen_moves, king_moves
from pieces import Piece


# game class
class Game:
    def __init__(self):

        self.board = create_board()
        self.turn = 'w'
        self.move_count = 0
        self.history = []
        self.enpassant = None
        self.state = None
        self.castling = {'wKR': True, 'wQR': True, 'bKR': True, 'bQR': True}

        self.w_king_pos = (7, 4)
        self.b_king_pos = (0, 4)

        self.move_clock = 0
        self.position_count = {}
        self.position_count[self.board_key()] = 1

    def pseudo_moves(self, row, col):
        piece = self.board[row][col]
        if not piece: return []
        if piece.name == 'P':
            return pawn_moves(self.board, row, col, self.enpassant)
        elif piece.name == 'N':
            return knight_moves(self.board, row, col)
        elif piece.name == 'B':
            return bishop_moves(self.board, row, col)
        elif piece.name == 'R':
            return rook_moves(self.board, row, col)
        elif piece.name == 'Q':
            return queen_moves(self.board, row, col)
        elif piece.name == 'K':
            moves = king_moves(self.board, row, col)
            moves.extend(self.castling_moves(piece, row, col))
            return moves
        return []

    def castling_moves(self, king, row, col):
        moves = []
        if king.colour == 'w':
            if row != 7 or col != 4: return []
            if self.castling['wKR'] and self.board[7][5] is None and self.board[7][6] is None:
                if not self.square_attacked((7, 4), 'b') and not self.square_attacked((7, 5),
                                                                                      'b') and not self.square_attacked(
                    (7, 6), 'b'):
                    moves.append((7, 6))
            if self.castling['wQR'] and self.board[7][3] is None and self.board[7][2] is None and self.board[7][
                1] is None:
                if not self.square_attacked((7, 4), 'b') and not self.square_attacked((7, 3),
                                                                                      'b') and not self.square_attacked(
                    (7, 2), 'b'):
                    moves.append((7, 2))
        elif king.colour == 'b':
            if row != 0 or col != 4: return []
            if self.castling['bKR'] and self.board[0][5] is None and self.board[0][6] is None:
                if not self.square_attacked((0, 4), 'w') and not self.square_attacked((0, 5),
                                                                                      'w') and not self.square_attacked(
                    (0, 6), 'w'):
                    moves.append((0, 6))
            if self.castling['bQR'] and self.board[0][3] is None and self.board[0][2] is None and self.board[0][
                1] is None:
                if not self.square_attacked((0, 4), 'w') and not self.square_attacked((0, 3),
                                                                                      'w') and not self.square_attacked(
                    (0, 2), 'w'):
                    moves.append((0, 2))
        return moves

    def attack_moves(self, row, col):
        piece = self.board[row][col]
        if not piece: return []
        if piece.name == 'P':
            direction = -1 if piece.colour == 'w' else 1
            moves = []
            nrow = row + direction
            if 0 <= nrow < 8:
                if col - 1 >= 0: moves.append((nrow, col - 1))
                if col + 1 < 8: moves.append((nrow, col + 1))
            return moves
        elif piece.name == 'N':
            return knight_moves(self.board, row, col)
        elif piece.name == 'B':
            return bishop_moves(self.board, row, col)
        elif piece.name == 'R':
            return rook_moves(self.board, row, col)
        elif piece.name == 'Q':
            return queen_moves(self.board, row, col)
        elif piece.name == 'K':
            return king_moves(self.board, row, col)
        return []

    def get_moves(self, row, col):
        piece = self.board[row][col]
        if not piece or piece.colour != self.turn: return []

        moves = self.pseudo_moves(row, col)
        legal_moves = []

        king_pos_to_check = self.w_king_pos if piece.colour == 'w' else self.b_king_pos

        for r, c in moves:
            move_details = self.force_move_and_save((row, col), (r, c))

            if piece.name == 'K':
                king_pos_to_check = (r, c)

            if not self.square_attacked(king_pos_to_check, 'b' if piece.colour == 'w' else 'w'):
                legal_moves.append((r, c))

            self.unmake_move((row, col), (r, c), move_details)

            if piece.name == 'K':
                king_pos_to_check = (row, col)

        return legal_moves

    def force_move_and_save(self, start, end, promotion='Q'):
        r1, c1 = start;
        r2, c2 = end
        piece = self.board[r1][c1]

        old_castling = self.castling.copy()
        old_enpassant = self.enpassant
        captured = self.board[r2][c2]
        captured_pos = None

        old_w_king = self.w_king_pos
        old_b_king = self.b_king_pos

        # --- FIX 1 (Castling): Check captured piece ---
        if captured and captured.name == 'R':
            if captured.colour == 'w':
                if (r2, c2) == (7, 0):
                    self.castling['wQR'] = False
                elif (r2, c2) == (7, 7):
                    self.castling['wKR'] = False
            else:  # 'b'
                if (r2, c2) == (0, 0):
                    self.castling['bQR'] = False
                elif (r2, c2) == (0, 7):
                    self.castling['bKR'] = False
        # --- END OF FIX ---

        if piece.name == 'P' and self.enpassant and end == self.enpassant:
            captured_pos = (r2 + 1, c2) if piece.colour == 'w' else (r2 - 1, c2)
            captured = self.board[captured_pos[0]][captured_pos[1]]
            self.board[captured_pos[0]][captured_pos[1]] = None

        self.board[r2][c2] = piece
        self.board[r1][c1] = None

        if piece.name == 'K':
            if piece.colour == 'w':
                self.w_king_pos = (r2, c2)
            else:
                self.b_king_pos = (r2, c2)

        if piece.name == 'K' and abs(c2 - c1) == 2:
            if c2 == 6:  # Kingside
                rook = self.board[r2][7];
                self.board[r2][5] = rook;
                self.board[r2][7] = None
                rook_move = ((r2, 7), (r2, 5), rook)
            elif c2 == 2:  # Queenside
                rook = self.board[r2][0];
                self.board[r2][3] = rook;
                self.board[r2][0] = None
                rook_move = ((r2, 0), (r2, 3), rook)
            else:
                rook_move = None
        else:
            rook_move = None

        promoted_to = None
        if piece.name == 'P':
            if (piece.colour == 'w' and r2 == 0) or (piece.colour == 'b' and r2 == 7):
                promo_piece_name = promotion if promotion is not None else 'Q'
                self.board[r2][c2] = Piece(piece.colour, promo_piece_name)
                promoted_to = self.board[r2][c2]

        if piece.name == 'K':
            if piece.colour == 'w':
                self.castling['wKR'] = False;
                self.castling['wQR'] = False
            else:
                self.castling['bKR'] = False;
                self.castling['bQR'] = False
        elif piece.name == 'R':
            if piece.colour == 'w':
                if (r1, c1) == (7, 0):
                    self.castling['wQR'] = False
                elif (r1, c1) == (7, 7):
                    self.castling['wKR'] = False
            else:
                if (r1, c1) == (0, 0):
                    self.castling['bQR'] = False
                elif (r1, c1) == (0, 7):
                    self.castling['bKR'] = False

        # --- UPDATED EN PASSANT LOGIC ---
        new_enpassant = None
        if piece.name == 'P' and abs(r2 - r1) == 2:
            # This is the potential target square (e.g., d3)
            ep_target_square = ((r1 + r2) // 2, c1)

            # Check for enemy pawns on adjacent files *at the destination rank* (e.g., c4 or e4)
            for dc in [-1, 1]:
                check_c = c1 + dc
                if 0 <= check_c < 8:
                    # r2 is the destination rank (e.g., rank 4)
                    enemy_pawn = self.board[r2][check_c]
                    if (enemy_pawn and
                            enemy_pawn.name == 'P' and
                            enemy_pawn.colour != piece.colour):
                        # Found a pawn that can capture, so the EP square is valid
                        new_enpassant = ep_target_square
                        break  # No need to check the other side
        self.enpassant = new_enpassant


        return (piece, captured, captured_pos, rook_move, promoted_to, old_enpassant, old_castling, old_w_king,
                old_b_king)

    def unmake_move(self, start, end, move_details):
        (piece, captured, captured_pos, rook_move, promoted_to, old_enpassant, old_castling, old_w_king,
         old_b_king) = move_details
        r1, c1 = start;
        r2, c2 = end

        self.board[r1][c1] = piece
        self.board[r2][c2] = None

        if captured:
            restore_pos = captured_pos if captured_pos else end
            self.board[restore_pos[0]][restore_pos[1]] = captured

        if rook_move:
            r_start, r_end, rook = rook_move
            self.board[r_start[0]][r_start[1]] = rook
            self.board[r_end[0]][r_end[1]] = None

        self.enpassant = old_enpassant
        self.castling = old_castling

        self.w_king_pos = old_w_king
        self.b_king_pos = old_b_king

    def _force_move(self, start, end, promotion='Q'):
        r1, c1 = start;
        r2, c2 = end;
        piece = self.board[r1][c1]

        if not piece: return False

        old_castling = self.castling.copy()
        old_enpassant = self.enpassant
        old_move_clock = self.move_clock
        old_turn = self.turn

        old_w_king = self.w_king_pos
        old_b_king = self.b_king_pos

        captured = self.board[r2][c2];
        captured_pos = None

        if captured and captured.name == 'R':
            if captured.colour == 'w':
                if (r2, c2) == (7, 0):
                    self.castling['wQR'] = False
                elif (r2, c2) == (7, 7):
                    self.castling['wKR'] = False
            else:  # 'b'
                if (r2, c2) == (0, 0):
                    self.castling['bQR'] = False
                elif (r2, c2) == (0, 7):
                    self.castling['bKR'] = False

        if piece.name == 'P' and self.enpassant and end == self.enpassant:
            captured_pos = (r2 + 1, c2) if piece.colour == 'w' else (r2 - 1, c2)
            captured = self.board[captured_pos[0]][captured_pos[1]]
            self.board[captured_pos[0]][captured_pos[1]] = None

        self.board[r2][c2] = piece;
        self.board[r1][c1] = None

        if piece.name == 'K':
            if piece.colour == 'w':
                self.w_king_pos = (r2, c2)
            else:
                self.b_king_pos = (r2, c2)

        rook_move = None
        if piece.name == 'K' and abs(c2 - c1) == 2:
            if c2 == 6:
                rook = self.board[r2][7];
                self.board[r2][5] = rook;
                self.board[r2][7] = None
                rook_move = ((r2, 7), (r2, 5), rook)
            elif c2 == 2:
                rook = self.board[r2][0];
                self.board[r2][3] = rook;
                self.board[r2][0] = None
                rook_move = ((r2, 0), (r2, 3), rook)

        promoted_to = None
        if piece.name == 'P':
            if (piece.colour == 'w' and r2 == 0) or (piece.colour == 'b' and r2 == 7):
                self.board[r2][c2] = Piece(piece.colour, promotion)
                promoted_to = self.board[r2][c2]

        if piece.name == 'K':
            if piece.colour == 'w':
                self.castling['wKR'] = False;
                self.castling['wQR'] = False
            else:
                self.castling['bKR'] = False;
                self.castling['bQR'] = False
        elif piece.name == 'R':
            if piece.colour == 'w':
                if (r1, c1) == (7, 0):
                    self.castling['wQR'] = False
                elif (r1, c1) == (7, 7):
                    self.castling['wKR'] = False
            else:
                if (r1, c1) == (0, 0):
                    self.castling['bQR'] = False
                elif (r1, c1) == (0, 7):
                    self.castling['bKR'] = False

        new_enpassant = None
        if piece.name == 'P' and abs(r2 - r1) == 2:
            ep_target_square = ((r1 + r2) // 2, c1)

            # Check for enemy pawns on adjacent files *at the destination rank* (e.g., c4 or e4)
            for dc in [-1, 1]:
                check_c = c1 + dc
                if 0 <= check_c < 8:
                    # r2 is the destination rank (e.g., rank 4)
                    enemy_pawn = self.board[r2][check_c]
                    if (enemy_pawn and
                            enemy_pawn.name == 'P' and
                            enemy_pawn.colour != piece.colour):
                        # Found a pawn that can capture, so the EP square is valid
                        new_enpassant = ep_target_square
                        break  # No need to check the other side
        self.enpassant = new_enpassant

        return (piece, captured, captured_pos, rook_move, promoted_to, old_enpassant, old_castling, old_move_clock,
                old_turn, old_w_king, old_b_king)

    def make_move(self, start, end, promotion='Q'):
        if start is None or end is None: return False
        r1, c1 = start;
        piece = self.board[r1][c1]
        if not piece or piece.colour != self.turn: return False

        moves = self.get_moves(r1, c1)
        if end not in moves: return False

        move_details = self._force_move(start, end, promotion)
        (moved_piece, captured_piece, captured_pos, rook_move, promoted_to, old_enpassant, old_castling, old_move_clock,
         old_turn, old_w_king, old_b_king) = move_details

        history_entry = (start, end, moved_piece, captured_piece, captured_pos, rook_move, promoted_to,
                         old_enpassant, old_castling, old_move_clock, old_turn, old_w_king, old_b_king)
        self.history.append(history_entry)

        self.turn = 'b' if self.turn == 'w' else 'w'
        if self.turn == 'w': self.move_count += 1

        if moved_piece.name == 'P' or captured_piece:
            self.move_clock = 0
        else:
            self.move_clock += 1

        key = self.board_key()
        self.position_count[key] = self.position_count.get(key, 0) + 1

        if self.is_checkmate(self.turn):
            self.state = "Checkmate"
        elif self.is_stalemate(self.turn):
            self.state = "Stalemate"
        elif self.move_clock >= 100:
            self.state = "Draw (50-move rule)"
        elif self.position_count.get(key, 0) >= 3:
            self.state = "Draw (Threefold repetition)"
        elif self.is_check(self.turn):
            self.state = "Check"
        else:
            self.state = None
        return True

    def full_unmake_move(self):
        if not self.history: return False
        (start, end, piece, captured, captured_pos, rook_move, promoted_to,
         old_enpassant, old_castling, old_move_clock, old_turn,
         old_w_king, old_b_king) = self.history.pop()

        r1, c1 = start;
        r2, c2 = end
        self.board[r1][c1] = piece
        if promoted_to:
            self.board[r2][c2] = None
        else:
            self.board[r2][c2] = None

        if captured:
            restore_pos = captured_pos if captured_pos else end
            self.board[restore_pos[0]][restore_pos[1]] = captured

        if rook_move:
            r_start, r_end, rook = rook_move
            self.board[r_start[0]][r_start[1]] = rook
            self.board[r_end[0]][r_end[1]] = None

        self.enpassant = old_enpassant
        self.castling = old_castling
        self.move_clock = old_move_clock
        self.turn = old_turn
        if self.turn == 'b': self.move_count -= 1

        self.w_king_pos = old_w_king
        self.b_king_pos = old_b_king

        # FIX: Ensure key exists before attempting to decrement/delete
        key = self.board_key()
        if key in self.position_count:
            self.position_count[key] -= 1
            if self.position_count[key] == 0:
                del self.position_count[key]
        # END FIX

        self.state = None
        if self.is_check(self.turn): self.state = "Check"
        return True

    def square_attacked(self, square, colour):
        for row in range(8):
            for col in range(8):
                piece = self.board[row][col]
                if piece and piece.colour == colour:
                    if square in self.attack_moves(row, col): return True
        return False

    def is_check(self, colour):
        king_pos = self.w_king_pos if colour == 'w' else self.b_king_pos
        return self.square_attacked(king_pos, 'b' if colour == 'w' else 'w')

    def is_checkmate(self, colour):
        if not self.is_check(colour): return False
        for r in range(8):
            for c in range(8):
                piece = self.board[r][c]
                if piece and piece.colour == colour:
                    if self.get_moves(r, c): return False
        return True

    def is_stalemate(self, colour):
        if self.is_check(colour): return False
        for r in range(8):
            for c in range(8):
                piece = self.board[r][c]
                if piece and piece.colour == colour:
                    if self.get_moves(r, c): return False
        return True

    def copy(self):
        return copy.deepcopy(self)

    def light_copy(self):
        copy = Game()
        copy.board = [[piece for piece in row] for row in self.board]
        copy.turn = self.turn
        copy.enpassant = self.enpassant
        copy.castling = self.castling.copy()
        copy.move_count = self.move_count
        copy.move_clock = self.move_clock
        copy.state = self.state
        copy.w_king_pos = self.w_king_pos
        copy.b_king_pos = self.b_king_pos
        copy.position_count = self.position_count.copy()

        return copy

    def board_key(self):
        rows = []
        for r in range(8):
            row = []
            for c in range(8):
                p = self.board[r][c]
                # FIX: Ensure proper string for empty squares
                row.append(str(p) if p else ".")
            rows.append("".join(row))
        return "/".join(rows) + " " + self.turn

    def perft(self, depth):
        """
        Recursively calculates the number of legal leaf nodes at a given depth.
        Used for testing the integrity of move generation and board state updates.
        """
        # Base Case: When depth reaches 0, we count the node
        if depth == 0:
            return 1

        nodes = 0
        current_colour = self.turn

        # 1. Gather all legal moves for the current position
        all_moves = []
        for r_start in range(8):
            for c_start in range(8):
                piece = self.board[r_start][c_start]
                if piece and piece.colour == current_colour:
                    try:
                        # get_moves returns the list of legal end squares
                        legal_ends = self.get_moves(r_start, c_start)

                        # Handle promotion options explicitly for each end square
                        for r_end, c_end in legal_ends:
                            if piece.name == 'P' and (r_end == 0 or r_end == 7):
                                # If it's a promotion, list all promotion types
                                for promo in ['Q', 'R', 'B', 'N']:
                                    all_moves.append(((r_start, c_start), (r_end, c_end), promo))
                            else:
                                # Not a promotion
                                all_moves.append(((r_start, c_start), (r_end, c_end), None))
                    except Exception:
                        # Skip if any piece/move function raises an unexpected error
                        continue

        # 2. Iterate and recurse through moves
        for start, end, promotion in all_moves:

            # Since get_moves already ensures the move is legal, make_move should return True
            if self.make_move(start, end, promotion):
                # Recursive call
                nodes += self.perft(depth - 1)

                # Unmake the move to restore the board for the next iteration
                self.full_unmake_move()

        return nodes
    # --- END OF PERFT FUNCTION ---


def notation_to_index(move):
    col_map = {'a': 0, 'b': 1, 'c': 2, 'd': 3,
               'e': 4, 'f': 5, 'g': 6, 'h': 7}
    col = col_map.get(move[0].lower())
    if col is None: return None
    try:
        row = 8 - int(move[1])
    except ValueError:
        return None
    return (row, col)


def index_to_notation(pos):
    col_map = 'abcdefgh'
    row, col = pos
    return col_map[col] + str(8 - row)


def move_to_algebraic(game, start, end, promotion=None):
    piece = game.board[start[0]][start[1]]
    target_piece = game.board[end[0]][end[1]]

    if piece.name == 'K' and abs(start[1] - end[1]) == 2:
        return 'O-O' if end[1] == 6 else 'O-O-O'

    move_str = ''
    if piece.name != 'P': move_str += piece.name

    if target_piece or (piece.name == 'P' and start[1] != end[1]):
        if piece.name == 'P': move_str += index_to_notation(start)[0]
        move_str += 'x'

    move_str += index_to_notation(end)

    if promotion: move_str += f"={promotion}"

    return move_str