server.py

"""
MCP server exposing chess analysis tools (Stockfish + Maia2) for use with
Claude Desktop, Claude Code, or any MCP-compatible client.

Configuration via environment variables:
    STOCKFISH_PATH   — path to Stockfish binary (required)
    STOCKFISH_DEPTH  — search depth (default 24)
    STOCKFISH_TOP_N  — default number of top moves (default 3)
    MAIA2_TYPE       — "rapid" or "blitz" (default "rapid")
    MAIA2_DEVICE     — "cpu" or "cuda" (default "cuda")
    MAIA2_TOP_N      — default top-N predictions (default 5)
    MAIA2_SAVE_ROOT  — directory for auto-downloaded model weights (default "./models")
"""

import atexit
import io
import os
import re
import sys

import chess
import chess.engine
import chess.pgn
import chess.svg
import torch
from mcp.server.fastmcp import FastMCP

from maia2 import model as maia2_model_mod, inference as maia2_inf
from maia2.utils import map_to_category, mirror_move

# ── Configuration ──────────────────────────────────────────────────────────────

STOCKFISH_PATH = os.environ.get("STOCKFISH_PATH")
if not STOCKFISH_PATH:
    raise RuntimeError(
        "STOCKFISH_PATH environment variable is not set. "
        "Set it to the path of your Stockfish binary "
        "(e.g. export STOCKFISH_PATH=/usr/local/bin/stockfish)."
    )
STOCKFISH_DEPTH = int(os.environ.get("STOCKFISH_DEPTH", "24"))
STOCKFISH_TOP_N = int(os.environ.get("STOCKFISH_TOP_N", "3"))

MAIA2_TYPE = os.environ.get("MAIA2_TYPE", "rapid")
MAIA2_DEVICE = os.environ.get("MAIA2_DEVICE", "cpu")
MAIA2_TOP_N = int(os.environ.get("MAIA2_TOP_N", "5"))
MAIA2_SAVE_ROOT = os.environ.get("MAIA2_SAVE_ROOT", "./models")

ELO_LEVELS = [1050, 1150, 1250, 1350, 1450, 1550, 1650, 1750, 1850, 1950, 2050]

# ── Maia2 model (loaded once at startup) ───────────────────────────────────────

print("Loading Maia2 model…", file=sys.stderr)
maia2_model = maia2_model_mod.from_pretrained(
    type=MAIA2_TYPE, device=MAIA2_DEVICE, save_root=MAIA2_SAVE_ROOT
)
# Move model to the requested device if needed
if MAIA2_DEVICE == "cuda" and not next(maia2_model.parameters()).is_cuda:
    maia2_model = maia2_model.cuda()
maia2_prepared = maia2_inf.prepare()
print("Maia2 model ready.", file=sys.stderr)

# ── Stockfish engine (persistent, reused across calls) ───────────────────────

print("Starting Stockfish engine…", file=sys.stderr)
_stockfish_engine = chess.engine.SimpleEngine.popen_uci(STOCKFISH_PATH)
print("Stockfish engine ready.", file=sys.stderr)
atexit.register(_stockfish_engine.quit)

# ── MCP Server ─────────────────────────────────────────────────────────────────

mcp = FastMCP("Chess Analysis Tools")


# -- Board / move utilities --------------------------------------------------

@mcp.tool()
def is_move_legal(fen: str, move_uci: str) -> str:
    """Check if a move (UCI format) is legal in the given position.

    Args:
        fen: FEN string of the current position.
        move_uci: Move in UCI format (e.g. 'e2e4').
    """
    try:
        board = chess.Board(fen)
        return "true" if chess.Move.from_uci(move_uci) in board.legal_moves else "false"
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def get_legal_moves(fen: str) -> str:
    """List all legal moves grouped by piece type, showing both UCI and SAN notation.

    Args:
        fen: FEN string of the current position.
    """
    try:
        board = chess.Board(fen)
        moves = list(board.legal_moves)
        if not moves:
            return "No legal moves (checkmate or stalemate)"
        # Group moves by piece type
        piece_names = {
            chess.PAWN: "Pawns", chess.KNIGHT: "Knights", chess.BISHOP: "Bishops",
            chess.ROOK: "Rooks", chess.QUEEN: "Queen", chess.KING: "King",
        }
        groups: dict[int, list[str]] = {}
        for m in moves:
            pt = board.piece_at(m.from_square).piece_type
            groups.setdefault(pt, []).append(f"{m.uci()} ({board.san(m)})")
        parts = []
        for pt in [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING]:
            if pt in groups:
                parts.append(f"{piece_names[pt]}: {', '.join(groups[pt])}")
        return f"{len(moves)} legal moves. " + " | ".join(parts)
    except ValueError as e:
        return f"Error: {e}"


