"""
Student MCP Server for Text Adventure Games

This is your MCP server submission. Implement the tools that your agent
will use to play text adventure games.

Required tool:
    play_action(action: str) -> str
        Execute a game command and return the result.

Recommended tools:
    memory() -> str
        Return current game state, score, and recent history.
    
    inventory() -> str  
        Return the player's current inventory.
    
    get_map() -> str
        Return a map of explored locations.

Test your server with:
    fastmcp dev submission_template/mcp_server.py

Then open the MCP Inspector in your browser to test the tools interactively.
"""

import sys
import os

# 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


# =============================================================================
# Create the MCP Server
# =============================================================================

mcp = FastMCP("Student Text Adventure Server")


# =============================================================================
# Game State Management
# =============================================================================

class GameManager:
    """
    Manages the text adventure game state.
    
    Tracks:
    - Action history (for memory tool, though agent manages its own memory)
    - Current location (using Jericho API)
    - Current score and moves
    
    Note: The agent handles its own location_memory system and doesn't rely
    on server-side tracking beyond the core MCP tools.
    """
    
    def __init__(self):
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""
        # State tracking (for optional tools - agent manages its own memory)
        self.history: list[tuple[str, str]] = []              # (action, observation)
        self.current_location: str = "Unknown"
    
    def initialize(self, game: str = "zork1"):
        """Initialize or reset the game."""
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()

        # Reset tracking data
        self.history = []
        self.current_location = self._get_player_location_internal()
        return self.state.observation
    
    def _get_player_location_internal(self) -> str:
        """
        Get current player location using Jericho API.
        """
        if self.env and hasattr(self.env, 'env') and self.env.env:
            try:
                # Access Jericho's get_player_location() which returns a ZObject
                loc_obj = self.env.env.get_player_location()
                # ZObject has a .name attribute
                if hasattr(loc_obj, 'name'):
                    return loc_obj.name
            except Exception:
                pass
        return "Unknown"

    def step(self, action: str) -> str:
        """Execute an action and return the result."""
        if self.env is None:
            self.initialize()

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

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

        # Update current location using Jericho API
        self.current_location = self._get_player_location_internal()
        
        return obs

    
    def get_score(self) -> int:
        """Get current score."""
        return self.state.score if self.state else 0
    
    def get_moves(self) -> int:
        """Get number of moves taken."""
        return self.state.moves if self.state else 0


# Global game manager
_game = GameManager()

# =============================================================================
# MCP Tools - IMPLEMENT THESE
# =============================================================================

@mcp.tool()
def play_action(action: str) -> str:
    """
    Execute a game command and return the result.
    
    This is the main tool for interacting with the game.
    
    Args:
        action: The command to execute (e.g., "north", "take lamp", "open mailbox")
        
    Returns:
        The game's response to the action
        
    Valid commands include:
        - Movement: north, south, east, west, up, down, enter, exit
        - Objects: take <item>, drop <item>, open <thing>, examine <thing>
        - Other: look, inventory, read <thing>, turn on lamp
    """
    game = get_game()
    
    obs = game.step(action)

    # Append score/moves info for the agent
    score_info = f"\n\n[Score: {game.get_score()} | Moves: {game.get_moves()}]"

    # If the environment exposes reward/done, include them (optional but helpful)
    try:
        if getattr(game.state, "reward", 0) and game.state.reward > 0:
            score_info = f"\n\n+{game.state.reward} points! (Total: {game.get_score()})"
    except Exception:
        pass

    done_info = ""
    try:
        if getattr(game.state, "done", False):
            done_info = "\n\nGAME OVER"
    except Exception:
        pass

    return obs + score_info + done_info


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

    Returns:
        A summary including current location, score, moves, and recent history
    """
    game = get_game()

    recent = game.history[-5:] if game.history else []
    if recent:
        recent_str = "\n".join([f"  > {a} -> {obs[:80]}..." for a, obs in recent])
    else:
        recent_str = "  (none yet)"

    return (
        "Current State:\n"
        f"- Location: {game.current_location}\n"
        f"- Score: {game.get_score()} points\n"
        f"- Moves: {game.get_moves()}\n"
        f"- Game: {game.game_name}\n\n"
        "Recent Actions:\n"
        f"{recent_str}\n\n"
        "Current Observation:\n"
        f"{game.state.observation if game.state else ''}"
    )

@mcp.tool()
def get_map() -> str:
    """
    Get a map of explored locations.
    
    Note: This tool is not used by the current agent implementation.
    The agent manages its own location memory internally.

    Returns:
        A message indicating the agent doesn't use this tool
    """
    game = get_game()
    return f"Current location: {game.current_location}\\n\\nNote: The agent manages location tracking internally via its location_memory system."

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

    items = []
    try:
        if getattr(game.state, "inventory", None):
            items = game.state.inventory
    except Exception:
        items = []

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

    # Convert items to readable names
    item_names = []
    for item in items:
        s = str(item)
        s_lower = s.lower()
        if "parent" in s_lower:
            idx = s_lower.index("parent")
            name = s[:idx].strip()
            if ":" in name:
                name = name.split(":", 1)[1].strip()
            item_names.append(name)
        elif ":" in s:
            item_names.append(s.split(":", 1)[1].strip())
        else:
            item_names.append(s)

    return "Inventory: " + ", ".join(item_names)

def get_game() -> GameManager:
    """Get or initialize the game manager."""
    global _game

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

    if _game.env is None:
        _game.initialize(game_name)
    elif _game.game_name != game_name:
        _game.initialize(game_name)

    return _game


@mcp.tool()
def get_player_location() -> str:
    """
    Get the current player location name from Jericho's location tracking.

    Returns:
        The name of the current location (e.g., "West of House", "Forest")
    """
    game = get_game()

    if game.env and hasattr(game.env, 'env') and game.env.env:
        try:
            # Access Jericho's get_player_location() which returns a ZObject
            loc_obj = game.env.env.get_player_location()

            # ZObject has a .name attribute
            if hasattr(loc_obj, 'name'):
                return loc_obj.name
        except Exception as e:
            pass

    # Fallback: use heuristic location extraction
    return game.current_location


@mcp.tool()
def get_valid_actions() -> str:
    """
    Get valid actions from Jericho's action space at the current location.

    Returns:
        JSON string with valid actions: {"available": true, "actions": [...], "count": N}
    """
    game = get_game()

    if game.env and hasattr(game.env, 'env') and game.env.env:
        try:
            # CRITICAL: use_parallel=False to prevent deadlock on Lost Pig
            valid_actions = game.env.env.get_valid_actions(
                use_object_tree=True,
                use_ctypes=True,
                use_parallel=False  # Prevent multiprocessing deadlock
            )

            import json
            return json.dumps({
                "available": True,
                "actions": valid_actions,
                "count": len(valid_actions),
                "source": "jericho"
            })
        except Exception as e:
            import json
            return json.dumps({
                "available": False,
                "error": str(e),
                "actions": [],
                "count": 0
            })

    import json
    return json.dumps({
        "available": False,
        "error": "Game environment not initialized",
        "actions": [],
        "count": 0
    })


# =============================================================================
# Run the server
# =============================================================================

if __name__ == "__main__":
    # This runs the server with stdio transport (for MCP clients)
    mcp.run()