chess_engine/board.py

# chess_engine/board.py

import chess
import chess.engine
from typing import Optional, List, Tuple, Dict, Any
from enum import Enum
from dataclasses import dataclass
from .promotion import PromotionMoveHandler

class GameState(Enum):
    PLAYING = "playing"
    CHECK = "check"
    CHECKMATE = "checkmate"
    STALEMATE = "stalemate"
    DRAW = "draw"

class MoveResult(Enum):
    VALID = "valid"
    INVALID = "invalid"
    ILLEGAL = "illegal"

@dataclass
class MoveInfo:
    move: chess.Move
    san: str  # Standard Algebraic Notation
    uci: str  # Universal Chess Interface notation
    is_capture: bool
    is_check: bool
    is_checkmate: bool
    is_castling: bool
    promoted_piece: Optional[chess.PieceType] = None
    is_promotion: bool = False
    promotion_required: bool = False
    available_promotions: Optional[List[str]] = None

class ChessBoard:
    """
    Enhanced chess board class that wraps python-chess with additional functionality
    """
    
    def __init__(self, fen: Optional[str] = None):
        """
        Initialize chess board
        
        Args:
            fen: FEN string to initialize position, defaults to starting position
        """
        self.board = chess.Board(fen) if fen else chess.Board()
        self.move_history: List[MoveInfo] = []
        self.position_history: List[str] = [self.board.fen()]
        
    def get_board_state(self) -> Dict[str, Any]:
        """Get current board state as dictionary"""
        return {
            'fen': self.board.fen(),
            'turn': 'white' if self.board.turn == chess.WHITE else 'black',
            'game_state': self._get_game_state(),
            'castling_rights': {
                'white_kingside': self.board.has_kingside_castling_rights(chess.WHITE),
                'white_queenside': self.board.has_queenside_castling_rights(chess.WHITE),
                'black_kingside': self.board.has_kingside_castling_rights(chess.BLACK),
                'black_queenside': self.board.has_queenside_castling_rights(chess.BLACK),
            },
            'en_passant': self.board.ep_square,
            'halfmove_clock': self.board.halfmove_clock,
            'fullmove_number': self.board.fullmove_number,
            'legal_moves': [move.uci() for move in self.board.legal_moves],
            'in_check': self.board.is_check(),
            'move_count': len(self.move_history)
        }
    
    def get_piece_at(self, square: str) -> Optional[Dict[str, Any]]:
        """
        Get piece information at given square
        
        Args:
            square: Square in algebraic notation (e.g., 'e4')
            
        Returns:
            Dictionary with piece info or None if empty square
        """
        try:
            square_index = chess.parse_square(square)
            piece = self.board.piece_at(square_index)
            
            if piece is None:
                return None
                
            return {
                'type': piece.piece_type,
                'color': 'white' if piece.color == chess.WHITE else 'black',
                'symbol': piece.symbol(),
                'unicode': piece.unicode_symbol(),
                'square': square
            }
        except ValueError:
            return None
    
    def get_all_pieces(self) -> Dict[str, Dict[str, Any]]:
        """Get all pieces on the board"""
        pieces = {}
        for square in chess.SQUARES:
            square_name = chess.square_name(square)
            piece_info = self.get_piece_at(square_name)
            if piece_info:
                pieces[square_name] = piece_info
        return pieces
    
    def make_move(self, move_str: str) -> Tuple[MoveResult, Optional[MoveInfo]]:
        """
        Make a move on the board
        
        Args:
            move_str: Move in UCI notation (e.g., 'e2e4', 'e7e8q') or SAN (e.g., 'e4', 'e8=Q')
            
        Returns:
            Tuple of (MoveResult, MoveInfo if successful)
        """
        try:
            # Try to parse as UCI first, then SAN
            try:
                move = chess.Move.from_uci(move_str)
            except ValueError:
                move = self.board.parse_san(move_str)
            
            # Check if this is a promotion move that requires piece specification
            from_square = chess.square_name(move.from_square)
            to_square = chess.square_name(move.to_square)
            
            # If this is a promotion move but no promotion piece is specified, return promotion required
            if self.is_promotion_move(from_square, to_square) and move.promotion is None:
                # Get available promotion moves for this pawn
                available_promotions = self.get_promotion_moves(from_square)
                
                move_info = MoveInfo(
                    move=move,
                    san="",  # Will be empty since move is incomplete
                    uci=move.uci(),
                    is_capture=self.board.is_capture(move),
                    is_check=False,  # Cannot determine without promotion piece
                    is_checkmate=False,
                    is_castling=False,
                    promoted_piece=None,
                    is_promotion=True,
                    promotion_required=True,
                    available_promotions=available_promotions
                )
                
                return MoveResult.INVALID, move_info
            
            # Check if move is legal
            if move not in self.board.legal_moves:
                return MoveResult.ILLEGAL, None
            
            # Store move information before making the move
            is_promotion = move.promotion is not None
            move_info = MoveInfo(
                move=move,
                san=self.board.san(move),
                uci=move.uci(),
                is_capture=self.board.is_capture(move),
                is_check=self.board.gives_check(move),
                is_checkmate=False,  # Will be updated after move
                is_castling=self.board.is_castling(move),
                promoted_piece=move.promotion,
                is_promotion=is_promotion,
                promotion_required=False,
                available_promotions=None
            )
            
            # Make the move
            self.board.push(move)
            
            # Update move info with post-move state
            move_info.is_checkmate = self.board.is_checkmate()
            
            # Store in history
            self.move_history.append(move_info)
            self.position_history.append(self.board.fen())
            
            return MoveResult.VALID, move_info
            
        except ValueError:
            return MoveResult.INVALID, None
    
    def undo_move(self) -> bool:
        """
        Undo the last move
        
        Returns:
            True if successful, False if no moves to undo
        """
        if not self.move_history:
            return False
            
        self.board.pop()
        self.move_history.pop()
        self.position_history.pop()
        return True
    
    def get_legal_moves(self, square: Optional[str] = None) -> List[str]:
        """
        Get legal moves, optionally filtered by starting square
        
        Args:
            square: Starting square to filter moves (e.g., 'e2')
            
        Returns:
            List of legal moves in UCI notation
        """
        legal_moves = []
        
        for move in self.board.legal_moves:
            if square is None:
                legal_moves.append(move.uci())
            else:
                try:
                    square_index = chess.parse_square(square)
                    if move.from_square == square_index:
                        legal_moves.append(move.uci())
                except ValueError:
                    continue
                    
        return legal_moves
    
    def get_move_history(self) -> List[Dict[str, Any]]:
        """Get move history as list of dictionaries"""
        return [
            {
                'move_number': i + 1,
                'san': move.san,
                'uci': move.uci,
                'is_capture': move.is_capture,
                'is_check': move.is_check,
                'is_checkmate': move.is_checkmate,
                'is_castling': move.is_castling,
                'promoted_piece': move.promoted_piece,
                'is_promotion': move.is_promotion,
                'promotion_required': move.promotion_required,
                'available_promotions': move.available_promotions
            }
            for i, move in enumerate(self.move_history)
        ]
    
    def reset_board(self):
        """Reset board to starting position"""
        self.board = chess.Board()
        self.move_history.clear()
        self.position_history = [self.board.fen()]
    
    def load_position(self, fen: str) -> bool:
        """
        Load position from FEN string
        
        Args:
            fen: FEN string
            
        Returns:
            True if successful, False if invalid FEN
        """
        try:
            self.board = chess.Board(fen)
            self.move_history.clear()
            self.position_history = [fen]
            return True
        except ValueError:
            return False
    
    def _get_game_state(self) -> GameState:
        """Determine current game state"""
        if self.board.is_checkmate():
            return GameState.CHECKMATE
        elif self.board.is_stalemate():
            return GameState.STALEMATE
        elif self.board.is_insufficient_material() or \
             self.board.is_seventyfive_moves() or \
             self.board.is_fivefold_repetition():
            return GameState.DRAW
        elif self.board.is_check():
            return GameState.CHECK
        else:
            return GameState.PLAYING
    
    def get_board_array(self) -> List[List[Optional[str]]]:
        """
        Get board as 2D array for easier frontend rendering
        
        Returns:
            8x8 array where each cell contains piece symbol or None
        """
        board_array = []
        for rank in range(8):
            row = []
            for file in range(8):
                square = chess.square(file, 7-rank)  # Flip rank for display
                piece = self.board.piece_at(square)
                row.append(piece.symbol() if piece else None)
            board_array.append(row)
        return board_array
    
    def get_attacked_squares(self, color: chess.Color) -> List[str]:
        """Get squares attacked by given color"""
        attacked = []
        for square in chess.SQUARES:
            if self.board.is_attacked_by(color, square):
                attacked.append(chess.square_name(square))
        return attacked
    
    def is_square_attacked(self, square: str, by_color: str) -> bool:
        """Check if square is attacked by given color"""
        try:
            square_index = chess.parse_square(square)
            color = chess.WHITE if by_color.lower() == 'white' else chess.BLACK
            return self.board.is_attacked_by(color, square_index)
        except ValueError:
            return False
    
    def is_promotion_move(self, from_square: str, to_square: str) -> bool:
        """
        Detect if a move from one square to another would result in pawn promotion.
        
        Args:
            from_square: Starting square in algebraic notation (e.g., 'e7')
            to_square: Destination square in algebraic notation (e.g., 'e8')
            
        Returns:
            True if the move is a pawn promotion, False otherwise
        """
        return PromotionMoveHandler.is_promotion_move(self.board, from_square, to_square)
    
    def get_promotion_moves(self, square: str) -> List[str]:
        """
        Get all possible promotion moves for a pawn at the given square.
        
        Args:
            square: Square in algebraic notation (e.g., 'e7')
            
        Returns:
            List of promotion moves in UCI notation (e.g., ['e7e8q', 'e7e8r', 'e7e8b', 'e7e8n'])
        """
        return PromotionMoveHandler.get_promotion_moves(self.board, square)