def _parse_move(board: chess.Board, move_str: str) -> chess.Move | None:
    """Try to parse a move as UCI first, then SAN. Returns None if invalid."""
    try:
        m = chess.Move.from_uci(move_str)
        if m in board.legal_moves:
            return m
    except ValueError:
        pass
    try:
        return board.parse_san(move_str)
    except (ValueError, chess.IllegalMoveError, chess.AmbiguousMoveError):
        return None


@mcp.tool()
def make_move(fen: str, move: str) -> str:
    """Apply a move and return the resulting FEN.
    Accepts UCI (e.g. 'e2e4') or SAN (e.g. 'Nf3') notation.

    Args:
        fen: FEN string of the current position.
        move: Move in UCI or SAN format.
    """
    try:
        board = chess.Board(fen)
        m = _parse_move(board, move)
        if m is None:
            return f"Invalid or illegal move: {move}"
        board.push(m)
        return board.fen()
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def make_moves(fen: str, moves: str) -> str:
    """Apply a sequence of moves and return the resulting FEN.
    Accepts UCI (e.g. 'e2e4,e7e5,g1f3') or SAN (e.g. 'e4,e5,Nf3') — formats can be mixed.
    Stops at the first illegal move and reports an error.

    Args:
        fen: FEN string of the starting position.
        moves: Comma-separated moves in UCI or SAN format.
    """
    try:
        board = chess.Board(fen)
        move_list = [m.strip() for m in moves.split(",") if m.strip()]
        for i, move_str in enumerate(move_list, 1):
            m = _parse_move(board, move_str)
            if m is None:
                return f"Error: move {i} '{move_str}' is invalid or illegal"
            board.push(m)
        return board.fen()
    except ValueError as e:
        return f"Error: {e}"


# -- Notation conversion -----------------------------------------------------

@mcp.tool()
def uci_to_san(fen: str, move_uci: str) -> str:
    """Convert a single UCI move to SAN notation.

    Args:
        fen: FEN string of the current position.
        move_uci: Move in UCI format (e.g. 'e2e4').
    """
    try:
        board = chess.Board(fen)
        return board.san(chess.Move.from_uci(move_uci))
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def uci_to_san_batch(fen: str, moves_uci: str) -> str:
    """Convert multiple UCI moves to SAN in one call (comma-separated).

    Args:
        fen: FEN string of the current position.
        moves_uci: Comma-separated UCI moves (e.g. 'e2e4,g1f3,d2d4').
    """
    try:
        board = chess.Board(fen)
        results = []
        for uci in moves_uci.split(","):
            uci = uci.strip()
            results.append(f"{uci}={board.san(chess.Move.from_uci(uci))}")
        return ", ".join(results)
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def san_to_uci(fen: str, move_san: str) -> str:
    """Convert a single SAN move to UCI format.

    Args:
        fen: FEN string of the current position.
        move_san: Move in SAN notation (e.g. 'e4', 'Nf3', 'O-O').
    """
    try:
        board = chess.Board(fen)
        return board.parse_san(move_san).uci()
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def san_to_uci_batch(fen: str, moves_san: str) -> str:
    """Convert multiple SAN moves to UCI in one call (comma-separated).

    Args:
        fen: FEN string of the current position.
        moves_san: Comma-separated SAN moves (e.g. 'e4,Nf3,O-O').
    """
    try:
        board = chess.Board(fen)
        results = []
        for san in moves_san.split(","):
            san = san.strip()
            results.append(f"{san}={board.parse_san(san).uci()}")
        return ", ".join(results)
    except ValueError as e:
        return f"Error: {e}"


# -- PGN / FEN ---------------------------------------------------------------

