mcp_server.py

"""
MCP Server for Text Adventures

Exposes text adventure games via MCP tools with per-location tracking,
valid actions from Jericho, and exploration context.
"""

import sys
import os
from dataclasses import dataclass, field

# Add parent directory to path to import games module
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fastmcp import FastMCP
from games.zork_env import TextAdventureEnv, list_available_games


# Get game from environment variable (default: zork1)
INITIAL_GAME = os.environ.get("GAME", "zork1")

# Create the MCP server
mcp = FastMCP("Text Adventure Server")


@dataclass
class LocationLog:
    """Tracks what has been tried at a specific location."""
    name: str
    actions_tried: dict[str, str] = field(default_factory=dict)  # action -> short result
    valid_actions: list[str] = field(default_factory=list)
    visits: int = 0


class GameState:
    """Manages the text adventure game state and exploration data."""

    def __init__(self, game: str = "zork1"):
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()
        self.history: list[tuple[str, str]] = []

        # Per-location tracking using Jericho location IDs
        self.location_logs: dict[int, LocationLog] = {}
        self.current_loc_num: int = -1
        self.current_loc_name: str = "Unknown"
        self.steps_at_current: int = 0

        # Map: location_num -> {direction -> destination_num}
        self.map_connections: dict[int, dict[str, int]] = {}

        # Initialize location
        self._update_location()

    def _get_jericho_location(self) -> tuple[int, str]:
        """Get current location from Jericho API."""
        try:
            loc = self.env.env.get_player_location()
            return loc.num, loc.name
        except Exception:
            return -1, "Unknown"

    def _get_valid_actions(self) -> list[str]:
        """Get valid actions from Jericho."""
        try:
            return self.env.get_valid_actions()
        except Exception:
            return []

    def _update_location(self) -> bool:
        """Update current location. Returns True if location changed."""
        loc_num, loc_name = self._get_jericho_location()
        changed = loc_num != self.current_loc_num

        self.current_loc_num = loc_num
        self.current_loc_name = loc_name

        if changed:
            self.steps_at_current = 0
            # Create log for new location if needed
            if loc_num not in self.location_logs:
                valid = self._get_valid_actions()
                self.location_logs[loc_num] = LocationLog(
                    name=loc_name,
                    valid_actions=valid,
                )
            self.location_logs[loc_num].visits += 1
        else:
            self.steps_at_current += 1

        return changed

    def _summarize_result(self, result: str) -> str:
        """Short summary of action result for logging."""
        first_line = result.strip().split('\n')[0]
        return first_line[:80]

    def take_action(self, action: str) -> str:
        """Execute a game action and return enriched result."""
        prev_loc_num = self.current_loc_num

        self.state = self.env.step(action)
        result = self.state.observation

        # Track history
        self.history.append((action, result))
        if len(self.history) > 50:
            self.history = self.history[-50:]

        # Log action at previous location
        if prev_loc_num in self.location_logs:
            self.location_logs[prev_loc_num].actions_tried[action] = self._summarize_result(result)

        # Update location and check if changed
        location_changed = self._update_location()

        # Track map connections for directional moves
        direction_words = {"north", "south", "east", "west", "up", "down",
                          "enter", "exit", "n", "s", "e", "w", "u", "d",
                          "northeast", "northwest", "southeast", "southwest",
                          "ne", "nw", "se", "sw"}
        if action.lower() in direction_words and location_changed:
            if prev_loc_num not in self.map_connections:
                self.map_connections[prev_loc_num] = {}
            self.map_connections[prev_loc_num][action.lower()] = self.current_loc_num

        return result

    def get_location_context(self) -> str:
        """Get context about current location for the agent."""
        log = self.location_logs.get(self.current_loc_num)
        if not log:
            return ""

        parts = []
        parts.append(f"[Location: {self.current_loc_name}]")

        # Valid actions not yet tried
        tried = set(log.actions_tried.keys())
        untried = [a for a in log.valid_actions if a not in tried]
        if untried:
            parts.append(f"[Untried valid actions: {', '.join(untried)}]")

        # Actions already tried here
        if log.actions_tried:
            tried_str = "; ".join(f"{a} -> {r}" for a, r in list(log.actions_tried.items())[-5:])
            parts.append(f"[Already tried here: {tried_str}]")

        # Exploration pressure
        if self.steps_at_current >= 3:
            parts.append(f"[WARNING: You've spent {self.steps_at_current} steps here. MOVE to a new location!]")

        return "\n".join(parts)

    def get_memory(self) -> str:
        """Get a summary of current game state."""
        recent = self.history[-5:] if self.history else []
        recent_str = "\n".join([f"  > {a} -> {r[:60]}..." for a, r in recent]) if recent else "  (none yet)"

        loc_context = self.get_location_context()

        return f"""Current State:
- Location: {self.current_loc_name}
- Score: {self.state.score} points
- Moves: {self.state.moves}
- Locations visited: {len(self.location_logs)}

{loc_context}

Recent Actions:
{recent_str}

Current Observation:
{self.state.observation}"""

    def get_map(self) -> str:
        """Get a map of explored locations with connections."""
        if not self.location_logs:
            return "Map: No locations explored yet. Try moving around!"

        lines = [f"Explored Locations ({len(self.location_logs)} total):"]
        for loc_num, log in sorted(self.location_logs.items(), key=lambda x: x[1].name):
            marker = " [YOU ARE HERE]" if loc_num == self.current_loc_num else ""
            lines.append(f"\n* {log.name}{marker}")
            # Show connections from map
            if loc_num in self.map_connections:
                for direction, dest_num in sorted(self.map_connections[loc_num].items()):
                    dest_name = self.location_logs.get(dest_num, LocationLog(name="?")).name
                    lines.append(f"    {direction} -> {dest_name}")

        return "\n".join(lines)

    def get_inventory(self) -> str:
        """Get current inventory."""
        items = self.state.inventory if hasattr(self.state, 'inventory') and self.state.inventory else []

        if not items:
            return "Inventory: You are empty-handed."

        item_names = []
        for item in items:
            item_str = str(item)
            item_lower = item_str.lower()
            if "parent" in item_lower:
                idx = item_lower.index("parent")
                name = item_str[:idx].strip()
                if ":" in name:
                    name = name.split(":", 1)[1].strip()
                item_names.append(name)
            elif ":" in item_str:
                name = item_str.split(":")[1].strip()
                item_names.append(name)
            else:
                item_names.append(item_str)

        return f"Inventory: {', '.join(item_names)}"


