Design a Chess Game.

# chess.py, part 1: colours, squares and moves
from __future__ import annotations

import threading
from abc import ABC, abstractmethod
from collections.abc import Iterator
from dataclasses import dataclass
from enum import Enum


class Color(Enum):
    WHITE = "white"
    BLACK = "black"

    @property
    def other(self) -> Color:
        return Color.BLACK if self is Color.WHITE else Color.WHITE


Square = tuple[int, int]  # (file, rank), both 0-7: "e4" is (4, 3)


def sq(name: str) -> Square:
    return "abcdefgh".index(name[0]), int(name[1]) - 1


def name(s: Square) -> str:
    return "abcdefgh"[s[0]] + str(s[1] + 1)


def on_board(s: Square) -> bool:
    return 0 <= s[0] < 8 and 0 <= s[1] < 8


@dataclass(frozen=True)
class Move:
    frm: Square
    to: Square
    promotion: str | None = None  # "q", "r", "b" or "n"
    castle: bool = False
    en_passant: bool = False

    def __str__(self) -> str:
        return name(self.frm) + name(self.to) + (self.promotion or "")

# part 2: each piece knows the squares it attacks; legality is the game's job
class Piece(ABC):
    symbol = "?"

    def __init__(self, color: Color) -> None:
        self.color = color

    def __repr__(self) -> str:
        return self.symbol.upper() if self.color is Color.WHITE else self.symbol

    @abstractmethod
    def attacks(self, board: Board, at: Square) -> Iterator[Square]: ...

    def moves(self, board: Board, at: Square, ep: Square | None) -> Iterator[Move]:
        for to in self.attacks(board, at):  # most pieces: any attacked square not holding a friend
            target = board.squares.get(to)
            if target is None or target.color is not self.color:
                yield Move(at, to)


class Slider(Piece):
    dirs: tuple[tuple[int, int], ...] = ()

    def attacks(self, board: Board, at: Square) -> Iterator[Square]:
        for df, dr in self.dirs:
            s = (at[0] + df, at[1] + dr)
            while on_board(s):
                yield s
                if s in board.squares:
                    break  # the first piece in the way is attacked; nothing beyond it is
                s = (s[0] + df, s[1] + dr)


class Stepper(Piece):
    steps: tuple[tuple[int, int], ...] = ()

    def attacks(self, board: Board, at: Square) -> Iterator[Square]:
        for df, dr in self.steps:
            if on_board(s := (at[0] + df, at[1] + dr)):
                yield s


class Rook(Slider):
    symbol, dirs = "r", ((1, 0), (-1, 0), (0, 1), (0, -1))


class Bishop(Slider):
    symbol, dirs = "b", ((1, 1), (1, -1), (-1, 1), (-1, -1))


class Queen(Slider):
    symbol, dirs = "q", Rook.dirs + Bishop.dirs


class Knight(Stepper):
    symbol, steps = "n", ((1, 2), (2, 1), (2, -1), (1, -2), (-1, -2), (-2, -1), (-2, 1), (-1, 2))


class King(Stepper):
    symbol, steps = "k", Queen.dirs  # castling needs history, so the game adds it


class Pawn(Piece):
    symbol = "p"

    def attacks(self, board: Board, at: Square) -> Iterator[Square]:
        d = 1 if self.color is Color.WHITE else -1
        for df in (-1, 1):
            if on_board(s := (at[0] + df, at[1] + d)):
                yield s

    def moves(self, board: Board, at: Square, ep: Square | None) -> Iterator[Move]:
        d, home, last = (1, 1, 7) if self.color is Color.WHITE else (-1, 6, 0)
        one, two = (at[0], at[1] + d), (at[0], at[1] + 2 * d)
        if on_board(one) and one not in board.squares:
            yield from self._maybe_promote(at, one, last)
            if at[1] == home and two not in board.squares:
                yield Move(at, two)
        for s in self.attacks(board, at):
            target = board.squares.get(s)
            if target is not None and target.color is not self.color:
                yield from self._maybe_promote(at, s, last)
            elif s == ep:
                yield Move(at, s, en_passant=True)

    @staticmethod
    def _maybe_promote(at: Square, to: Square, last: int) -> Iterator[Move]:
        if to[1] == last:
            for p in "qrbn":
                yield Move(at, to, promotion=p)
        else:
            yield Move(at, to)


class Board:
    def __init__(self) -> None:
        self.squares: dict[Square, Piece] = {}

    def king(self, color: Color) -> Square:
        return next(s for s, p in self.squares.items() if isinstance(p, King) and p.color is color)

    def attacked(self, target: Square, by: Color) -> bool:
        return any(p.color is by and target in p.attacks(self, s) for s, p in list(self.squares.items()))