@mcp.tool()
def pgn_to_fen(pgn: str, move_number: int = 0) -> str:
    """Parse PGN and return the FEN after a given number of half-moves (plies).
    If move_number is 0, returns the final position.

    Args:
        pgn: PGN string of the game.
        move_number: Number of half-moves to play. 0 = all moves.
    """
    game = chess.pgn.read_game(io.StringIO(pgn))
    if game is None:
        return "Error: Invalid PGN"
    board = game.board()
    moves = list(game.mainline_moves())
    limit = len(moves) if move_number <= 0 else min(move_number, len(moves))
    for m in moves[:limit]:
        board.push(m)
    return board.fen()


@mcp.tool()
def get_pgn_moves(raw_pgn: str) -> str:
    """Extract clean move list from raw PGN, stripping headers/comments/annotations/results.

    Args:
        raw_pgn: Raw PGN string with headers, comments, clock times, etc.
    """
    game = chess.pgn.read_game(io.StringIO(raw_pgn))
    if game is None:
        return "Error: Invalid PGN"
    exporter = chess.pgn.StringExporter(headers=False, variations=False, comments=False)
    result = game.accept(exporter)
    return re.sub(r"\s*(1-0|0-1|1/2-1/2|\*)\s*$", "", result).strip()


# -- Board inspection --------------------------------------------------------

@mcp.tool()
def get_piece_at(fen: str, square: str) -> str:
    """Get the piece on a given square.

    Args:
        fen: FEN string of the current position.
        square: Square in algebraic notation (e.g. 'e4').
    """
    try:
        board = chess.Board(fen)
        piece = board.piece_at(chess.parse_square(square))
        if piece is None:
            return "empty"
        color = "white" if piece.color else "black"
        return f"{color} {chess.piece_name(piece.piece_type)}"
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def get_game_status(fen: str) -> str:
    """Get position status: whose turn, check, checkmate, stalemate, draw conditions,
    move number, halfmove clock, and castling rights.

    Args:
        fen: FEN string of the current position.
    """
    try:
        board = chess.Board(fen)
        turn = _color_name(board.turn)

        if board.is_checkmate():
            return f"Checkmate. {_color_name(not board.turn)} wins."
        if board.is_stalemate():
            return "Stalemate. Draw."
        if board.is_insufficient_material():
            return "Insufficient material. Draw."

        parts = [f"{turn} to move."]
        if board.is_check():
            parts.append("In check.")
        if board.can_claim_fifty_moves():
            parts.append("Draw claimable (fifty-move rule).")
        parts.append(f"Move {board.fullmove_number}.")
        parts.append(f"Halfmove clock: {board.halfmove_clock}.")
        parts.append(f"Castling: {_castling_str(board, chess.WHITE)} vs {_castling_str(board, chess.BLACK)}.")

        return " ".join(parts)
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def board_to_ascii(fen: str) -> str:
    """Render a position as an ASCII board diagram.

    Args:
        fen: FEN string of the position to display.
    """
    try:
        return str(chess.Board(fen))
    except ValueError as e:
        return f"Error: {e}"


@mcp.tool()
def board_to_svg(
    fen: str,
    arrows: str = "",
    highlights: str = "",
    orientation: str = "white",
    size: int = 400,
) -> str:
    """Render a position as an SVG board diagram with optional arrows and highlights.

    Args:
        fen: FEN string of the position to display.
        arrows: Comma-separated arrows as 'from-to' (e.g. 'e2-e4,red:g1-f3'). Optional color prefix.
        highlights: Comma-separated squares to highlight (e.g. 'e4,red:d5'). Optional color prefix.
        orientation: 'white' or 'black' — which side at bottom.
        size: SVG size in pixels (default 400).
    """
    try:
        board = chess.Board(fen)
        orient = chess.WHITE if orientation.lower() == "white" else chess.BLACK

        # Parse arrows
        arrow_list = []
        if arrows.strip():
            for a in arrows.split(","):
                color, a = _parse_svg_color(a.strip(), "#15781B")
                parts = a.split("-")
                if len(parts) == 2:
                    tail = chess.parse_square(parts[0].strip())
                    head = chess.parse_square(parts[1].strip())
                    arrow_list.append(chess.svg.Arrow(tail, head, color=color))

        # Parse highlights
        fill_dict = {}
        if highlights.strip():
            for h in highlights.split(","):
                color, sq_str = _parse_svg_color(h.strip(), "#FF6600")
                fill_dict[chess.parse_square(sq_str.strip())] = color

        svg = chess.svg.board(
            board, orientation=orient, arrows=arrow_list, fill=fill_dict, size=size
        )
        return svg
    except Exception as e:
        return f"Error: {e}"


