import os
import random
from typing import Dict, List, Optional

import chess
import chess.pgn
from mcp.server.fastmcp import FastMCP


mcp = FastMCP(name="ChessServer", stateless_http=True)


# Simple in-memory game state keyed by session_id
games: Dict[str, Dict] = {}


def _get_session_id() -> str:
    # Provide a default single-session id when not given by client
    return "default"


def _ensure_game(session_id: str) -> Dict:
    if session_id not in games:
        games[session_id] = {
            "board": chess.Board(),
            "pgn_game": chess.pgn.Game(),
            "node": None,
        }
        games[session_id]["node"] = games[session_id]["pgn_game"]
    return games[session_id]


@mcp.tool(description="Start a new chess game. Returns initial board state.")
def start_game(session_id: Optional[str] = None, player_color: str = "white") -> Dict:
    sid = session_id or _get_session_id()
    game = {
        "board": chess.Board(),
        "pgn_game": chess.pgn.Game(),
        "node": None,
    }
    if player_color.lower() not in ("white", "black"):
        player_color = "white"
    # If AI should play first (player is black), make an AI opening move
    if player_color.lower() == "black":
        # Let engine choose a random legal move
        ai_move = random.choice(list(game["board"].legal_moves))
        game["board"].push(ai_move)
    game["node"] = game["pgn_game"]
    games[sid] = game
    return _board_state(game["board"]) | {"session_id": sid, "player_color": player_color}


def _board_state(board: chess.Board) -> Dict:
    return {
        "fen": board.fen(),
        "unicode": board.unicode(borders=True),
        "turn": "white" if board.turn else "black",
        "is_game_over": board.is_game_over(),
        "result": board.result() if board.is_game_over() else None,
        "legal_moves": [board.san(m) for m in board.legal_moves],
    }


@mcp.tool(description="Make a player move in SAN or UCI notation.")
def player_move(move: str, session_id: Optional[str] = None) -> Dict:
    sid = session_id or _get_session_id()
    g = _ensure_game(sid)
    board: chess.Board = g["board"]
    try:
        try:
            chess_move = board.parse_san(move)
        except ValueError:
            chess_move = chess.Move.from_uci(move)
            if chess_move not in board.legal_moves:
                raise ValueError("Illegal move")
        board.push(chess_move)
        # Update PGN
        g["node"] = g["node"].add_variation(chess_move)
        return _board_state(board) | {"last_move": board.san(chess_move)}
    except Exception as e:
        return {"error": f"Invalid move: {e}"}


@mcp.tool(description="Have the AI make a move (random legal move).")
def ai_move(session_id: Optional[str] = None) -> Dict:
    sid = session_id or _get_session_id()
    g = _ensure_game(sid)
    board: chess.Board = g["board"]
    if board.is_game_over():
        return _board_state(board)
    move = random.choice(list(board.legal_moves))
    board.push(move)
    g["node"] = g["node"].add_variation(move)
    return _board_state(board) | {"last_move": board.san(move)}


@mcp.tool(description="Return current board state.")
def board(session_id: Optional[str] = None) -> Dict:
    sid = session_id or _get_session_id()
    g = _ensure_game(sid)
    return _board_state(g["board"]) | {"session_id": sid}


@mcp.tool(description="List legal moves in SAN notation.")
def legal_moves(session_id: Optional[str] = None) -> List[str]:
    sid = session_id or _get_session_id()
    g = _ensure_game(sid)
    b: chess.Board = g["board"]
    return [b.san(m) for m in b.legal_moves]


@mcp.tool(description="Game status including check, checkmate, stalemate.")
def status(session_id: Optional[str] = None) -> Dict:
    sid = session_id or _get_session_id()
    g = _ensure_game(sid)
    b: chess.Board = g["board"]
    return {
        "turn": "white" if b.turn else "black",
        "is_check": b.is_check(),
        "is_checkmate": b.is_checkmate(),
        "is_stalemate": b.is_stalemate(),
        "is_insufficient_material": b.is_insufficient_material(),
        "is_game_over": b.is_game_over(),
        "result": b.result() if b.is_game_over() else None,
    }


@mcp.tool(description="Export game PGN.")
def pgn(session_id: Optional[str] = None) -> str:
    sid = session_id or _get_session_id()
    g = _ensure_game(sid)
    game = g["pgn_game"]
    exporter = chess.pgn.StringExporter(headers=True, variations=False, comments=False)
    return game.accept(exporter)