import ctypes as C


class Board(C.Structure):
    _fields_ = [
        ("pieces",          C.c_uint64 * 12),
        ("occ",             C.c_uint64 * 3),
        ("king_square",     C.c_uint64 * 2),
        ("castling_rights", C.c_uint8),
        ("ep_square",       C.c_int),
        ("side_to_move",    C.c_int),
        ("halfmove_clock",  C.c_int),
        ("fullmove_number", C.c_int),
    ]


class Move(C.Structure):
    _fields_ = [
        ("from",  C.c_uint16),
        ("to",    C.c_uint16),
        ("piece", C.c_uint8),
        ("promo", C.c_uint8),
        ("flags", C.c_uint8),
    ]


FILES = {c:i for i,c in enumerate("abcdefgh")}
WHITE, BLACK, BOTH = 0, 1, 2
P, N, B, R, Q, K, p, n, b, r, q, k = range(12)


"""
    Register C function bindings and interfaces in a dictionary.
"""
PAWN_MOVE_SIG = ((C.POINTER(Board), C.POINTER(Move), C.POINTER(C.c_int)), None)
PIECE_MOVE_SIG = ((C.POINTER(Board), C.POINTER(Move), C.POINTER(C.c_int), C.c_bool), None)
FFI_SPEC = {
    "create_knight_attack_cache": None,
    "create_king_attack_cache": None,
    "create_pawn_attack_cache": None,
    "gen_white_pawn_quiet_pushes": PAWN_MOVE_SIG,
    "gen_black_pawn_quiet_pushes": PAWN_MOVE_SIG,
    "gen_white_pawn_push_promotions": PAWN_MOVE_SIG,
    "gen_black_pawn_push_promotions": PAWN_MOVE_SIG,
    "gen_white_pawn_captures": PAWN_MOVE_SIG,
    "gen_black_pawn_captures": PAWN_MOVE_SIG,
    "gen_white_pawn_capture_promotions": PAWN_MOVE_SIG,
    "gen_black_pawn_capture_promotions": PAWN_MOVE_SIG,
    "gen_knight_moves": PIECE_MOVE_SIG,
    "gen_bishop_moves": PIECE_MOVE_SIG,
    "gen_rook_moves": PIECE_MOVE_SIG,
    "gen_queen_moves": PIECE_MOVE_SIG,
    "gen_king_moves": PIECE_MOVE_SIG,
    "square_attacked": ((C.POINTER(Board), C.c_int, C.c_int), C.c_bool),
    "in_check": ((C.POINTER(Board), C.c_int), C.c_bool),
    "attackers_to": ((C.POINTER(Board), C.c_int, C.c_int), C.c_uint64),
    "get_legal_moves":((C.POINTER(Board), C.POINTER(Move)), C.c_int),
    "perft": ((C.POINTER(Board), C.c_int), C.c_uint64),
    "load_fen": ((C.POINTER(Board), C.c_char_p), C.c_int),
    "gen_pseudo_moves": ((C.POINTER(Board), C.POINTER(Move), C.c_bool), C.c_int),
}


def draw_bb(mask, origin=None):
    print("\n")
    lines = []
    for r in range(7, -1, -1):
        row = []
        for f in range(8):
            sqi = r * 8 + f
            bit = (mask >> sqi) & 1
            if origin is not None and sqi == origin:
                ch = 'O'
            elif bit:
                ch = 'x'
            else:
                ch = '.'
            row.append(ch)
        lines.append(f"{r+1} " + " ".join(row))
    lines.append("  " + " ".join(FILES))
    lines = "\n".join(lines)
    print(lines, "\n") 


def sq_to_coord(sq):
    return chr(ord('a') + (sq % 8)) + chr(ord('1') + (sq // 8))


def sq(name):
    f = FILES[name[0].lower()]
    r = int(name[1]) - 1
    return f + 8 * r


class ChessFFI:
    def __init__(self):
        self.lib_path = "./build/libchess.so"
        self._load_lib()
        self._load_constants()
        self._bind_functions()


    def init_attack_cache(self):
        self._c_create_knight_attack_cache()
        self._c_create_king_attack_cache()
        self._c_create_pawn_attack_cache()

    
    def load_fen(self, board, fen_string):
        return self._c_load_fen(board, fen_string.encode("ascii"))
            

    def gen_pseudo_moves(self, board, captures_only=False, cap=256):
        buf = (Move * cap)()
        n = self._c_gen_pseudo_moves(C.byref(board), buf, captures_only)
        return buf, n


    def get_legal_moves(self, board, out):
        return int(self._c_get_legal_moves(C.byref(board), out))


    def square_attacked(self, board, sq, by):
        return bool(self._c_square_attacked(C.byref(board), int(sq), int(by)))


    def in_check(self, board, side):
        return bool(self._c_in_check(C.byref(board), int(side)))


    def attackers_to(self, board, sq, by):
        return self._c_attackers_to(board, sq, by)


    def popcount(self, x):
        return x.bit_count()
    

    def _load_lib(self):
        self._lib = C.CDLL(self.lib_path)


    def _load_constants(self):
        piece_array = (C.c_uint64 * 64)

        # Account for both sets of pawn colors because
        # they have different attack patterns.
        pawn_array = piece_array * 2

        self.KNIGHT_ATTACKS = self._in_dll_or_none("KNIGHT_ATTACKS", piece_array)
        self.KING_ATTACKS = self._in_dll_or_none("KING_ATTACKS",   piece_array)
        self.PAWN_ATTACKS = self._in_dll_or_none("PAWN_ATTACKS",   pawn_array)


    def _bind_functions(self):
        for name, spec in FFI_SPEC.items():
            if spec is None:
                argtypes, restype = (), None
            else:
                argtypes, restype = spec
            fn = getattr(self._lib, name, None)
            if not fn:
                print(f"missing function {name}")

            if fn is not None:
                fn.argtypes = argtypes
                fn.restype = restype

            # Prepend each function name with _c to allow us to reuse the 
            # function names as methods on this class. This means that we
            # can use this class as interface to the C functions without
            # modifying the names too much.
            #
            # Example:
            #
            # def load_fen(self):
            #    self._c_load_fen()
            #
            # ChessFFI.load_fen()
            #
            #
            setattr(self, f"_c_{name}", fn)


    def _in_dll_or_none(self, symbol, ctype):
        try:
            return ctype.in_dll(self._lib, symbol)
        except ValueError:
            print("Constant not found.")
            return None