# chessboard.py

import pygame
import re
# Import constants from the constants module
from .constants import COLS, ROWS, LIGHT_BROWN, DARK_BROWN, SQUARE_SIZE, PIECES

class ChessBoard:
    def __init__(self):
        self.reset()

    def reset(self):
        self.board = self.create_board()
        self.current_player = 'white'
        # Return the initial observation (e.g., the initial board state)
        return self.board

    def step(self, action):
        """
        action: A string containing the natural language input, e.g., 'Move pawn from e2 to e4'
        Returns: observation, reward, done, info
        """
        success = self.handle_input(action)
        observation = self.board
        reward = 1 if success else -1
        done = self.check_game_over()
        info = {}
        return observation, reward, done, info

    def render(self, win):
        self.draw(win)
        pygame.display.update()

    def close(self):
        pygame.quit()

    # Include all your existing methods here (create_board, draw, is_legal, etc.)

    def handle_input(self, input_str):
        # Parse the natural language input and perform the move
        try:
            input_str = input_str.lower()
            piece_names = ['pawn', 'rook', 'knight', 'bishop', 'queen', 'king']

            # Find squares in the input (e.g., e2, d4)
            squares = re.findall(r'\b[a-h][1-8]\b', input_str)

            # Find the piece name
            piece_type = None
            for piece in piece_names:
                if piece in input_str:
                    piece_type = piece
                    break

            if not piece_type or len(squares) < 2:
                return False

            start_square = squares[0]
            dest_square = squares[1]

            if self.move_piece(piece_type, start_square, dest_square):
                # Switch player after a successful move
                self.current_player = 'black' if self.current_player == 'white' else 'white'
                return True
            else:
                return False

        except Exception as e:
            print("Invalid input format or move!", e)
            return False

    def check_game_over(self):
        # Implement logic to check if the game is over
        return False  # Placeholder