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


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] 

@dataclass
class Piece:
    type: PieceType
    color: Color

@dataclass
class BoardMove:
    m_from: BoardPos
    m_to: BoardPos

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


@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)

    # wir lieben python <3                                                          :(
    @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

    def moves_basic_checked(self) -> list[BoardMove]:
        moves: list[BoardMove] = []
        for i, row in enumerate(self.fields):
            for j, field in enumerate(row):
                if field.piece is None:
                    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 possible_moves(self) -> list[BoardMove]:

        return self.moves_basic_checked()
    
# used only for testing purposes
def main(): 
    default_brd = ChessBoard.init_default()
    mvs = default_brd.possible_moves()
    for mv in mvs:
        print(mv)
    print(f"len of moves: { len(mvs)}")

main()