"""
Student MCP Server for Text Adventure Games
Integrates: Automated map mapping, deep location extraction, 
Jericho valid actions, and state monitoring.
"""

import sys
import os
import re
import json
from collections import defaultdict

# 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:
    def __init__(self):
        self.env: TextAdventureEnv = None
        self.state = None
        self.game_name: str = ""
        # Core features: Map mapping and location tracking
        self.current_location = "START"
        self.explored_locations = defaultdict(dict) # {old_loc: {action: new_loc}}
    
    def initialize(self, game: str = "zork1"):
        """Resets the environment and initializes the starting location."""
        self.game_name = game
        self.env = TextAdventureEnv(game)
        self.state = self.env.reset()
        
        # Extract initial location
        self.current_location = self._extract_location(self.state.observation)
        self.explored_locations.clear()
        return self.state.observation
    
    def step(self, action: str) -> str:
        """Executes a move, updates the map topology, and tracks location changes."""
        if self.env is None:
            self.initialize()
        
        raw_action = action.strip().lower()
        old_loc = self.current_location
        
        # Execute action in game environment
        self.state = self.env.step(raw_action)
        obs = self.state.observation
        
        # Core Logic: Extract location and update automated map
        new_loc = self._extract_location(obs)
        if old_loc and new_loc and old_loc != new_loc:
            self.explored_locations[old_loc][raw_action] = new_loc
            
        self.current_location = new_loc
        return obs
    
    def _extract_location(self, obs: str) -> str:
        """Accurately extracts room names from observation text across different game styles."""
        # 1. Match [Room Name] format (Classic Zork style)
        match = re.search(r"\[([^\]]+)\]", obs)
        if match:
            return match.group(1).upper()
            
        # 2. Match short uppercase lines at the start (Standard header style)
        lines = [l.strip() for l in obs.split("\n") if l.strip()]
        for line in lines[:5]:
            if (line[0].isupper() and 
                not line.endswith(".") and 
                not line.endswith("!") and 
                len(line) < 40):
                return line.upper()
        
        return self.current_location # Fallback to current location if extraction fails

    def get_score(self) -> int:
        return self.state.score if self.state else 0
    
    def get_moves(self) -> int:
        return self.state.moves if self.state else 0

# Global game manager instance
_game = GameManager()

def get_game() -> GameManager:
    """Singleton-style accessor for the global game manager."""
    global _game
    if _game.env is None:
        game = os.environ.get("GAME", "zork1")
        _game.initialize(game)
    return _game

# =============================================================================
# MCP Tools - Full Functionality Implementation
# =============================================================================

@mcp.tool()
def play_action(action: str) -> str:
    """Executes an action and returns results, automatically injecting a location tag."""
    game = get_game()
    obs = game.step(action)
    # Inject location tag to ensure the Agent always maintains spatial awareness
    return f"[{game.current_location}] {obs}"

@mcp.tool()
def memory() -> str:
    """Retrieves current state summary: Location, Score, and Move count."""
    game = get_game()
    return f"LOC: {game.current_location} | SCORE: {game.get_score()} | MOVES: {game.get_moves()}"

@mcp.tool()
def inventory() -> str:
    """Queries the current items in the player's inventory."""
    game = get_game()
    return game.step("inventory")

@mcp.tool()
def get_map() -> str:
    """Returns the JSON topology of the explored map for BFS pathfinding."""
    game = get_game()
    if not game.explored_locations:
        return "Map: Empty"
    return json.dumps(game.explored_locations)

@mcp.tool()
def get_valid_actions() -> str:
    """Fetches currently valid action candidates from the Jericho engine."""
    game = get_game()
    # Access the underlying Jericho interface via the TextAdventureEnv wrapper
    if game.env and hasattr(game.env, 'env'):
        valid = game.env.env.get_valid_actions()
        return json.dumps(valid[:50])
    return "[]"

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