from test.base import ChessLibTestBase
from test.chess_ffi import Move
from test.chess_ffi import Board
from test.chess_ffi import BLACK
from test.chess_ffi import WHITE
from test.chess_ffi import is_in_check
from test.chess_ffi import gen_legal_moves
from test.chess_ffi import sq

MAX_MOVES = 256


class TestLegalMoveGen(ChessLibTestBase):
    def _gen_legal(self, board):
        """Support either return-count or out-parameter signatures."""
        moves = (Move * MAX_MOVES)()
        n = gen_legal_moves(board, moves)
        return n, moves


    def test_start_position(self):
        # 20 legal moves from the initial chess position
        fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
        b = Board()
        self.load_fen(fen, board=b)

        n, _ = self._gen_legal(b)
        self.assertEqual(n, 20, "Start position must have 20 legal moves for White")


    def test_kiwipete_depth1_count(self):
        # Kiwipete: perft(1) = 48
        # position 2: https://www.chessprogramming.org/Perft_Results
        fen = "r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq - 0 1"
        b = Board(); self.load_fen(fen, board=b)

        n, moves = self._gen_legal(b)
        self.assertEqual(n, 48, "Kiwipete perft(1) should be 48")


    def test_perft_position_3(self):
        fen = "8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - - 0 1"
        b = Board(); self.load_fen(fen, board=b)

        n, moves = self._gen_legal(b)
        self.assertEqual(n, 14)


    def test_perft_position_4(self):
        fen = "r3k2r/Pppp1ppp/1b3nbN/nP6/BBP1P3/q4N2/Pp1P2PP/R2Q1RK1 w kq - 0 1"
        b = Board(); self.load_fen(fen, board=b)

        n, moves = self._gen_legal(b)
        self.assertEqual(n, 6)


    def test_perft_position_5(self):
        fen = "rnbq1k1r/pp1Pbppp/2p5/8/2B5/8/PPP1NnPP/RNBQK2R w KQ - 1 8"
        b = Board(); self.load_fen(fen, board=b)

        n, moves = self._gen_legal(b)
        self.assertEqual(n, 44)

    
    def test_perft_position_6(self):
        fen = "r4rk1/1pp1qppp/p1np1n2/2b1p1B1/2B1P1b1/P1NP1N2/1PP1QPPP/R4RK1 w - - 0 10"
        b = Board(); self.load_fen(fen, board=b)

        n, moves = self._gen_legal(b)
        self.assertEqual(n, 46)


    def test_stalemate_black_to_move(self):
        # Classic stalemate: Black to move, no legal moves, not in check
        fen = "7k/5Q2/6K1/8/8/8/8/8 b - - 0 1"
        b = Board()
        self.load_fen(fen, board=b)

        moves = (Move * MAX_MOVES)()
        n = gen_legal_moves(b, moves)

        self.assertEqual(n, 0, "Stalemate should have 0 legal moves")
        self.assertFalse(is_in_check(b, BLACK), "Side to move must not be in check in stalemate")
    

    def test_simple_checkmate(self):
        # Black to move; mated (Qg7# with Kg6 support)
        fen = "7k/6Q1/6K1/8/8/8/8/8 b - - 0 1"
        b = Board(); self.load_fen(fen, board=b)
        moves = (Move * MAX_MOVES)()
        n = gen_legal_moves(b, moves)

        self.assertEqual(n, 0, "Checkmated side should have 0 legal moves")
        self.assertTrue(is_in_check(b, BLACK), "Mated side must be in check")


    def test_pinned_piece_cannot_move(self):
            # White knight e2 pinned by rook e8 vs king e1
            fen = "4r3/8/8/8/8/8/4N3/4K3 w - - 0 1"
            b = Board(); self.load_fen(fen, board=b)
            moves = (Move * MAX_MOVES)()
            n = gen_legal_moves(b, moves)

            for move in moves:
                # Pinned knight cannot move from it's square.
                self.assertFalse(getattr(move, "from") == sq("e2"))