mcp_server.py

"""
MCP Server for Text Adventure Games

Enhanced server with exploration tracking, per-location failed action memory,
Jericho valid actions, and a full exploration graph.
"""

import re
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from fastmcp import FastMCP
from games.zork_env import TextAdventureEnv

INITIAL_GAME = os.environ.get("GAME", "zork1")

mcp = FastMCP("Text Adventure Server")

DIRECTIONS = {
    "north", "south", "east", "west", "up", "down",
    "enter", "exit", "n", "s", "e", "w", "u", "d",
    "ne", "nw", "se", "sw",
    "northeast", "northwest", "southeast", "southwest",
}

BASIC_DIRS = {"north", "south", "east", "west", "up", "down"}

REVERSE_DIR = {
    "north": "south", "south": "north",
    "east": "west", "west": "east",
    "up": "down", "down": "up",
    "n": "s", "s": "n", "e": "w", "w": "e", "u": "d", "d": "u",
    "ne": "sw", "sw": "ne", "nw": "se", "se": "nw",
    "northeast": "southwest", "southwest": "northeast",
    "northwest": "southeast", "southeast": "northwest",
    "enter": "exit", "exit": "enter",
}

FAIL_PATTERNS = [
    "you can't", "that's not", "i don't understand", "nothing happens",
    "you don't see", "that doesn't", "impossible", "you aren't",
    "you see nothing", "what do you want", "i beg your pardon",
    "you already", "that's not something", "huh?", "beg your pardon",
    "isn't something", "don't know", "can't go", "not a verb",
    "doesn't seem", "no effect", "not want", "grunk not",
]

# Patterns that tell us which exits ARE valid (so we can mark everything else as failed)
# e.g. "There only doorway to east and west" -> east, west are valid, rest are failed
EXIT_HINT_PATTERNS = [
    r'(?:there (?:only )?doorway to|only (?:way|exit|doorway) (?:is |to )?|only see (?:way to go )?(?:back )?)([\w\s,and]+?)(?:\.|$|from)',
    r'tunnel go ([\w\s,and]+?)(?:\.|$|,\s*and)',
    r'only see one place.*?that (?:tunnel |doorway |way )?(?:back |is )?([\w\s]+?)(?:\.|$)',
]

ALL_DIR_WORDS = {
    "north", "south", "east", "west", "up", "down",
    "northeast", "northwest", "southeast", "southwest",
    "enter", "exit",
}


def _parse_exit_hints(observation: str) -> set[str] | None:
    """Parse observation for explicit exit listings. Returns set of valid dirs, or None if no hint found."""
    obs_lower = observation.lower()
    all_found = set()
    for pattern in EXIT_HINT_PATTERNS:
        for m in re.finditer(pattern, obs_lower):
            text = m.group(1)
            for d in sorted(ALL_DIR_WORDS, key=len, reverse=True):
                if re.search(r'\b' + d + r'\b', text):
                    all_found.add(d)
    return all_found if all_found else None


def _clean_jericho_name(raw: str) -> str:
    """Extract a clean room name from Jericho's object repr.

    Jericho returns e.g. 'Obj93: Outside Parent0 Sibling0 Child87 Attributes...'
    We want just 'Outside'.
    """
    if not raw or raw == "None":
        return ""
    # Pattern: ObjN: <name> Parent...
    m = re.match(r'Obj\d+:\s*(.+?)\s+Parent', raw)
    if m:
        return m.group(1).strip()
    # Fallback: strip Obj prefix
    if raw.startswith("Obj"):
        parts = raw.split(":", 1)
        if len(parts) > 1:
            name = parts[1].strip().split()[0]
            return name
    return raw


