from test.base import ChessLibTestBase
from test.chess_ffi import Board
from test.chess_ffi import BLACK
from test.chess_ffi import WHITE
from test.chess_ffi import sq
from test.chess_ffi import popcount
from test.chess_ffi import load_fen


class TestFenLoading(ChessLibTestBase):


    def rank_mask(self, r):
        return sum(1 << (r*8 + f) for f in range(8))
    

    def test_startpos_fields_and_occupancies(self):
        fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
        rc = self.load_fen(fen)
        
        self.assertEqual(rc, 0)
        self.assertEqual(self.board.side_to_move, WHITE)
        self.assertEqual(self.board.castling_rights, 0b1111)
        self.assertEqual(self.board.ep_square, -1)
        self.assertEqual(self.board.halfmove_clock, 0)
        self.assertEqual(self.board.fullmove_number, 1)

        white_expected = self.rank_mask(0) | self.rank_mask(1)
        black_expected = self.rank_mask(6) | self.rank_mask(7)
        self.assertEqual(int(self.board.occ[WHITE]), white_expected)
        self.assertEqual(int(self.board.occ[BLACK]), black_expected)
        self.assertEqual(int(self.board.occ[2]),     white_expected | black_expected)

        self.assertEqual(popcount(int(self.board.occ[WHITE])), 16)
        self.assertEqual(popcount(int(self.board.occ[BLACK])), 16)


    def test_castling_flags_parsing(self):
        fen = "r3k2r/8/8/8/8/8/8/R3K2R b KQkq - 5 42"

        rc = self.load_fen(fen)
        self.assertEqual(self.board.side_to_move, BLACK)
        self.assertEqual(self.board.castling_rights, 0b1111)
        self.assertEqual(self.board.halfmove_clock, 5)
        self.assertEqual(self.board.fullmove_number, 42)


    def test_en_passant_targets(self):
        # EP at e6, black to move
        fen = "8/8/8/3pP3/8/8/8/8 b KQkq e6 0 1"
        rc = self.load_fen(fen)
        self.assertEqual(rc, 0)
        self.assertEqual(self.board.side_to_move, BLACK)
        self.assertEqual(self.board.ep_square, sq("e6"))

        # EP at d3, white to move — use a fresh Board
        fen = "8/8/8/8/3Pp3/8/8/8 w KQkq d3 12 7"
        b = Board()
        rc = self.load_fen(fen, board=b)
        self.assertEqual(rc, 0)
        self.assertEqual(b.side_to_move, WHITE)
        self.assertEqual(b.ep_square, sq("d3"))


    def test_malformed_piece_field(self):
        # rank sums wrong (9 on a rank), bad char, missing separators
        bad = [
            "9/8/8/8/8/8/8/8 w - - 0 1",
            "8/8/8/8/8/8/8/7x w - - 0 1",
            "8/8/8/8/8/8/8/8w - - 0 1",
        ]
        for fen in bad:
            r = self.load_fen(fen, board=Board())
            self.assertEqual(r, -1)


    def test_malformed_side_castling_ep(self):
        bad = [
            "8/8/8/8/8/8/8/8 x - - 0 1",     # bad side
            "8/8/8/8/8/8/8/8 w KX - 0 1",    # bad castling char
            "8/8/8/8/8/8/8/8 w - j9 0 1",    # bad EP square
            "8/8/8/8/8/8/8/8 w - e4 0 1",    # EP not rank 3/6 (your code allows any 1..8; consider tightening)
        ]
        for fen in bad:
            r = self.load_fen(fen, board=Board())
            self.assertEqual(r, -1)