mcp_server.py

"""
MCP (Model Context Protocol) Server for RTS Game
Exposes game state and actions to AI agents via the Model Context Protocol
"""

import asyncio
import json
from typing import Any, Dict, List, Optional
from mcp.server import FastMCP
from mcp.types import ToolAnnotations

# Import game components
from app import manager  # ConnectionManager instance from app.py


class RTSGameMCP:
    """MCP server for the RTS game that exposes game state and actions."""
    
    def __init__(self):
        """Initialize the MCP server with tools and resources."""
        self.mcp = FastMCP(
            name="RTS Commander MCP Server",
            instructions="""You are an assistant for an RTS game. You can access game state and perform actions.
            
            Available commands:
            - get_game_state(): Get the current state of the RTS game
            - get_ai_analysis(language): Get AI tactical analysis of the current game state
            - move_units(unit_ids, target_x, target_y): Move selected units to a target position
            - attack_unit(attacker_ids, target_id): Attack an enemy unit
            - build_building(building_type, position_x, position_y, player_id): Build a building at a specific position
            - send_game_command(command_type, **kwargs): Send a generic command to the game
            
            Game state includes:
            - Units (infantry, tanks, helicopters, harvesters)
            - Buildings (HQ, power plants, barracks, war factories, refineries, defense turrets)
            - Player resources (credits, power)
            - Map terrain (grass, ore, gems, water)
            """,
            host="0.0.0.0",
            port=8001,  # Different port from the main game server
        )
        
        # Register tools and resources
        self._register_tools()
        self._register_resources()
    
    def _register_tools(self):
        """Register tools that allow AI to interact with the game."""
        # Tool to get current game state
        @self.mcp.tool(
            name="get_game_state",
            description="Get the current state of the RTS game"
        )
        async def get_game_state() -> str:
            """Get the current game state as JSON."""
            game_state = manager.game_state.to_dict()
            return json.dumps(game_state, indent=2)
        
        # Tool to get AI analysis
        @self.mcp.tool(
            name="get_ai_analysis",
            description="Get AI tactical analysis of the current game state"
        )
        async def get_ai_analysis(language: str = "en") -> str:
            """Get AI analysis of the current game state."""
            # Get current game state
            game_state = manager.game_state.to_dict()
            
            # Get AI analysis
            analysis = manager.ai_analyzer.summarize_combat_situation(
                game_state, 
                language_code=language
            )
            
            return json.dumps(analysis, indent=2)
        
        # Tool to send command to game
        @self.mcp.tool(
            name="send_game_command",
            description="Send a command to the game"
        )
        async def send_game_command(command_type: str, **kwargs) -> str:
            """Send a command to the game.
            
            Args:
                command_type: Type of command to send (e.g., 'move_unit', 'attack_unit', 'build_building')
                **kwargs: Command-specific parameters
            """
            try:
                # Create command dictionary
                command = {"type": command_type}
                command.update(kwargs)
                
                # Handle the command through the connection manager
                await manager.handle_command(command)
                
                return f"Command '{command_type}' sent successfully"
            except Exception as e:
                return f"Error sending command: {str(e)}"
        
        # Specific tools for common actions
        @self.mcp.tool(
            name="move_units",
            description="Move selected units to a target position"
        )
        async def move_units(unit_ids: List[str], target_x: float, target_y: float) -> str:
            """Move units to a target position."""
            command = {
                "type": "move_unit",
                "unit_ids": unit_ids,
                "target": {"x": target_x, "y": target_y}
            }
            return await send_game_command(command["type"], **{k: v for k, v in command.items() if k != "type"})
        
        @self.mcp.tool(
            name="attack_unit",
            description="Attack an enemy unit"
        )
        async def attack_unit(attacker_ids: List[str], target_id: str) -> str:
            """Attack an enemy unit."""
            command = {
                "type": "attack_unit",
                "attacker_ids": attacker_ids,
                "target_id": target_id
            }
            return await send_game_command(command["type"], **{k: v for k, v in command.items() if k != "type"})
        
        @self.mcp.tool(
            name="build_building",
            description="Build a building at a specific position"
        )
        async def build_building(building_type: str, position_x: float, position_y: float, player_id: int = 0) -> str:
            """Build a building at a specific position."""
            command = {
                "type": "build_building",
                "building_type": building_type,
                "position": {"x": position_x, "y": position_y},
                "player_id": player_id
            }
            return await send_game_command(command["type"], **{k: v for k, v in command.items() if k != "type"})
    
    def _register_resources(self):
        """Register resources that provide information about the game."""
        # Resource for game documentation
        @self.mcp.resource(
            name="game_documentation",
            description="Documentation for the RTS game",
            uri="file://docs/README.md"
        )
        async def game_documentation() -> str:
            """Get game documentation."""
            try:
                with open("docs/README.md", "r") as f:
                    return f.read()
            except FileNotFoundError:
                return "Game documentation not found"
        
        # Resource for game rules
        @self.mcp.resource(
            name="game_rules",
            description="Rules and mechanics of the RTS game",
            uri="file://docs/ARCHITECTURE.md"
        )
        async def game_rules() -> str:
            """Get game rules and mechanics."""
            try:
                with open("docs/ARCHITECTURE.md", "r") as f:
                    return f.read()
            except FileNotFoundError:
                return "Game rules documentation not found"
    
    async def run(self):
        """Start the MCP server."""
        await self.mcp.run()


# Create and run the MCP server
async def main():
    """Main entry point for the MCP server."""
    server = RTSGameMCP()
    await server.run()


if __name__ == "__main__":
    asyncio.run(main())