# -- Position analysis -------------------------------------------------------

_PIECE_VALUES = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3,
                 chess.ROOK: 5, chess.QUEEN: 9, chess.KING: 0}
_PIECE_VALUES_TACTICS = {**_PIECE_VALUES, chess.KING: 100}
_PIECE_SYMBOLS = {chess.PAWN: "P", chess.KNIGHT: "N", chess.BISHOP: "B",
                  chess.ROOK: "R", chess.QUEEN: "Q", chess.KING: "K"}
_CENTER_SQUARES = [chess.E4, chess.D4, chess.E5, chess.D5]
_SVG_COLOR_MAP = {"red": "#FF0000", "blue": "#003088", "yellow": "#E6D100", "green": "#15781B"}


def _color_name(color: chess.Color) -> str:
    return "White" if color == chess.WHITE else "Black"


def _castling_str(board: chess.Board, color: chess.Color) -> str:
    parts = []
    if board.has_kingside_castling_rights(color):
        parts.append("K" if color == chess.WHITE else "k")
    if board.has_queenside_castling_rights(color):
        parts.append("Q" if color == chess.WHITE else "q")
    return "".join(parts) or "none"


def _parse_svg_color(token: str, default: str) -> tuple[str, str]:
    """Parse optional 'color:value' prefix. Returns (color_hex, value)."""
    if ":" in token:
        c, val = token.split(":", 1)
        return _SVG_COLOR_MAP.get(c.lower(), c), val
    return default, token


def _eval_material(board: chess.Board) -> str:
    mat = {chess.WHITE: 0, chess.BLACK: 0}
    pieces_str = {chess.WHITE: [], chess.BLACK: []}
    for color in [chess.WHITE, chess.BLACK]:
        for pt in [chess.QUEEN, chess.ROOK, chess.BISHOP, chess.KNIGHT, chess.PAWN]:
            count = len(board.pieces(pt, color))
            mat[color] += count * _PIECE_VALUES[pt]
            if count > 0:
                sym = f"{count}{_PIECE_SYMBOLS[pt]}" if count > 1 else _PIECE_SYMBOLS[pt]
                pieces_str[color].append(sym)
    diff = mat[chess.WHITE] - mat[chess.BLACK]
    balance = f"White +{diff}" if diff > 0 else f"Black +{-diff}" if diff < 0 else "Equal"
    result = (
        f"Material: White {mat[chess.WHITE]} ({' '.join(pieces_str[chess.WHITE])}) "
        f"vs Black {mat[chess.BLACK]} ({' '.join(pieces_str[chess.BLACK])}). {balance}."
    )
    # Bishop pair
    bp = []
    for color in [chess.WHITE, chess.BLACK]:
        if len(board.pieces(chess.BISHOP, color)) >= 2:
            bp.append(f"{_color_name(color)} has bishop pair")
    if bp:
        result += f"\nBishop pair: {', '.join(bp)}."
    return result


def _eval_pawn_structure(board: chess.Board) -> str:
    pawn_info = {chess.WHITE: [], chess.BLACK: []}
    for color in [chess.WHITE, chess.BLACK]:
        pawns = board.pieces(chess.PAWN, color)
        files_with_pawns = {chess.square_file(sq) for sq in pawns}

        for f in range(8):
            if len(pawns & chess.SquareSet(chess.BB_FILES[f])) >= 2:
                pawn_info[color].append(f"doubled on {chr(ord('a') + f)}-file")

        for f in files_with_pawns:
            adj = {f - 1, f + 1} & set(range(8))
            if not adj & files_with_pawns:
                pawn_info[color].append(f"isolated on {chr(ord('a') + f)}-file")

        opp_pawns = board.pieces(chess.PAWN, not color)
        for sq in pawns:
            f, r = chess.square_file(sq), chess.square_rank(sq)
            passed = True
            for cf in [f - 1, f, f + 1]:
                if cf < 0 or cf > 7:
                    continue
                for opp_sq in opp_pawns:
                    if chess.square_file(opp_sq) == cf:
                        opp_r = chess.square_rank(opp_sq)
                        if (color == chess.WHITE and opp_r > r) or \
                           (color == chess.BLACK and opp_r < r):
                            passed = False
            if passed:
                pawn_info[color].append(f"passed pawn on {chess.square_name(sq)}")

    parts = []
    for color in [chess.WHITE, chess.BLACK]:
        if pawn_info[color]:
            parts.append(f"{_color_name(color)}: {', '.join(pawn_info[color])}")
    return f"Pawns: {'. '.join(parts)}." if parts else "Pawns: no notable features."


