from enum import Enum, auto
from dataclasses import dataclass
from typing import Optional, List, Tuple


class PieceType(Enum):
    PAWN = auto()
    KNIGHT = auto()
    BISHOP = auto()
    ROOK = auto()
    QUEEN = auto()
    KING = auto()


class Color(Enum):
    WHITE = auto()
    BLACK = auto()


@dataclass
class BoardPos:
    p: Tuple[int, int]

    def __str__(self) -> str:
        return f"({self.p[0]},{self.p[1]})"


@dataclass
class Piece:
    type: PieceType
    color: Color

    def char(self) -> str:
        """Single-character representation. Uppercase = White, lowercase = Black."""
        mapping = {
            PieceType.PAWN: "p",
            PieceType.KNIGHT: "n",
            PieceType.BISHOP: "b",
            PieceType.ROOK: "r",
            PieceType.QUEEN: "q",
            PieceType.KING: "k",
        }
        ch = mapping[self.type]
        return ch.upper() if self.color == Color.WHITE else ch

    def __str__(self) -> str:
        return self.char()

    def __repr__(self) -> str:
        return f"Piece({self.type.name},{self.color.name})"


@dataclass
class BoardMove:
    m_from: BoardPos
    m_to: BoardPos

    def __str__(self) -> str:
        return f"{self.m_from}->{self.m_to}"


@dataclass
class BoardField:
    piece: Optional[Piece] = None

    def __str__(self) -> str:
        return str(self.piece) if self.piece is not None else "."