# part 3: the rules that need history, as tables and records
PIECES: dict[str, type[Piece]] = {"k": King, "q": Queen, "r": Rook, "b": Bishop, "n": Knight, "p": Pawn}
START = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq -"
CASTLES = {  # right: king from, king to, rook from, rook to, squares that must be empty
    "K": (sq("e1"), sq("g1"), sq("h1"), sq("f1"), (sq("f1"), sq("g1"))),
    "Q": (sq("e1"), sq("c1"), sq("a1"), sq("d1"), (sq("b1"), sq("c1"), sq("d1"))),
    "k": (sq("e8"), sq("g8"), sq("h8"), sq("f8"), (sq("f8"), sq("g8"))),
    "q": (sq("e8"), sq("c8"), sq("a8"), sq("d8"), (sq("b8"), sq("c8"), sq("d8"))),
}
ROOK_JUMP = {(kf, kt): (rf, rt) for kf, kt, rf, rt, _ in CASTLES.values()}
LOSES_RIGHTS = {  # any move from or onto these squares ends these castling rights
    sq("e1"): {"K", "Q"}, sq("h1"): {"K"}, sq("a1"): {"Q"},
    sq("e8"): {"k", "q"}, sq("h8"): {"k"}, sq("a8"): {"q"},
}


class Status(Enum):
    ONGOING = "ongoing"
    CHECK = "check"
    CHECKMATE = "checkmate"
    STALEMATE = "stalemate"


@dataclass(frozen=True)
class Undo:  # everything a move destroys, so pop() can put it back
    move: Move
    piece: Piece
    captured: Piece | None
    captured_at: Square
    rights: frozenset[str]
    ep: Square | None


class IllegalMove(Exception): ...

