"""
Enhanced MCP Server for Text Adventure Games

Features:
- play_action: Execute game commands with score delta reporting
- memory: Rich state summary with failed-action tracking
- get_map: Graph-based explored location map
- inventory: Current inventory via direct game state
- get_failed_actions: Returns actions that failed at current location
"""

import sys
import os
import logging

# ── Logging setup ──────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[
        logging.FileHandler("mcp_server.log", mode="w"),
        logging.StreamHandler(sys.stderr),
    ],
)
logger = logging.getLogger("mcp_server")

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

from fastmcp import FastMCP
from games.zork_env import TextAdventureEnv

# ── MCP server instance ────────────────────────────────────────────────────────
mcp = FastMCP("Enhanced Text Adventure Server")

INITIAL_GAME = os.environ.get("GAME", "zork1")
logger.info(f"Server starting. GAME={INITIAL_GAME}")


# ==============================================================================
# Game State Manager
# ==============================================================================

class GameState:
    """
    Manages game state with rich tracking:
    - Full action/observation history (capped)
    - Per-location failed actions (for anti-loop context)
    - Graph-based map of explored locations
    - Score delta tracking
    """

    MAX_HISTORY = 100  # Keep last N (action, observation) pairs

    def __init__(self, game: str = "zork1"):
        self.game_name = game
        logger.info(f"[GameState] Initializing game: {game}")

        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()

        # History: list of {"action": str, "observation": str, "score": int, "reward": int}
        self.history: list[dict] = []

        # Map: location -> set of "direction -> destination" strings
        self.explored_map: dict[str, set[str]] = {}

        # Per-location failed actions: location -> set of actions that changed nothing
        self.failed_actions: dict[str, set[str]] = {}

        # All-time attempted actions per location (to avoid any repetition)
        self.attempted_actions: dict[str, set[str]] = {}

        # Score tracking
        self.prev_score: int = 0
        self.current_score: int = self.state.score

        # Location tracking
        self.current_location: str = self._extract_location(self.state.observation) or "Start"
        logger.debug(f"[GameState] Initial location: {self.current_location}")

    # ── Internal helpers ───────────────────────────────────────────────────────

    def _extract_location(self, observation: str) -> str:
        """Extract room name from observation.

        Z-machine games output room descriptions as:
            [optional flavor text]\\n\\n<Room Name>\\n<Description>
        A room name appears after a blank line (or at start), is short,
        and does NOT end with sentence-ending punctuation.
        Returns empty string if no room name detected.
        """
        lines = observation.split("\n")
        prev_blank = True  # Treat start of text as "after blank"
        for line in lines:
            stripped = line.strip()
            if not stripped:
                prev_blank = True
                continue
            if prev_blank:
                if (2 < len(stripped) < 60
                    and not stripped.endswith(('.', '!', '?', ':', ',', ';', '*', '"'))
                    and not stripped.startswith(('[', '>'))
                ):
                    return stripped
            prev_blank = False
        return ""

    def _is_movement(self, action: str) -> bool:
        exact_moves = {
            "north", "south", "east", "west", "up", "down",
            "n", "s", "e", "w", "u", "d",
            "northeast", "northwest", "southeast", "southwest",
            "ne", "nw", "se", "sw",
            "enter", "exit", "out", "in",
            "go north", "go south", "go east", "go west",
            "go up", "go down",
        }
        action_lower = action.strip().lower()
        if action_lower in exact_moves:
            return True
        first_word = action_lower.split()[0] if action_lower else ""
        return first_word in {"go", "walk", "run", "climb", "enter", "exit"}

    def _is_failed_observation(self, observation: str) -> bool:
        """
        Heuristic: detect when the game didn't accept the action
        (no state change, parser confusion, or negative feedback).
        """
        fail_phrases = [
            "i don't understand",
            "i don't know the word",
            "that's not a verb",
            "you can't go that way",
            "there is no",
            "nothing happens",
            "that doesn't work",
            "you can't",
            "it is fixed",
            "already",
            "doesn't open",
            "you don't see",
            "i beg your pardon",
            "what do you want to",
            "huh?",
            "you can't do that",
            "not see that",
            "not know how",
            "not have",
            "won't budge",
            "there's nothing",
            "that's not something",
            "not allowed",
        ]
        obs_lower = observation.lower()
        return any(p in obs_lower for p in fail_phrases)

    # ── Core action execution ──────────────────────────────────────────────────

    def take_action(self, action: str) -> tuple[str, int, bool]:
        """
        Execute an action. Returns (observation, score_delta, done).
        Also updates all tracking state.
        """
        action = action.strip()
        logger.info(f"[GameState] ACTION @ '{self.current_location}': {repr(action)}")

        self.prev_score = self.current_score
        prev_location = self.current_location

        # Track attempted actions per location
        self.attempted_actions.setdefault(self.current_location, set()).add(action.lower())

        self.state = self.env.step(action)
        observation = self.state.observation
        self.current_score = self.state.score
        score_delta = self.current_score - self.prev_score
        done = self.state.done

        logger.debug(f"[GameState] OBS: {observation[:120]!r}")
        logger.debug(f"[GameState] Score: {self.prev_score} -> {self.current_score} (Δ{score_delta}), Done={done}")

        # Extract room name from observation
        new_location = self._extract_location(observation)

        # Detect failure — only when NO room name detected
        # (presence of a room name means the game accepted the action)
        is_fail = (not new_location) and self._is_failed_observation(observation)
        if is_fail:
            self.failed_actions.setdefault(self.current_location, set()).add(action.lower())
            logger.warning(f"[GameState] FAILED action '{action}' at '{self.current_location}'")

        # Update location — only when a real room name is detected
        if new_location:
            if new_location != self.current_location:
                if self._is_movement(action):
                    self.explored_map.setdefault(self.current_location, set()).add(
                        f"{action.lower()} → {new_location}"
                    )
                logger.info(f"[GameState] LOCATION: '{self.current_location}' → '{new_location}'")
            self.current_location = new_location
        # If no room name detected, keep self.current_location unchanged

        # Update history
        entry = {
            "action": action,
            "observation": observation,
            "score": self.current_score,
            "reward": score_delta,
        }
        self.history.append(entry)
        if len(self.history) > self.MAX_HISTORY:
            self.history = self.history[-self.MAX_HISTORY :]

        return observation, score_delta, done

    # ── Tool helpers ───────────────────────────────────────────────────────────

    def get_memory(self) -> str:
        """Rich summary of game state for the agent."""
        recent = self.history[-8:]
        recent_str = ""
        for e in recent:
            obs_short = e["observation"].replace("\n", " ")[:80]
            reward_tag = f" [+{e['reward']} pts]" if e["reward"] > 0 else ""
            recent_str += f"  > {e['action']}{reward_tag} → {obs_short}...\n"

        failed_here = self.failed_actions.get(self.current_location, set())
        attempted_here = self.attempted_actions.get(self.current_location, set())
        failed_str = ", ".join(sorted(failed_here)) if failed_here else "(none)"
        attempted_str = ", ".join(sorted(attempted_here)) if attempted_here else "(none)"

        return (
            f"=== GAME STATE ===\n"
            f"Game      : {self.game_name}\n"
            f"Location  : {self.current_location}\n"
            f"Score     : {self.current_score}\n"
            f"Moves     : {self.state.moves}\n"
            f"Done      : {self.state.done}\n\n"
            f"=== RECENT ACTIONS (last 8) ===\n"
            f"{recent_str}\n"
            f"=== AT THIS LOCATION ===\n"
            f"Failed actions  : {failed_str}\n"
            f"All attempted   : {attempted_str}\n\n"
            f"=== CURRENT OBSERVATION ===\n"
            f"{self.state.observation}"
        )

    def get_map(self) -> str:
        """Text map of all explored locations and connections."""
        if not self.explored_map:
            return "Map: No connections discovered yet. Move around to build the map."

        lines = ["=== EXPLORED MAP ==="]
        for loc in sorted(self.explored_map):
            marker = " ◄ (you are here)" if loc == self.current_location else ""
            lines.append(f"\n[{loc}]{marker}")
            for conn in sorted(self.explored_map[loc]):
                lines.append(f"    {conn}")

        # Current location may not be in map yet
        if self.current_location not in self.explored_map:
            lines.append(f"\n[{self.current_location}] ◄ (you are here, unexplored exits)")

        return "\n".join(lines)

    def get_inventory(self) -> str:
        """Return inventory via direct game command."""
        logger.debug("[GameState] Fetching inventory via 'inventory' command")
        inv_state = self.env.step("inventory")
        # Restore state (inventory command in Frotz doesn't change game state but step() updates it)
        # We re-read from the response directly; the score/done remain valid
        obs = inv_state.observation
        return f"=== INVENTORY ===\n{obs}"

    def get_failed_actions(self, location: str | None = None) -> str:
        """Return failed and attempted actions for a location (default: current)."""
        loc = location or self.current_location
        failed = self.failed_actions.get(loc, set())
        attempted = self.attempted_actions.get(loc, set())
        return (
            f"Location: {loc}\n"
            f"Failed actions (parser rejected): {', '.join(sorted(failed)) or '(none)'}\n"
            f"All attempted actions: {', '.join(sorted(attempted)) or '(none)'}"
        )