def _eval_mobility(board: chess.Board) -> str:
    own_moves = len(list(board.legal_moves))
    opp_board = board.copy()
    opp_board.turn = not board.turn
    opp_board.ep_square = None
    opp_moves = len(list(opp_board.legal_moves))
    side = _color_name(board.turn)
    other = _color_name(not board.turn)
    return f"Mobility: {side} {own_moves} moves, {other} ~{opp_moves} moves."


@mcp.tool()
def evaluate_position(fen: str) -> str:
    """Evaluate static positional features: material balance, pawn structure,
    center control, castling rights, and piece mobility.

    Args:
        fen: FEN string of the position to evaluate.
    """
    try:
        board = chess.Board(fen)
        w_ctrl = sum(len(board.attackers(chess.WHITE, sq)) for sq in _CENTER_SQUARES)
        b_ctrl = sum(len(board.attackers(chess.BLACK, sq)) for sq in _CENTER_SQUARES)
        sections = [
            _eval_material(board),
            _eval_pawn_structure(board),
            f"Center (e4 d4 e5 d5): White {w_ctrl} vs Black {b_ctrl}.",
            f"Castling: White {_castling_str(board, chess.WHITE)}, Black {_castling_str(board, chess.BLACK)}.",
            _eval_mobility(board),
        ]
        return "\n".join(sections)
    except Exception as e:
        return f"Error: {e}"


def _find_pins(board: chess.Board) -> str:
    pin_strs = []
    for color in [chess.WHITE, chess.BLACK]:
        opp = not color
        for sq in chess.SQUARES:
            piece = board.piece_at(sq)
            if not (piece and piece.color == color and board.is_pinned(color, sq)):
                continue
            pin_ray = board.pin(color, sq)
            for ray_sq in pin_ray:
                p = board.piece_at(ray_sq)
                if p and p.color == opp and ray_sq != sq and \
                   p.piece_type in [chess.BISHOP, chess.ROOK, chess.QUEEN]:
                    pin_strs.append(
                        f"{_color_name(color)} {chess.piece_name(piece.piece_type)} on "
                        f"{chess.square_name(sq)} pinned by "
                        f"{chess.piece_name(p.piece_type)} on {chess.square_name(ray_sq)}"
                    )
                    break
    return f"Pins: {'; '.join(pin_strs)}." if pin_strs else "Pins: none."


def _find_hanging(board: chess.Board) -> str:
    strs = []
    side, opp = board.turn, not board.turn
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if piece is None or piece.color == side or piece.piece_type == chess.KING:
            continue
        attackers = board.attackers(side, sq)
        if not attackers:
            continue
        defenders = board.attackers(opp, sq)
        opp_name = _color_name(opp)
        piece_name = chess.piece_name(piece.piece_type)
        sq_name = chess.square_name(sq)
        if not defenders:
            strs.append(f"{opp_name} {piece_name} on {sq_name} undefended")
        else:
            min_val = min(_PIECE_VALUES_TACTICS[board.piece_at(a).piece_type] for a in attackers)
            if min_val < _PIECE_VALUES_TACTICS[piece.piece_type]:
                atk_type = next(
                    board.piece_at(a).piece_type for a in attackers
                    if _PIECE_VALUES_TACTICS[board.piece_at(a).piece_type] == min_val
                )
                strs.append(f"{opp_name} {piece_name} on {sq_name} attacked by {chess.piece_name(atk_type)}")
    return f"Hanging: {'; '.join(strs)}." if strs else "Hanging: none."


def _find_checks(board: chess.Board) -> str:
    checks = [board.san(m) for m in board.legal_moves if board.gives_check(m)]
    return f"Checking moves: {', '.join(checks)}." if checks else "Checking moves: none."


