from copy import deepcopy

from .pieces import Color, Piece, PieceType, Position


class Board:
    COLS = 9
    ROWS = 10

    def __init__(self):
        self.grid: list[list[Piece | None]] = [
            [None] * self.COLS for _ in range(self.ROWS)
        ]
        self.pieces: list[Piece] = []
        self._setup_initial_position()

    def _setup_initial_position(self):
        back_rank = [
            PieceType.CHARIOT,
            PieceType.HORSE,
            PieceType.ELEPHANT,
            PieceType.ADVISOR,
            PieceType.GENERAL,
            PieceType.ADVISOR,
            PieceType.ELEPHANT,
            PieceType.HORSE,
            PieceType.CHARIOT,
        ]

        for col, ptype in enumerate(back_rank):
            self._place(ptype, Color.BLACK, Position(col, 0))
        self._place(PieceType.CANNON, Color.BLACK, Position(1, 2))
        self._place(PieceType.CANNON, Color.BLACK, Position(7, 2))
        for col in range(0, 9, 2):
            self._place(PieceType.SOLDIER, Color.BLACK, Position(col, 3))

        for col, ptype in enumerate(back_rank):
            self._place(ptype, Color.RED, Position(col, 9))
        self._place(PieceType.CANNON, Color.RED, Position(1, 7))
        self._place(PieceType.CANNON, Color.RED, Position(7, 7))
        for col in range(0, 9, 2):
            self._place(PieceType.SOLDIER, Color.RED, Position(col, 6))

    def _place(self, piece_type: PieceType, color: Color, pos: Position):
        piece = Piece(piece_type, color, pos)
        self.grid[pos.row][pos.col] = piece
        self.pieces.append(piece)

    def get_piece(self, pos: Position) -> Piece | None:
        if 0 <= pos.col < self.COLS and 0 <= pos.row < self.ROWS:
            return self.grid[pos.row][pos.col]
        return None

    def move_piece(self, from_pos: Position, to_pos: Position) -> Piece | None:
        piece = self.grid[from_pos.row][from_pos.col]
        if piece is None:
            return None

        captured = self.grid[to_pos.row][to_pos.col]
        if captured:
            captured.position = to_pos
            self.pieces.remove(captured)

        self.grid[from_pos.row][from_pos.col] = None
        self.grid[to_pos.row][to_pos.col] = piece
        piece.position = to_pos

        return captured

    def clone(self) -> "Board":
        return deepcopy(self)

    def get_pieces(self, color: Color | None = None) -> list[Piece]:
        if color is not None:
            return [p for p in self.pieces if p.color == color]
        return list(self.pieces)

    def find_general(self, color: Color) -> Piece | None:
        for p in self.pieces:
            if p.piece_type == PieceType.GENERAL and p.color == color:
                return p
        return None

    def position_key(self) -> tuple:
        key = []
        for row in range(self.ROWS):
            for col in range(self.COLS):
                piece = self.grid[row][col]
                if piece:
                    key.append((col, row, piece.piece_type.value, piece.color.value))
        return tuple(key)
