"""
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
import re

# 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")


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

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

class GameManager:
    """
    Manages the text adventure game state.
    
    TODO: Extend this class to track:
    - Action history (for memory tool)
    - Explored locations (for mapping)
    - Current score and moves
    """
    
    def __init__(self):
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""
        # TODO: Add more state tracking
        # self.history: list[tuple[str, str]] = []
        # self.explored_locations: dict[str, set[str]] = {}
        # self.current_location: str = ""

        self.history: list[tuple[str, str]] = []
        self.explored_locations: dict[str, set[str]] = {}
        self.exploration_count = {}
        self.current_location: str = ""
    
    def initialize(self, game: str = "zork1"):
        """Initialize or reset the game."""
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()
        # TODO: Reset your state tracking here
        self.history: list[tuple[str, str]] = []
        self.explored_locations: dict[str, set[str]] = {}
        self.exploration_count = {}
        self.current_location: str = ""
        return self.state.observation
    
    def _extract_location(self, observation: str, max_length: int = 25) -> str:
        """
        Extracts the location by finding the first line that:
        1. Is shorter than max_length.
        2. Contains ONLY alphanumeric characters and spaces (no punctuation).
        """
        lines = observation.strip().split('\n')

        for line in lines:
            cleaned_line = line.strip()
            
            if not cleaned_line:
                continue
                
            if len(cleaned_line) >= max_length:
                continue
                
            if re.match(r'^[a-zA-Z0-9 ]+$', cleaned_line):
                return cleaned_line
                
        return "Blocked"
    
    def step(self, action: str) -> str:
        """Execute an action and return the result."""
        self.state = self.env.step(action)
        result = self.state.observation
        
        action = action.strip("go ")
        # Track history
        self.history.append((action, result))
        if len(self.history) > 50:
            self.history = self.history[-50:]
        
        if self.current_location == "":
            new_location = self._extract_location(result)
            self.current_location = new_location
            self.explored_locations[new_location] = set()
            self.exploration_count[new_location] = 1

        # Update map
        elif action in [
                "north", "south", "east", "west", "up", "down", 
                "enter", "exit", "n", "s", "e", "w", "u", "d", "northeast", "northwest",
                "southeast", "southwest", "ne", "nw", "se", "sw"
            ]:
            new_location = self._extract_location(result)
            self.exploration_count[new_location] = self.exploration_count.get(new_location, 0) + 1

            if new_location != self.current_location:
                self.explored_locations[self.current_location].add(f"{action} -> {new_location}")
                if new_location != "Blocked":
                    if new_location not in self.explored_locations:
                        self.explored_locations[new_location] = set()
                    self.explored_locations[new_location].add(f"{REVERSE_ACTIONS[action]} -> {self.current_location}")
                    self.current_location = new_location
                    
        
        return result
    

    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}..." for a, r in recent]) if recent else "  (none yet)"
        
        return f"""Current State:
        - Location: {self.current_location}
        - Score: {self.state.score} points
        - Moves: {self.state.moves}
        - Game: {self.game_name}

        Recent Actions:
        {recent_str}

        Current Observation:
        {self.state.observation}"""
    
    
    def get_map(self) -> str:
        """Get a map of explored locations."""
        if not self.explored_locations:
            return "Map: No locations explored yet. Try moving around!"
        
        lines = ["Explored Locations and Exits:"]
        for loc, exits in sorted(self.explored_locations.items()):
            lines.append(f"\n* {loc} | Visit count : {self.exploration_count[loc]}")
            for exit_info in sorted(exits):
                lines.append(f"    > {exit_info}")
        
        lines.append(f"\n[Current] {self.current_location}")
        return "\n".join(lines)
    
    def get_mini_map(self) -> str:
        parts = [f"CONNECTIONS FROM {self.current_location}:"]
        for exit in self.explored_locations[self.current_location]:
            parts.append(f"  > {exit}")

        return "\n".join(parts)
    
    def get_current_location(self) -> str:
        return self.current_location
    
    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)}"
    
    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()


def get_game() -> GameManager:
    """Get or initialize the game manager."""
    global _game
    if _game.env is None:
        # Get game from environment variable (set by evaluator)
        game = os.environ.get("GAME", "zork1")
        _game.initialize(game)
    return _game


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

@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 to your action
    """
    game = get_game()
    result = game.step(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})"
    
    done_info = ""
    if game.state.done:
        done_info = "\n\nGAME OVER"
    
    return result + score_info + done_info


# TODO: Implement additional tools to help your agent

@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 get_mini_map() -> str:
    """
    Get a mini map showing explored locations and connections.
    
    Useful for navigation and avoiding getting lost.
    """
    return get_game().get_mini_map()

@mcp.tool()
def get_location() -> str:
    return get_game().get_current_location()

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

@mcp.tool()
def get_valid_actions() -> str:
    """
    Get a list of likely valid actions from the current location.
    
    Returns:
        List of actions that might work here
    """
    # This is a hint: Jericho provides get_valid_actions()
    game = get_game()
    if game.env and game.env:
        valid = game.env.get_valid_actions()
        return "Valid actions: " + ", ".join(valid[:20])
    return "Could not determine valid actions"


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

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