def _find_forks(board: chess.Board) -> str:
    fork_strs = []
    for m in board.legal_moves:
        board.push(m)
        moved_piece = board.piece_at(m.to_square)
        if moved_piece:
            attacker_val = _PIECE_VALUES_TACTICS[moved_piece.piece_type]
            targets = []
            for t_sq in board.attacks(m.to_square):
                t_piece = board.piece_at(t_sq)
                if t_piece and t_piece.color != moved_piece.color:
                    if _PIECE_VALUES_TACTICS[t_piece.piece_type] >= attacker_val:
                        targets.append(
                            f"{chess.piece_name(t_piece.piece_type).capitalize()} "
                            f"on {chess.square_name(t_sq)}"
                        )
            if len(targets) >= 2:
                board.pop()
                fork_strs.append(f"{board.san(m)} forks {', '.join(targets)}")
                board.push(m)
        board.pop()
    return f"Forks: {'; '.join(fork_strs)}." if fork_strs else "Forks: none."


@mcp.tool()
def find_tactics(fen: str) -> str:
    """Find tactical patterns: pins, forks, checks, hanging pieces,
    and discovered attack opportunities.

    Args:
        fen: FEN string of the position to analyze.
    """
    try:
        board = chess.Board(fen)
        if board.is_check():
            checkers = board.checkers()
            check_str = "In check by: " + ", ".join(
                f"{chess.piece_name(board.piece_at(sq).piece_type)} on {chess.square_name(sq)}"
                for sq in checkers
            ) + "."
        else:
            check_str = "Not in check."
        sections = [
            check_str,
            _find_pins(board),
            _find_hanging(board),
            _find_checks(board),
            _find_forks(board),
        ]
        return "\n".join(sections)
    except Exception as e:
        return f"Error: {e}"


# -- Stockfish ----------------------------------------------------------------

def _format_pv(board: chess.Board, pv: list[chess.Move]) -> tuple[str, str]:
    """Convert a PV move list to numbered SAN string. Returns (first_move_san, full_pv_str)."""
    pv_board = board.copy()
    san_moves = []
    for mv in pv:
        san_moves.append(pv_board.san(mv))
        pv_board.push(mv)
    if not san_moves:
        return pv[0].uci() if pv else "", ""
    pv_parts = []
    start_ply = board.ply()
    for i, san in enumerate(san_moves):
        ply = start_ply + i
        move_num = ply // 2 + 1
        if ply % 2 == 0:
            pv_parts.append(f"{move_num}. {san}")
        elif i == 0:
            pv_parts.append(f"{move_num}... {san}")
        else:
            pv_parts.append(san)
    return san_moves[0], " ".join(pv_parts)


def _format_score(score: chess.engine.PovScore) -> tuple[str, str]:
    """Format engine score and WDL into strings."""
    s = score.white()
    if s.is_mate():
        score_str = f"mate in {s.mate()}"
    else:
        score_str = f"{'+' if s.score() >= 0 else ''}{s.score()} cp"
    wdl = s.wdl()
    wdl_str = f"W:{wdl.wins/10:.1f}% D:{wdl.draws/10:.1f}% L:{wdl.losses/10:.1f}%"
    return score_str, wdl_str


@mcp.tool()
def stockfish_analyse(fen: str, n: int = 0, depth: int = 0) -> str:
    """Analyse a position with Stockfish. Returns top N moves with scores, WDL
    probabilities, and full principal variation lines in SAN notation.

    Args:
        fen: FEN string of the position to analyze.
        n: Number of top moves to return. 0 = server default.
        depth: Search depth. 0 = server default.
    """
    try:
        count = n if n > 0 else STOCKFISH_TOP_N
        search_depth = depth if depth > 0 else STOCKFISH_DEPTH
        board = chess.Board(fen)
        infos = _stockfish_engine.analyse(
            board, chess.engine.Limit(depth=search_depth), multipv=count
        )
        lines = []
        for rank, info in enumerate(infos, 1):
            score_str, wdl_str = _format_score(info["score"])
            first_san, pv_str = _format_pv(board, info.get("pv", []))
            lines.append(f"{rank}. {first_san} ({score_str}, {wdl_str}) PV: {pv_str}")
        return "\n".join(lines)
    except Exception as e:
        return f"Error: {e}"


# -- Maia2 human move prediction ----------------------------------------------