# Global state
_game_state: GameState | None = None


def get_game() -> GameState:
    global _game_state
    if _game_state is None:
        _game_state = GameState(INITIAL_GAME)
        logger.info(f"[Server] Created fresh GameState for '{INITIAL_GAME}'")
    return _game_state


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

@mcp.tool()
def play_action(action: str) -> str:
    """
    Execute a game command and return the result.

    Args:
        action: The command (e.g., "north", "take lamp", "open mailbox", "examine sword")

    Returns:
        Game response + score/move context

    Valid command patterns:
    - Movement : north | south | east | west | up | down | ne | nw | se | sw
    - Objects  : take <item> | drop <item> | open <thing> | close <thing> | examine <thing>
    - Query    : look | inventory | read <thing>
    - Light    : turn on lamp | turn off lamp
    - Combat   : attack <enemy> with <weapon>
    """
    logger.info(f"[Tool:play_action] Called with action={repr(action)}")
    game = get_game()

    observation, score_delta, done = game.take_action(action)

    # Compose rich response
    delta_str = f"+{score_delta} pts! " if score_delta > 0 else (f"{score_delta} pts " if score_delta < 0 else "")
    done_str = "\n*** GAME OVER ***" if done else ""

    result = (
        f"{observation}\n"
        f"[Score: {game.current_score} | {delta_str}Moves: {game.state.moves}]{done_str}"
    )
    logger.debug(f"[Tool:play_action] Returning: {result[:120]!r}")
    return result


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

    Returns:
        Current location, score, move count, recent action history,
        failed actions at current location, and full current observation.

    Use this to orient yourself after re-entering context or when unsure
    what has already been tried.
    """
    logger.info("[Tool:memory] Called")
    result = get_game().get_memory()
    logger.debug(f"[Tool:memory] Returning {len(result)} chars")
    return result


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

    Returns:
        Map showing location names and how they connect.
        Marks your current location with ◄.

    Use before moving to avoid backtracking unnecessarily.
    """
    logger.info("[Tool:get_map] Called")
    result = get_game().get_map()
    logger.debug(f"[Tool:get_map] Returning {len(result)} chars")
    return result


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

    Returns:
        List of carried items from the game engine.
    """
    logger.info("[Tool:inventory] Called")
    result = get_game().get_inventory()
    logger.debug(f"[Tool:inventory] {result!r}")
    return result


@mcp.tool()
def get_failed_actions(location: str = "") -> str:
    """
    Get all actions that failed (parser-rejected or had no effect) at a location.

    Args:
        location: Location name to query. Empty string = current location.

    Returns:
        Set of failed actions and all attempted actions at that location.

    Use this to avoid wasting turns on commands already proven useless.
    """
    loc = location.strip() or None
    logger.info(f"[Tool:get_failed_actions] location={repr(loc)}")
    result = get_game().get_failed_actions(loc)
    logger.debug(f"[Tool:get_failed_actions] {result!r}")
    return result


# ==============================================================================
# Entry point
# ==============================================================================

if __name__ == "__main__":
    logger.info("[Server] Running MCP server via stdio")
    mcp.run()