# Global game state
_game_state: GameState | None = None


def get_game() -> GameState:
    """Get or initialize the game state."""
    global _game_state
    if _game_state is None:
        _game_state = GameState(INITIAL_GAME)
    return _game_state


# =============================================================================
# MCP Tools
# =============================================================================

@mcp.tool()
def play_action(action: str) -> str:
    """
    Execute a game action in the text adventure.

    Args:
        action: The command to execute (e.g., 'north', 'take lamp', 'open mailbox')

    Returns:
        The game's response with location context, valid actions, and score
    """
    game = get_game()
    result = game.take_action(action)

    # Add score info
    score_info = f"\n\n[Score: {game.state.score} | Moves: {game.state.moves}]"
    if game.state.reward > 0:
        score_info = f"\n\n+{game.state.reward} points! (Total: {game.state.score})"

    # Add location context (valid actions, tried actions, exploration pressure)
    loc_context = game.get_location_context()

    done_info = ""
    if game.state.done:
        done_info = "\n\nGAME OVER"

    return result + score_info + "\n" + loc_context + done_info


@mcp.tool()
def memory() -> str:
    """
    Get a summary of the current game state.

    Returns location, score, moves, recent actions, and current observation.
    """
    return get_game().get_memory()


@mcp.tool()
def get_map() -> str:
    """
    Get a map showing explored locations and connections.

    Useful for navigation and avoiding getting lost.
    """
    return get_game().get_map()


@mcp.tool()
def inventory() -> str:
    """
    Check what items you are currently carrying.
    """
    return get_game().get_inventory()


# =============================================================================
# Main
# =============================================================================

if __name__ == "__main__":
    mcp.run()