class GameManager:
    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, int]] = []
        self.current_location: str = self._get_location()
        self.prev_location: str = ""

        self.location_graph: dict[str, dict[str, str]] = {}
        self.visited_locations: set[str] = {self.current_location}
        self.location_descriptions: dict[str, str] = {
            self.current_location: self.state.observation[:200]
        }

        self.failed_actions: dict[str, set[str]] = {}
        self.total_steps: int = 0

    def _get_location(self) -> str:
        """Get clean location name from Jericho or observation text."""
        try:
            raw = str(self.env.env.get_player_location())
            clean = _clean_jericho_name(raw)
            if clean:
                return clean
        except Exception:
            pass
        # Fallback: first non-empty line of observation
        lines = self.state.observation.strip().split("\n")
        for line in lines:
            line = line.strip()
            if line:
                return line
        return "Unknown"

    def _is_failed(self, observation: str) -> bool:
        obs_lower = observation.lower().strip()
        return any(p in obs_lower for p in FAIL_PATTERNS)

    def _record_failed(self, location: str, action: str):
        if location not in self.failed_actions:
            self.failed_actions[location] = set()
        self.failed_actions[location].add(action.lower().strip())

    def _learn_exits_from_observation(self, observation: str, location: str, tried_dir: str):
        """When a direction fails and the game tells us which exits exist, mark all others as failed."""
        valid_exits = _parse_exit_hints(observation)
        if not valid_exits:
            return
        # Mark all basic + compound directions NOT in the valid set as failed
        known_exits = set(self.location_graph.get(location, {}).keys())
        for d in ALL_DIR_WORDS:
            if d not in valid_exits and d not in known_exits:
                self._record_failed(location, d)

    def step(self, action: str) -> str:
        prev_loc = self.current_location
        prev_score = self.state.score

        self.state = self.env.step(action)
        self.total_steps += 1

        new_loc = self._get_location()
        score_delta = self.state.score - prev_score

        self.history.append((action, self.state.observation, score_delta))
        if len(self.history) > 50:
            self.history = self.history[-50:]

        action_lower = action.lower().strip()

        if action_lower in DIRECTIONS:
            if new_loc != prev_loc:
                if prev_loc not in self.location_graph:
                    self.location_graph[prev_loc] = {}
                self.location_graph[prev_loc][action_lower] = new_loc

                # Record reverse as a hint (may not be valid for one-way paths)
                # Only record if the reverse direction hasn't been tried/failed at new_loc
                reverse = REVERSE_DIR.get(action_lower)
                if reverse:
                    failed_at_new = self.failed_actions.get(new_loc, set())
                    if reverse not in failed_at_new:
                        if new_loc not in self.location_graph:
                            self.location_graph[new_loc] = {}
                        if reverse not in self.location_graph[new_loc]:
                            self.location_graph[new_loc][reverse] = prev_loc
            else:
                self._record_failed(prev_loc, action)
                # Remove phantom graph edge if it exists
                if prev_loc in self.location_graph and action_lower in self.location_graph[prev_loc]:
                    del self.location_graph[prev_loc][action_lower]
                # Parse exit hints from failure response to bulk-mark invalid directions
                self._learn_exits_from_observation(self.state.observation, prev_loc, action_lower)
        else:
            if self._is_failed(self.state.observation):
                self._record_failed(new_loc, action)

        self.prev_location = prev_loc
        self.current_location = new_loc
        self.visited_locations.add(new_loc)

        if new_loc not in self.location_descriptions:
            self.location_descriptions[new_loc] = self.state.observation[:200]

        # Learn exits from room description when entering a new room
        if new_loc != prev_loc:
            self._learn_exits_from_observation(self.state.observation, new_loc, action_lower)

        return self.state.observation

    def get_untried_directions(self) -> list[str]:
        """Get directions not yet tried at current location."""
        known_exits = set(self.location_graph.get(self.current_location, {}).keys())
        failed_here = self.failed_actions.get(self.current_location, set())
        tried = known_exits | {a for a in failed_here if a in DIRECTIONS}
        return sorted(BASIC_DIRS - tried)

    def get_inventory_items(self) -> str:
        items = self.state.inventory if self.state.inventory else []
        if not items:
            return "Empty-handed."
        item_names = []
        for item in items:
            item_str = str(item)
            clean = _clean_jericho_name(item_str)
            if clean:
                item_names.append(clean)
            else:
                item_names.append(item_str)
        return ", ".join(item_names)


_game: GameManager | None = None


def get_game() -> GameManager:
    global _game
    if _game is None:
        _game = GameManager(INITIAL_GAME)
    return _game


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

@mcp.tool()
def play_action(action: str) -> str:
    """Execute a game command (movement, take/drop/open/examine, look, etc.)."""
    game = get_game()
    prev_score = game.state.score
    prev_loc = game.current_location

    result = game.step(action)

    meta = [f"[Score: {game.state.score} | Moves: {game.state.moves} | Location: {game.current_location}]"]

    if game.state.score > prev_score:
        meta.append(f"[+{game.state.score - prev_score} points!]")
    if game.current_location != prev_loc:
        meta.append(f"[Moved: {prev_loc} -> {game.current_location}]")
    if game.state.done:
        meta.append("[GAME OVER]")

    # Append untried directions (basic + Jericho-confirmed compound)
    untried = game.get_untried_directions()
    # Also check for compound directions and interactions via Jericho
    interactions = []
    try:
        valid = game.env.get_valid_actions()
        known_exits = game.location_graph.get(game.current_location, {})
        failed_here = game.failed_actions.get(game.current_location, set())
        for a in valid:
            a_lower = a.lower().strip()
            if a_lower in failed_here:
                continue
            words = a.split()
            if words and words[0].lower() in DIRECTIONS:
                d = words[0].lower()
                if d not in BASIC_DIRS and d not in untried and d not in known_exits and d not in failed_here:
                    untried.append(d)
            else:
                interactions.append(a)
    except BaseException:
        pass
    if untried:
        meta.append(f"[Untried exits: {', '.join(untried)}]")

    # Show known exits for backtracking
    known_exits = game.location_graph.get(game.current_location, {})
    if known_exits and not untried:
        exits_str = ", ".join(f"{d}->{dest}" for d, dest in sorted(known_exits.items()))
        meta.append(f"[Known exits: {exits_str}]")

    # Show top interactions (free info, no step cost)
    if interactions:
        meta.append(f"[Interactions: {', '.join(interactions[:8])}]")

    return result + "\n" + " ".join(meta)