@dataclass
class ChessBoard:
    fields: List[List[BoardField]]

    def place(self, row: int, col: int, piece_type: PieceType, color: Color):
        self.fields[row][col].piece = Piece(piece_type, color)

    # initialize default starting position
    @classmethod
    def init_default(cls) -> "ChessBoard":
        brd = cls(
            [[BoardField() for _ in range(8)] for _ in range(8)]
        )

        # place pawns
        for col in range(8):
            brd.place(1, col, PieceType.PAWN, Color.BLACK)
            brd.place(6, col, PieceType.PAWN, Color.WHITE)

        # back rank order
        back_rank = [
            PieceType.ROOK,
            PieceType.KNIGHT,
            PieceType.BISHOP,
            PieceType.QUEEN,
            PieceType.KING,
            PieceType.BISHOP,
            PieceType.KNIGHT,
            PieceType.ROOK,
        ]

        # place black back rank
        for col, piece_type in enumerate(back_rank):
            brd.place(0, col, piece_type, Color.BLACK)

        # place white back rank
        for col, piece_type in enumerate(back_rank):
            brd.place(7, col, piece_type, Color.WHITE)

        return brd
    
    def moves_unchecked(self, piece: Piece, pos: BoardPos) -> List[List[Tuple[int, int]]]:
        row, column = pos.p
        rays: List[List[Tuple[int, int]]] = []

        if piece.type == PieceType.KING:
            king_offsets = [(1, 1), (1, 0), (0, 1), (-1, -1), (-1, 0), (0, -1), (1, -1), (-1, 1)]
            rays.append([(row + dr, column + dc) for dr, dc in king_offsets])
        elif piece.type == PieceType.KNIGHT:
            knight_offsets = [
                (2, 1), (2, -1), (-2, 1), (-2, -1),
                (1, 2), (1, -2), (-1, 2), (-1, -2)
            ]
            rays.append([(row + dr, column + dc) for dr, dc in knight_offsets])
        elif piece.type == PieceType.PAWN:
            # correct forward direction and start row depending on color
            if piece.color == Color.WHITE:
                forward = -1
                start_row = 6
            else:
                forward = 1
                start_row = 1

            forward_ray: List[Tuple[int, int]] = []
            forward_ray.append((row + forward, column))
            # if the pawn is at the start row, include the double-step
            if row == start_row:
                forward_ray.append((row + 2 * forward, column))
            rays.append(forward_ray)

            # capture moves
            captures = [(row + forward, column - 1), (row + forward, column + 1)]
            rays.extend([[(rr, cc)] for rr, cc in captures])
        elif piece.type == PieceType.BISHOP:
            directions = [(1, 1), (1, -1), (-1, 1), (-1, -1)]
            for dr, dc in directions:
                ray = []
                for k in range(1, 8): 
                    ray.append((row + k * dr, column + k * dc))
                rays.append(ray)
        elif piece.type == PieceType.ROOK:
            directions = [(1, 0), (-1, 0), (0, 1), (0, -1)]
            for dr, dc in directions:
                ray = []
                for k in range(1, 8):
                    ray.append((row + k * dr, column + k * dc))
                rays.append(ray)
        elif piece.type == PieceType.QUEEN:
            directions = [
                (1, 0), (-1, 0), (0, 1), (0, -1),
                (1, 1), (1, -1), (-1, 1), (-1, -1)
            ]
            for dr, dc in directions:
                ray = []
                for k in range(1, 8):
                    ray.append((row + k * dr, column + k * dc))
                rays.append(ray)

        return rays

    def on_board(self, rr: int, cc: int) -> bool:
        return 0 <= rr < 8 and 0 <= cc < 8

    # takes the color of the player whos possible moves will be returned
    def moves_basic_checked(self, color: Color) -> List[BoardMove]:
        moves: List[BoardMove] = []
        for i, row in enumerate(self.fields):
            for j, field in enumerate(row):
                if field.piece is None:
                    continue
                if field.piece.color != color:
                    continue
                src_pos = BoardPos((i, j))
                piece = field.piece
                rays = self.moves_unchecked(piece, src_pos)

                for ray in rays:
                    # detect pawn capture ray and dont allow movement in that ray without capturing
                    is_pawn_capture_ray = (
                        piece.type == PieceType.PAWN
                        and len(ray) == 1
                        and ray[0][1] != j
                        and abs(ray[0][0] - i) == 1
                    )

                    # non sliding rays get treated individually
                    non_sliding = piece.type in (PieceType.KNIGHT, PieceType.KING) or is_pawn_capture_ray

                    for (tr, tc) in ray:
                        if not self.on_board(tr, tc):
                            if non_sliding:
                                continue
                            break
                        target_field = self.fields[tr][tc]
                        if target_field.piece is None:
                            # empty target
                            if is_pawn_capture_ray:
                                # pawn cant move diagonally into empty square
                                break
                            moves.append(BoardMove(src_pos, BoardPos((tr, tc))))
                            continue    
                        else:
                            # occupied
                            if target_field.piece.color == piece.color:
                                break
                            else:
                                # opponent piece
                                moves.append(BoardMove(src_pos, BoardPos((tr, tc))))
                                break
        return moves


    def __str__(self) -> str:
        """
        example string repr of starting pos
        rnbqkbnr
        pppppppp
        ........
        ........
        ........
        ........
        PPPPPPPP
        RNBQKBNR
        """
        lines: List[str] = []
        for r in range(8):
            row_chars = [str(self.fields[r][c]) for c in range(8)]
            lines.append("".join(row_chars))
        return "\n".join(lines)

    @classmethod
    def from_string(cls, s: str) -> "ChessBoard":
        lines = s.strip().split("\n")

        brd = cls([[BoardField() for _ in range(8)] for _ in range(8)])

        char_to_piece = {
            'p': PieceType.PAWN,
            'n': PieceType.KNIGHT,
            'b': PieceType.BISHOP,
            'r': PieceType.ROOK,
            'q': PieceType.QUEEN,
            'k': PieceType.KING,
        }

        for r, line in enumerate(lines):
            for c, ch in enumerate(line):
                if ch == '.':
                    continue
                lower = ch.lower()
                if lower not in char_to_piece:
                    raise ValueError(f"Invalid piece character '{ch}' at ({r},{c})")
                piece_type = char_to_piece[lower]
                color = Color.WHITE if ch.isupper() else Color.BLACK
                brd.place(r, c, piece_type, color)

        return brd


# used only for testing purposes
def main(): 
    default_brd = ChessBoard.init_default()

    # standardized expected string for initial position
    expected_start = (
        "rnbqkbnr\n"
        "pppppppp\n"
        "........\n"
        "........\n"
        "........\n"
        "........\n"
        "PPPPPPPP\n"
        "RNBQKBNR"
    )

    # test string representation of board
    actual_start = str(default_brd)
    if actual_start != expected_start:
        raise AssertionError(f"initial board mismatch:\nexpected:\n{expected_start}\n\nactual:\n{actual_start}")

    # test num moves
    mvs = default_brd.moves_basic_checked(Color.BLACK)
    expected_move_count = 20  # 16 pawn moves and 4 knight
    if len(mvs) != expected_move_count:
        raise AssertionError(f"Initial move count for black mismatch: expected {expected_move_count}, got {len(mvs)}")
    

    mvs = default_brd.moves_basic_checked(Color.WHITE)
    expected_move_count = 20  # 16 pawn moves and 4 knight
    if len(mvs) != expected_move_count:
        raise AssertionError(f"Initial move count for white mismatch: expected {expected_move_count}, got {len(mvs)}")


    # test loading 
    loaded_brd = ChessBoard.from_string(expected_start)
    if str(loaded_brd) != expected_start:
        raise AssertionError(
            "board loaded from string does not match original representation"
        )

    print(f"all tests passed")


if __name__ == '__main__':
    main()