import pygame
import re

from constants import WIDTH, HEIGHT, ROWS, COLS, SQUARE_SIZE, LIGHT_BROWN, DARK_BROWN, PIECES


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

    ### RL env part 
    def reset(self):
        """
        set up the board and the starting player  
        """
        self.board = self.create_board()
        self.current_player = 'white'

    def step(self,action):
        """
        action: A string using natural language and/or chess notation
        Returns: observation, reward, done, info  
        """
        self.handle_input(action)
        # self.move_piece()
        observation = self.board
        reward = self.calculate_reward()
        done = self.check_game_over()
        info = {}
        return observation, reward, done, info

    def render(self, win):
        """
        render the pygame window for visual feedback 
        """
        self.display(win)
        pygame.display.update()

    def close(self):
        pygame.quit()

    def check_game_over(self):
        if self.draw or self.checkmate or self.resign:
            return True
        else:
            return False

    ### Actual game
    def create_board(self):
        board = [['' for _ in range(COLS)] for _ in range(ROWS)]

        # Assign unique IDs to each piece
        # Pawns
        for i in range(COLS):
            board[1][i] = 'black_pawn'
            board[6][i] = 'white_pawn'

        # Rooks
        board[0][0] = board[0][7] = 'black_rook'
        board[7][0] = board[7][7] = 'white_rook'

        # Knights
        board[0][1] = board[0][6] = 'black_knight'
        board[7][1] = board[7][6] = 'white_knight'

        # Bishops
        board[0][2] = board[0][5] = 'black_bishop'
        board[7][2] = board[7][5] = 'white_bishop'

        # Queens
        board[0][3] = 'black_queen'
        board[7][3] = 'white_queen'

        # Kings
        board[0][4] = 'black_king'
        board[7][4] = 'white_king'

        return board

    def display(self, win):
        # draw squares
        for row in range(ROWS):
            for col in range(COLS):
                color = LIGHT_BROWN if (row + col) % 2 == 0 else DARK_BROWN
                pygame.draw.rect(win, color, (col * SQUARE_SIZE, row * SQUARE_SIZE, SQUARE_SIZE, SQUARE_SIZE))

        # draw pieces
        for row in range(ROWS):
            for col in range(COLS):
                piece = self.board[row][col]
                if piece != '':
                    base_piece_name = ''.join([i for i in piece])
                    piece_image = PIECES[base_piece_name]
                    piece_image = pygame.transform.scale(piece_image, (SQUARE_SIZE, SQUARE_SIZE))
                    win.blit(piece_image, (col * SQUARE_SIZE, row * SQUARE_SIZE))

    def handle_input(self, input_text):
        col_map = {'a':0, 'b':1, 'c':2, 'd':3, 'e':4, 'f':5, 'g':6, 'h':7}
        row_map = {'1':7, '2':6, '3':5, '4':4, '5':3, '6':2, '7':1, '8':0}

        # TODO: Implement logic for actions without squares
        # This includes statements such as
        # "I capture your queen with my rook"

        # TODO: Implement logic for actions with mix of squares and names
        # Ex: "Queen to e5"
        # Ex: "a2 captures knight"

        # TODO: Implement logic for actions with only squares
        # Example: "e4 to e5"
        
        # TODO: Probably should spread the different modes across functions for clarity
        # smth like 'figure_pieces()' which figures out the two pieces (and if two pieces interact)
        # and smth like 'figure_squares' which figures out the two squares


    def move_piece(self, start_square, end_square):

        start_col, start_row = start_square
        piece = self.board[start_col][start_row] 

        end_col, end_row = end_square
        end_spot = self.board[end_col][end_row] 
        if end_spot == '':
            self.board[end_col][end_row] = piece
        elif 'black' in end_spot:
            self.capture(end_spot)
            self.board[end_col][end_row] = piece

        self.check_promotion()