def _inference_batch_elos(fen, elo_pairs):
    """Run Maia2 inference for one FEN at multiple Elo pairs in a single batched forward pass."""
    all_moves_dict, elo_dict, all_moves_dict_reversed = maia2_prepared

    boards, elos_self, elos_oppo, legal_moves_list = [], [], [], []
    for es, eo in elo_pairs:
        board_input, es_cat, eo_cat, legal_moves = maia2_inf.preprocessing(
            fen, es, eo, elo_dict, all_moves_dict
        )
        boards.append(board_input)
        elos_self.append(es_cat)
        elos_oppo.append(eo_cat)
        legal_moves_list.append(legal_moves)

    device = next(maia2_model.parameters()).device
    boards_t = torch.stack(boards).to(device)
    elos_self_t = torch.tensor(elos_self).to(device)
    elos_oppo_t = torch.tensor(elos_oppo).to(device)
    legal_moves_t = torch.stack(legal_moves_list).to(device)

    maia2_model.eval()
    with torch.no_grad():
        logits_maia, _, logits_value = maia2_model(boards_t, elos_self_t, elos_oppo_t)
        logits_maia_legal = logits_maia * legal_moves_t
        probs = logits_maia_legal.softmax(dim=-1).cpu().tolist()
        logits_value = (logits_value / 2 + 0.5).clamp(0, 1).cpu().tolist()

    black_flag = fen.split(" ")[1] == "b"
    results = []
    for i in range(len(elo_pairs)):
        win_prob = round(1 - logits_value[i] if black_flag else logits_value[i], 4)
        move_probs = {}
        legal_idxs = legal_moves_t[i].nonzero().flatten().cpu().numpy().tolist()
        for idx in legal_idxs:
            move = all_moves_dict_reversed[idx]
            if black_flag:
                move = mirror_move(move)
            move_probs[move] = round(probs[i][idx], 4)
        move_probs = dict(sorted(move_probs.items(), key=lambda x: x[1], reverse=True))
        results.append((move_probs, win_prob))
    return results


@mcp.tool()
def predict_human_move(fen: str, elo_self: int, elo_oppo: int, n: int = 0) -> str:
    """Predict what move a human of a given Elo would play. Returns top N
    predicted moves with probabilities and win probability.

    Args:
        fen: FEN string of the current position.
        elo_self: Elo rating of the player to move (e.g. 1200, 1600, 2000).
        elo_oppo: Elo rating of the opponent.
        n: Number of top predicted moves. 0 = server default.
    """
    try:
        count = n if n > 0 else MAIA2_TOP_N
        move_probs, win_prob = maia2_inf.inference_each(
            maia2_model, maia2_prepared, fen, elo_self, elo_oppo
        )
        top = sorted(move_probs.items(), key=lambda x: -x[1])[:count]
        moves_str = "; ".join(f"{m} ({p:.1%})" for m, p in top)
        return f"Predicted moves: {moves_str} | Win probability: {win_prob:.1%}"
    except Exception as e:
        return f"Error: {e}"


@mcp.tool()
def predict_human_move_all_levels(fen: str, n: int = 0) -> str:
    """Predict moves at every skill level from ~1050 to ~2050 Elo. Useful for
    comparing how different-rated players handle the same position.

    Args:
        fen: FEN string of the current position.
        n: Number of top predicted moves per level. 0 = server default.
    """
    try:
        count = n if n > 0 else MAIA2_TOP_N
        elo_pairs = [(elo, elo) for elo in ELO_LEVELS]
        results = _inference_batch_elos(fen, elo_pairs)
        lines = []
        for elo, (move_probs, win_prob) in zip(ELO_LEVELS, results):
            top = sorted(move_probs.items(), key=lambda x: -x[1])[:count]
            moves_str = ", ".join(f"{m} {p:.1%}" for m, p in top)
            lines.append(f"Elo {elo}: {moves_str} (win {win_prob:.1%})")
        return "\n".join(lines)
    except Exception as e:
        return f"Error: {e}"


# ── Entry point ─────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--transport", choices=["stdio", "sse"], default="stdio",
                        help="MCP transport: stdio (local) or sse (network)")
    parser.add_argument("--host", default="0.0.0.0", help="SSE bind address")
    parser.add_argument("--port", type=int, default=8765, help="SSE port")
    args = parser.parse_args()

    if args.transport == "sse":
        mcp.run(transport="sse", host=args.host, port=args.port)
    else:
        mcp.run()