from __future__ import annotations

import pygame
from copy import deepcopy
from .constants import BLACK, ROWS, GREEN, SQUARE_SIZE, COLS, WHITE
from .piece import Piece


class Board:
    def __init__(self) -> None:
        """
        Constructor for the Board class
        :return: None
        """
        self.board = []
        self.greenLeft = self.whiteLeft = 12
        self.greenKings = self.whiteKings = 0
        self.green = (144, 184, 59)
        self._createBoard()

    def _drawSquares(self, win: pygame.display) -> None:
        """
        Draws the squares on the board
        :param win: The window
        """
        win.fill(BLACK)
        for row in range(ROWS):
            for col in range(row % 2, ROWS, 2):
                pygame.draw.rect(win, self.green, (row * SQUARE_SIZE, col * SQUARE_SIZE, SQUARE_SIZE, SQUARE_SIZE))

    def _createBoard(self) -> None:
        """
        Creates a board representation of the game
        :return: None
        """
        for row in range(ROWS):
            self.board.append([])
            for col in range(COLS):
                if col % 2 == ((row + 1) % 2):
                    if row < 3:
                        self.board[row].append(Piece(row, col, WHITE))
                        continue

                    if row > 4:
                        self.board[row].append(Piece(row, col, GREEN))
                        continue

                    self.board[row].append(0)
                    continue

                self.board[row].append(0)

    def draw(self, win: pygame.display) -> None:
        """
        Draws the pieces on the board
        :param win: The window
        :return: None
        """
        self._drawSquares(win)
        for row in range(ROWS):
            for col in range(COLS):
                piece = self.board[row][col]
                if piece != 0:
                    piece.draw(win)

    def move(self, piece: Piece, row: int, col: int) -> None:
        """
        Moves a piece and make it a king if it reaches the end of the board
        :param piece: Piece to move
        :param row: Row to move to
        :param col: Column to move to
        :return: None
        """
        self.board[piece.row][piece.col], self.board[row][col] = self.board[row][col], self.board[piece.row][piece.col]
        piece.move(row, col)

        if row == ROWS - 1 or row == 0:
            piece.makeKing()

            if piece.colour == WHITE:
                self.whiteKings += 1

            if piece.colour == GREEN:
                self.greenKings += 1

    def remove(self, skipped: tuple) -> None:
        """
        Removes a piece from the board
        :param skipped: A tuple of the piece to remove
        """
        for piece in skipped:
            self.board[piece.row][piece.col] = 0
            if piece != 0:
                if piece.colour == GREEN:
                    self.greenLeft -= 1
                    continue
                self.whiteLeft -= 1

    def getAllMoves(self, colour: int) -> list:
        """
        Gets all the possible moves for a player
        :param colour: colour of the player
        :return:
        """
        moves = []
        possibleMoves = []
        possiblePieces = []
        pieces = self.getAllPieces(colour)
        hasForcedCapture = False

        for piece in pieces:
            validMoves = self.getValidMoves(piece)

            # Check if there are forced capture moves for this piece
            forcedCaptureMoves = [move for move, skip in validMoves.items() if skip]
            if forcedCaptureMoves:
                hasForcedCapture = True
                possiblePieces.append(piece)
                possibleMoves.append({move: skip for move, skip in validMoves.items() if skip})

        if hasForcedCapture:
            # If there are forced capture moves, consider only those
            for i in range(len(possibleMoves)):
                for move, skip in possibleMoves[i].items():
                    tempBoard = deepcopy(self)
                    tempPiece = tempBoard.getPiece(possiblePieces[i].row, possiblePieces[i].col)
                    newBoard = self._simulateMove(tempPiece, move, tempBoard, skip)
                    moves.append(newBoard)
        else:
            # If no forced capture moves, consider all valid moves
            for piece in pieces:
                validMoves = self.getValidMoves(piece)
                for move, skip in validMoves.items():
                    tempBoard = deepcopy(self)
                    tempPiece = tempBoard.getPiece(piece.row, piece.col)
                    newBoard = self._simulateMove(tempPiece, move, tempBoard, skip)
                    moves.append(newBoard)

        return moves

    def _simulateMove(self, piece: Piece, move: list, board: Board, skip: tuple) -> Board:
        """
        Simulates a move on the board
        :param piece: Piece to move
        :param move:  Move to make
        :param board: Board to make the move on
        :param skip: Tuple of pieces to skip
        :return: Board after the move
        """
        board.move(piece, move[0], move[1])
        if skip:
            board.remove(skip)

        return board

    def getPiece(self, row: int, col: int) -> Piece:
        """
        Gets a piece from the board
        :param row: Row of the piece
        :param col: Column of the piece
        :return: Piece
        """
        return self.board[row][col]

    def winner(self):
        if self.greenLeft <= 0:
            return WHITE

        if self.whiteLeft <= 0:
            return GREEN

        return None

    def getValidMoves(self, piece: Piece) -> dict:
        """
        Gets all the valid moves for a piece
        :param piece: Piece to get the moves for
        :return: dictionary of moves
        """
        moves = {}
        forcedCapture = {}
        left = piece.col - 1
        right = piece.col + 1
        row = piece.row
        if piece.colour == GREEN or piece.king:
            moves.update(self._traverseLeft(row - 1, max(row - 3, -1), -1, piece.colour, left))
            moves.update(self._traverseRight(row - 1, max(row - 3, -1), -1, piece.colour, right))
        if piece.colour == WHITE or piece.king:
            moves.update(self._traverseLeft(row + 1, min(row + 3, ROWS), 1, piece.colour, left))
            moves.update(self._traverseRight(row + 1, min(row + 3, ROWS), 1, piece.colour, right))

        if len(moves.values()) <= 1:
            return moves

        movesValues = list(moves.values())
        movesKeys = list(moves.keys())

        forced = {}

        for i in range(len(movesKeys)):
            if not movesValues[i]:
                forced[movesKeys[i]] = moves[movesKeys[i]]
        if len(forced) != len(moves):
            forced.clear()
            for i in range(len(movesKeys)):
                if movesValues[i]:
                    forced[movesKeys[i]] = moves[movesKeys[i]]
            if len(forced) != len(moves):
                for i in range(len(movesKeys)):
                    if movesValues[i]:
                        forcedCapture[movesKeys[i]] = moves[movesKeys[i]]
            else:
                forcedCapture = forced
        else:
            forcedCapture = forced

        return forcedCapture

    def scoreOfTheBoard(self) -> int:
        """
        Calculates the score of the board
        :return: score of the board
        """
        return self.whiteLeft - self.greenLeft

    def getAllPieces(self, colour):
        """
        Gets all the pieces of a player
        :param colour: Piece colour
        :return: Pieces of the player
        """
        pieces = []
        for row in self.board:
            for piece in row:
                if piece != 0 and piece.colour == colour:
                    pieces.append(piece)
        return pieces

    def _traverseLeft(self, start: int, stop: int, step: int, colour: int, left: int, skipped: list = []) -> dict:
        """
        Traverses the left side of the board
        :param start: Start position
        :param stop:  Stop position
        :param step:  Step size
        :param colour: colour of the player
        :param left: Left position
        :param skipped: List of pieces to skip
        :return: dictionary of moves
        """
        moves = {}
        last = []
        for row in range(start, stop, step):
            if left < 0:
                break
            mvs = self._traverse(row, left, skipped, moves, step, last, colour)
            if mvs is None:
                break
            elif isinstance(mvs, list):
                last = mvs
            else:
                moves.update(mvs)
            left -= 1
        return moves

    def _traverseRight(self, start: int, stop: int, step: int, colour: int, right: int, skipped: list = []) -> dict:
        """
        Traverses the left side of the board
        :param start: Start position
        :param stop:  Stop position
        :param step:  Step size
        :param colour: colour of the player
        :param right: Right position
        :param skipped: List of pieces to skip
        :return: dictionary of moves
        """
        moves = {}
        last = []
        for row in range(start, stop, step):
            if right >= COLS:
                break

            mvs = self._traverse(row, right, skipped, moves, step, last, colour)
            if mvs is None:
                break
            elif isinstance(mvs, list):
                last = mvs
            else:
                moves.update(mvs)

            right += 1
        return moves

    def _traverse(self, row: int, col: int, skipped: list, moves: dict, step: int, last: list, colour: int) -> list or None:
        """
        Traverses the board
        :param row: Row to traverse
        :param col: Column to traverse
        :param skipped: List of pieces to jump
        :param moves: Dictionary of moves
        :param step: Step size
        :param last: List of last pieces
        :param colour: Colour of the player
        :return: list of last pieces or None
        """
        current = self.board[row][col]
        if current == 0:
            if skipped and not last:
                return None
            elif skipped:
                moves[(row, col)] = last + skipped
            else:
                moves[(row, col)] = last

            if last:
                if step == -1:
                    rowCalc = max(row - 3, 0)
                else:
                    rowCalc = min(row + 3, ROWS)
                moves.update(self._traverseLeft(row + step, rowCalc, step, colour, col - 1, skipped=last))
                moves.update(self._traverseRight(row + step, rowCalc, step, colour, col + 1, skipped=last))
            return None
        elif current.colour == colour:
            return None
        else:
            last = [current]
            return last

    def step(self, move: int, colour: int) -> None:
        """
        Takes a move and executes it
        :param move: The move to execute
        :param colour: The colour of the player
        :return: None
        """
        start, end = self._decode(move)
        start[0] = start[0] - 1
        start[1] = start[1] - 1
        end[0] = end[0] - 1
        end[1] = end[1] - 1
        reward = 0
        done = False
        piece = self.getPiece(start[0], start[1])
        if piece == 0:
            newStart = end
            end = start
            start = newStart
            piece = self.getPiece(start[0], start[1])
        moves = self.getValidMoves(piece)
        for move, skip in moves.items():
            if tuple(end) == move:
                self._simulateMove(piece, move, self, skip)
                if len(skip) == 1:
                    reward = 2
                    break
                if len(skip) > 1:
                    reward = 3 + len(skip) * 0.2
                    break
                reward = -0.5
                break

        if self.winner() == colour:
            done = True
            reward = 10
        return reward, self, done


    def _decode(self, move: int) -> tuple:
        """
        Decodes the move from a integer to a start and end tuple
        :param move: The move to decode
        :return: Start and end tuple
        """
        # Split digits back out
        str_code = str(move)
        # print(str_code)
        start_row = int(str_code[0])
        start_col = int(str_code[1])
        end_row = int(str_code[2])
        end_col = int(str_code[3])
        # Reconstruct positions
        start = [start_row, start_col]
        end = [end_row, end_col]
        return start, end

    # def reset(self):
    #     self.board = []
    #     self.whiteLeft = self.greenLeft = 12
    #     self.whiteKings = self.greenKings = 0
    #     self._createBoard()
    #     return self.board