# part 4: the game: turn, rights, en passant, history, and legality by make-and-test
class Game:
    def __init__(self, fen: str = START) -> None:
        placement, turn, rights, ep = fen.split()[:4]
        self.board = Board()
        for row, line in enumerate(placement.split("/")):
            f = 0
            for ch in line:
                if ch.isdigit():
                    f += int(ch)
                    continue
                color = Color.WHITE if ch.isupper() else Color.BLACK
                self.board.squares[(f, 7 - row)] = PIECES[ch.lower()](color)
                f += 1
        self.turn = Color.WHITE if turn == "w" else Color.BLACK
        self.rights = frozenset(rights) - {"-"}
        self.ep: Square | None = None if ep == "-" else sq(ep)
        self.history: list[Undo] = []
        self._lock = threading.RLock()

    def push(self, m: Move) -> None:
        sqs = self.board.squares
        piece = sqs.pop(m.frm)
        captured_at = (m.to[0], m.frm[1]) if m.en_passant else m.to
        captured = sqs.pop(captured_at, None)
        sqs[m.to] = PIECES[m.promotion](piece.color) if m.promotion else piece
        if m.castle:
            rf, rt = ROOK_JUMP[(m.frm, m.to)]
            sqs[rt] = sqs.pop(rf)
        self.history.append(Undo(m, piece, captured, captured_at, self.rights, self.ep))
        self.rights = self.rights - LOSES_RIGHTS.get(m.frm, set()) - LOSES_RIGHTS.get(m.to, set())
        double_step = isinstance(piece, Pawn) and abs(m.to[1] - m.frm[1]) == 2
        self.ep = (m.frm[0], (m.frm[1] + m.to[1]) // 2) if double_step else None
        self.turn = self.turn.other

    def pop(self) -> Move:
        u = self.history.pop()
        sqs, m = self.board.squares, u.move
        del sqs[m.to]
        sqs[m.frm] = u.piece  # a promoted queen goes back to being a pawn
        if u.captured is not None:
            sqs[u.captured_at] = u.captured
        if m.castle:
            rf, rt = ROOK_JUMP[(m.frm, m.to)]
            sqs[rf] = sqs.pop(rt)
        self.rights, self.ep, self.turn = u.rights, u.ep, self.turn.other
        return m

    def _castles(self) -> Iterator[Move]:
        for right in sorted(self.rights):
            k_from, k_to, r_from, r_to, empty = CASTLES[right]
            if (Color.WHITE if right.isupper() else Color.BLACK) is not self.turn:
                continue
            if any(s in self.board.squares for s in empty):
                continue
            # not out of, through, or into check (the king crosses the rook's landing square)
            if any(self.board.attacked(s, self.turn.other) for s in (k_from, r_to, k_to)):
                continue
            yield Move(k_from, k_to, castle=True)

    def legal_moves(self) -> list[Move]:
        with self._lock:  # push/pop below change the board, so even readers hold the lock
            mine = [(s, p) for s, p in self.board.squares.items() if p.color is self.turn]
            pseudo = [m for s, p in mine for m in p.moves(self.board, s, self.ep)] + list(self._castles())
            legal = []
            for m in pseudo:
                self.push(m)
                if not self.board.attacked(self.board.king(self.turn.other), self.turn):
                    legal.append(m)  # the side that just moved has not left its king in check
                self.pop()
            return legal

    def in_check(self) -> bool:
        return self.board.attacked(self.board.king(self.turn), self.turn.other)

    def status(self) -> Status:
        with self._lock:
            if self.legal_moves():
                return Status.CHECK if self.in_check() else Status.ONGOING
            return Status.CHECKMATE if self.in_check() else Status.STALEMATE

    def move(self, uci: str, by: Color | None = None) -> Status:
        frm, to, promo = sq(uci[:2]), sq(uci[2:4]), uci[4:] or None
        with self._lock:
            if by is not None and by is not self.turn:
                raise IllegalMove(f"it is {self.turn.value}'s turn")
            for m in self.legal_moves():
                if m.frm == frm and m.to == to and m.promotion == (promo or ("q" if m.promotion else None)):
                    self.push(m)
                    return self.status()
            raise IllegalMove(f"{uci} is not a legal move for {self.turn.value}")

    def undo(self) -> Move:
        with self._lock:
            if not self.history:
                raise IllegalMove("nothing to undo")
            return self.pop()

# part 5: short games that touch every rule, then perft against published counts
def perft(g: Game, depth: int) -> int:
    moves = g.legal_moves()
    if depth == 1:
        return len(moves)
    total = 0
    for m in moves:
        g.push(m)
        total += perft(g, depth - 1)
        g.pop()
    return total


if __name__ == "__main__":
    g = Game()
    for uci in ["f2f3", "e7e5", "g2g4"]:
        g.move(uci)
    for bad, by in [("d8h4", Color.WHITE), ("e8e6", None)]:
        try:
            g.move(bad, by)
        except IllegalMove as e:
            print("rejected:", e)
    print("d8h4:", g.move("d8h4").value)  # fool's mate
    print("undo", g.undo(), "->", g.status().value, "with", len(g.legal_moves()), "replies for black")

    ep = Game()
    for uci in ["e2e4", "a7a6", "e4e5", "d7d5"]:
        ep.move(uci)
    print("e5d6 en passant:", ep.move("e5d6").value, "| d5 now:", ep.board.squares.get(sq("d5")))
    ep.undo()
    print("undo: d5 is", ep.board.squares.get(sq("d5")), "| en passant square", name(ep.ep))

    c = Game()
    for uci in ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4", "g8f6", "e1g1"]:
        c.move(uci)
    print("castled: g1 =", c.board.squares[sq("g1")], "f1 =", c.board.squares[sq("f1")], "| rights:", "".join(sorted(c.rights)))
    try:
        Game("4kr2/8/8/8/8/8/8/4K2R w K -").move("e1g1")  # the black rook watches f1
    except IllegalMove as e:
        print("through check:", e)

    print("a7a8:", Game("7k/P7/8/8/8/8/8/K7 w - -").move("a7a8").value)  # the new queen checks along rank 8
    under = Game("7k/P7/8/8/8/8/8/K7 w - -")
    print("a7a8n:", under.move("a7a8n").value, under.board.squares[sq("a8")], "| undo:", under.undo(), under.board.squares[sq("a7")])
    print("f2f7:", Game("7k/8/6K1/8/8/8/5Q2/8 w - -").move("f2f7").value)

    print("perft start:", [perft(Game(), d) for d in (1, 2, 3)])
    kiwipete = Game("r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq -")
    print("perft kiwipete:", [perft(kiwipete, d) for d in (1, 2, 3)])

rejected: it is black's turn
rejected: e8e6 is not a legal move for black
d8h4: checkmate
undo d8h4 -> ongoing with 30 replies for black
e5d6 en passant: ongoing | d5 now: None
undo: d5 is p | en passant square d6
castled: g1 = K f1 = R | rights: kq
through check: e1g1 is not a legal move for white
a7a8: check
a7a8n: ongoing N | undo: a7a8n P
f2f7: stalemate
perft start: [20, 400, 8902]
perft kiwipete: [48, 2039, 97862]