@mcp.tool()
def get_valid_actions() -> str:
    """Get valid actions for the current state, excluding previously failed actions here."""
    game = get_game()

    # Get Jericho's suggested actions
    jericho_actions = []
    try:
        jericho_actions = game.env.get_valid_actions()
    except BaseException:
        pass

    failed_here = game.failed_actions.get(game.current_location, set())

    # Always include untried basic directions
    untried_dirs = game.get_untried_directions()

    # Also add any Jericho-confirmed movement directions not in basic set
    jericho_movement = []
    interaction = []
    for a in jericho_actions:
        if a.lower().strip() in failed_here:
            continue
        words = a.split()
        if words and words[0].lower() in DIRECTIONS:
            d = words[0].lower()
            if d not in BASIC_DIRS and d not in untried_dirs:
                known_exits = game.location_graph.get(game.current_location, {})
                if d not in known_exits and d not in failed_here:
                    jericho_movement.append(d)
        else:
            interaction.append(a)

    all_untried = untried_dirs + sorted(set(jericho_movement))

    parts = []
    if all_untried:
        parts.append(f"Untried directions: {', '.join(all_untried)}")

    # Known exits that lead to already-visited locations
    known_exits = game.location_graph.get(game.current_location, {})
    if known_exits:
        exits_str = ", ".join(f"{d} -> {dest}" for d, dest in sorted(known_exits.items()))
        parts.append(f"Known exits: {exits_str}")

    if interaction:
        parts.append(f"Interactions: {', '.join(interaction[:15])}")

    if not parts:
        parts.append("No known valid actions. Try: look, north, south, east, west")

    return "\n".join(parts)


@mcp.tool()
def get_status() -> str:
    """Get exploration status: location, score, unexplored exits, failed actions here."""
    game = get_game()

    known_exits = game.location_graph.get(game.current_location, {})
    failed_here = game.failed_actions.get(game.current_location, set())
    untried = game.get_untried_directions()

    lines = [
        f"Location: {game.current_location}",
        f"Score: {game.state.score}/{game.state.max_score}",
        f"Steps: {game.total_steps}",
        f"Locations visited: {len(game.visited_locations)}",
        f"Known exits: {dict(known_exits) if known_exits else 'none mapped'}",
        f"Untried directions: {', '.join(untried) if untried else 'all tried'}",
        f"Failed actions here: {', '.join(sorted(failed_here)) if failed_here else 'none'}",
    ]
    return "\n".join(lines)


@mcp.tool()
def get_map() -> str:
    """Get exploration map with visited locations, connections, and frontier."""
    game = get_game()

    if not game.visited_locations:
        return "No locations explored yet."

    lines = [f"=== MAP ({len(game.visited_locations)} locations) ==="]
    lines.append(f"Current: {game.current_location}\n")

    for loc in sorted(game.visited_locations):
        marker = ">>>" if loc == game.current_location else "   "
        exits = game.location_graph.get(loc, {})
        lines.append(f"{marker} {loc}")
        for direction, dest in sorted(exits.items()):
            v = "V" if dest in game.visited_locations else "?"
            lines.append(f"      {direction} -> {dest} [{v}]")

    frontier = set()
    for loc, exits in game.location_graph.items():
        for dest in exits.values():
            if dest not in game.visited_locations:
                frontier.add(dest)

    if frontier:
        lines.append(f"\nFrontier (unvisited): {', '.join(sorted(frontier))}")

    return "\n".join(lines)


@mcp.tool()
def inventory() -> str:
    """Check what items you are currently carrying (no step cost)."""
    game = get_game()
    return f"Inventory: {game.get_inventory_items()}"


@mcp.tool()
def get_history(n: int = 10) -> str:
    """Get the last N actions and their results."""
    game = get_game()
    recent = game.history[-n:]
    if not recent:
        return "No history yet."
    lines = []
    for action, obs, delta in recent:
        score_str = f" [+{delta}]" if delta > 0 else ""
        lines.append(f"> {action} -> {obs[:80]}{score_str}")
    return "\n".join(lines)


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