From 496207861e3cb8cd0f6ac0b4a50b8d86096e2897 Mon Sep 17 00:00:00 2001
From: Karma Riuk <riukkarma@gmail.com>
Date: Fri, 31 Jan 2025 16:35:41 +0100
Subject: [PATCH] implemented castling

---
 src/logic/board.py       | 48 ++++++++++++++++++++++++++++++-------
 src/logic/pieces/king.py | 51 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 90 insertions(+), 9 deletions(-)

diff --git a/src/logic/board.py b/src/logic/board.py
index c0e0ce5..480a97a 100644
--- a/src/logic/board.py
+++ b/src/logic/board.py
@@ -136,6 +136,9 @@ class Board:
                 return True
         return False
 
+    def castling_writes_for(self, colour: Colour) -> set[CastleSide]:
+        return self._white_castling_write if colour == Colour.WHITE else self._black_castling_write
+
     def make_move(self, move: Move) -> "Board":
         dest_piece = self.piece_at(move.pos.x, move.pos.y) 
 
@@ -151,16 +154,45 @@ class Board:
         ret._en_passant_target = self._en_passant_target
 
         piece = move.piece
+        pieces_moving, other_pieces = (ret._white, ret._black) if piece.colour == Colour.WHITE else (ret._black, ret._white)
+
+        del pieces_moving[piece.pos]
+        pieces_moving[move.pos] = piece.move_to(move.pos)
+        if move.pos in other_pieces:
+            del other_pieces[move.pos]
+
+        if move.castle_side == CastleSide.King:
+            rook_pos = Position(7, piece.pos.y)
+            assert rook_pos in pieces_moving and type(pieces_moving[rook_pos]) == Rook, "Either rook is absent from the king side or you are trying to castle with something else than a rook..."
+            del pieces_moving[rook_pos]
+            new_rook_pos = Position(5, piece.pos.y)
+            pieces_moving[new_rook_pos] = Rook(new_rook_pos, piece.colour)
+
+        elif move.castle_side == CastleSide.Queen:
+            rook_pos = Position(0, piece.pos.y)
+            assert rook_pos in pieces_moving and type(pieces_moving[rook_pos]) == Rook, "Either rook is absent from the queen side or you are trying to castle with something else than a rook..."
+            del pieces_moving[rook_pos]
+            new_rook_pos = Position(3, piece.pos.y)
+            pieces_moving[new_rook_pos] = Rook(new_rook_pos, piece.colour)
 
         if piece.colour == Colour.WHITE:
-            del ret._white[piece.pos]
-            ret._white[move.pos] = piece.move_to(move.pos)
-            if move.pos in ret._black:
-                del ret._black[move.pos]
+            if type(piece) == King:
+                ret._white_castling_write = set()
+
+            if type(piece) == Rook:
+                if piece.pos.x == 0 and CastleSide.Queen in ret._white_castling_write:
+                    ret._white_castling_write.remove(CastleSide.Queen)
+                elif piece.pos.x == 7 and CastleSide.King in ret._white_castling_write:
+                    ret._white_castling_write.remove(CastleSide.King)
         else:
-            del ret._black[piece.pos]
-            ret._black[move.pos] = piece.move_to(move.pos)
-            if move.pos in ret._white:
-                del ret._white[move.pos]
+            if type(piece) == King:
+                ret._black_castling_write = set()
+
+            if type(piece) == Rook:
+                if piece.pos.x == 0 and CastleSide.Queen in ret._black_castling_write:
+                    ret._black_castling_write.remove(CastleSide.Queen)
+                elif piece.pos.x == 7 and CastleSide.King in ret._black_castling_write:
+                    ret._black_castling_write.remove(CastleSide.King)
+
 
         return ret
diff --git a/src/logic/pieces/king.py b/src/logic/pieces/king.py
index 0c49de9..33ec573 100644
--- a/src/logic/pieces/king.py
+++ b/src/logic/pieces/king.py
@@ -1,8 +1,12 @@
+from logic.move import CastleSide, Move
+from logic.position import Position
 from .piece import Piece
 
 class King(Piece):
-    def legal_moves(self, board: "Board") -> list["Move"]:
+    def legal_moves(self, board: "Board") -> list[Move]:
         ret = []
+
+        # -- Regular moves
         for dx in [-1, 0, 1]:
             for dy in [-1, 0, 1]:
                 if dx == 0 and dy == 0: # skip current position
@@ -15,6 +19,51 @@ class King(Piece):
                     if not board_after_move.is_check_for(self.colour):
                         ret.append(move)
 
+        # -- Castles
+        castling_writes = board.castling_writes_for(self.colour)
+        if len(castling_writes) == 0:
+            return ret
+        print(castling_writes)
+
+        if CastleSide.King in castling_writes:
+            clear = True
+            for dx in range(1, 3):
+                x = self.pos.x + dx
+                y = self.pos.y
+                if board.piece_at(x, y) is not None:
+                    clear = False
+                    break
+
+                move = self._move_for_position(board, x, y)
+                board_after_move = board.make_move(move)
+                if board_after_move.is_check_for(self.colour):
+                    clear = False
+                    break
+
+            if clear:
+                ret.append(Move(self, Position(6, self.pos.y), castle_side=CastleSide.King))
+
+        if CastleSide.Queen in castling_writes:
+            clear = True
+            for dx in range(1, 3):
+                x = self.pos.x - dx
+                y = self.pos.y
+
+                if board.piece_at(x, y) is not None:
+                    clear = False
+                    break
+
+                move = self._move_for_position(board, x, y)
+                board_after_move = board.make_move(move)
+                if board_after_move.is_check_for(self.colour):
+                    clear = False
+                    break
+
+            if clear:
+                ret.append(Move(self, Position(2, self.pos.y), castle_side=CastleSide.Queen))
+
+        print(ret)
+
         return ret