"""
Student Agent for Text Adventure Games

This is your submission file. Implement the StudentAgent class to play
text adventure games using the MCP server you also implement.

Your agent should:
1. Connect to the MCP server via the provided client
2. Use the ReAct pattern (Thought -> Action -> Observation)
3. Call MCP tools to interact with the game
4. Maximize the game score within the step limit

Required method:
    async def run(self, client, game, max_steps, seed, verbose) -> RunResult

The 'client' is a FastMCP Client already connected to your MCP server.
Use it to call tools like: await client.call_tool("play_action", {"action": "look"})

Tips:
- Start by looking around and understanding your environment
- Keep track of visited locations to avoid loops
- Pick up useful items (lamp, sword, etc.)
- The seed parameter should be used to set your LLM's seed for reproducibility
"""

import json
import os
import re
from dataclasses import dataclass, field
from typing import Optional

from dotenv import load_dotenv
from huggingface_hub import InferenceClient

# Load environment variables
load_dotenv()

# =============================================================================
# LLM Configuration - DO NOT MODIFY
# =============================================================================

# Model to use (fixed for fair evaluation)
LLM_MODEL = "Qwen/Qwen2.5-72B-Instruct"

# Initialize the LLM client (uses HF_TOKEN from environment)
_hf_token = os.getenv("HF_TOKEN")
if not _hf_token:
    raise ValueError("HF_TOKEN not found. Set it in your .env file.")

LLM_CLIENT = InferenceClient(token=_hf_token)


def call_llm(prompt: str, system_prompt: str, seed: int, max_tokens: int = 300) -> str:
    """
    Call the LLM with the given prompt. Use this function in your agent.
    
    Args:
        prompt: The user prompt (current game state, history, etc.)
        system_prompt: The system prompt (instructions for the agent)
        seed: Random seed for reproducibility
        max_tokens: Maximum tokens in response (default: 300)
        
    Returns:
        The LLM's response text
        
    Example:
        response = call_llm(
            prompt="You are in a forest. What do you do?",
            system_prompt=SYSTEM_PROMPT,
            seed=42,
        )
    """
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ]
    
    response = LLM_CLIENT.chat.completions.create(
        model=LLM_MODEL,
        messages=messages,
        temperature=0.0,  # Deterministic for reproducibility
        max_tokens=max_tokens,
        seed=seed,
    )
    
    return response.choices[0].message.content


@dataclass
class RunResult:
    """Result of running the agent. Do not modify this class."""
    final_score: int
    max_score: int
    moves: int
    locations_visited: set[str]
    game_completed: bool
    error: Optional[str] = None
    history: list[tuple[str, str, str]] = field(default_factory=list)


# =============================================================================
# System Prompt - Customize this for your agent
# =============================================================================

SYSTEM_PROMPT = """You are an expert text adventure game player. Maximize your score.

TOOLS (only two):
1. play_action(action) — Execute ANY game command.
2. memory() — See full context. Use sparingly, never twice in a row.

VALID GAME COMMANDS — use ONLY these patterns:
Directions: north, south, east, west, up, down, northeast, northwest, southeast, southwest, enter, exit, in, out, 
Movement: <direction>, climb <object>, descend <object>, jump, jump over <thing>, swim, climb stairs
Take/Drop: take <item>, drop <item>
Examine: examine <thing>, read <thing>, look in/under/behind <thing>, search <thing>
Open/Close: open <thing>, close <thing>, unlock <thing> with <key>
Manipulate: push <thing>, pull <thing>, move <thing>, turn <thing>, light <thing> with <source>
Social: ask <person> about <topic>, give <item> to <person>, show <item> to <person>
Senses: look, listen, smell, wait
Other: inventory

Keep commands to 2-3 words maximum.

HOW TO PLAY WELL:

1. WHEN YOU ENTER A NEW ROOM:
   a) Take any items on the ground immediately.
   b) Try each exit that the prompt tells you is UNEXPLORED (don't guess exits from room text yourself).
   c) Do NOT go back the way you came until you've tried all other exits.
   d) If something interesting happened (sound, event), react to it before moving on: listen, search, examine.

2. WHEN A ROOM IS FULLY EXPLORED (no unexplored exits left):
   Search interesting objects (search fountain, search chest, look in/under things).
   Then navigate toward a FRONTIER room — one with unexplored exits.

3. WHEN YOU ARE STUCK:
   a) Check the FRONTIER section — go to a room that still has unexplored exits.
   b) Try: listen, wait, search <thing>, look under/behind <thing>.
   c) Try different verb forms for the same idea: push/pull/move/turn, put X on Y, light X with Y.
   d) Re-read the game text for clues you missed.

4. ANSWER DISAMBIGUATION QUESTIONS DIRECTLY.
   If game asks "X or Y?" → respond with just "X" or "Y".

5. ONE EXAMINE PER OBJECT IS ENOUGH — same text every time. Use search/open/look in next.

6. READ THE GAME TEXT. Every sentence is a clue. Follow sounds, take objects, try described exits or mentioned directions.

The [Location:X] tag tells you where you are. If you try a direction and Location stays the same, it didn't work.

OUTPUT FORMAT (no markdown):
THOUGHT:  <your reasoning - what you learned, what you know about your current situation and what to do next>
TOOL: <tool_name>
ARGS: <JSON>
"""


# =============================================================================
# Student Agent - IMPLEMENT THIS CLASS
# =============================================================================

ALL_DIRECTIONS = [
    "north", "south", "east", "west", "up", "down",
    "northeast", "northwest", "southeast", "southwest",
    "enter", "exit", "in", "out",
]

DIRECTION_SET = set(ALL_DIRECTIONS) | {
    "n", "s", "e", "w", "ne", "nw", "se", "sw", "u", "d",
}

DISAMBIGUATION_PATTERNS = [
    r"which do you mean",
    r"what that mean,?\s+.+\sor\s",
    r"do you mean .+ or",
    r"did you mean",
]

REVERSE_DIR = {
    "north": "south", "south": "north",
    "east": "west", "west": "east",
    "up": "down", "down": "up",
    "northeast": "southwest", "southwest": "northeast",
    "northwest": "southeast", "southeast": "northwest",
    "enter": "exit", "exit": "enter",
    "in": "out", "out": "in",
}


class StudentAgent:
    def __init__(self):
        self.history = []
        self.current_location = "Unknown"
        self.visited_locations = set()
        self.tried_actions = {}
        self.world_map = {}
        self.inventory = set()
        self.location_notes = {}
        self.global_summary = ""
        self.step_counter = 0
        self.consecutive_same_loc = 0
        self.last_location = ""
        self.consecutive_memory = 0
        self.room_descriptions = {}
        self.exits_taken = {}
        self.exits_failed = {}
        self.arrived_via = {}
        self.just_entered_new_room = False
        self.room_exits = {}

    def _extract_location(self, observation: str) -> Optional[str]:
        m = re.search(r'\[.*?Location:\s*(.+?)\]', observation)
        if m:
            return m.group(1).strip()
        m = re.search(r'\[Moved from .+? to (.+?)\]', observation)
        if m:
            return m.group(1).strip()
        return None

    def _norm(self, action: str) -> str:
        return action.lower().strip()

    def _is_direction(self, action: str) -> bool:
        return self._norm(action) in DIRECTION_SET

    def _record_action(self, location: str, action: str, observation: str):
        key = (location, self._norm(action))
        snippet = observation[:150].replace("\n", " ")
        if key not in self.tried_actions:
            self.tried_actions[key] = []
        self.tried_actions[key].append(snippet)

    def _action_count_here(self, action: str) -> int:
        key = (self.current_location, self._norm(action))
        return len(self.tried_actions.get(key, []))

    def _get_recent_actions_at(self, location: str, n: int = 20) -> list[str]:
        results = []
        for (loc, act), obs_list in self.tried_actions.items():
            if loc == location:
                last_obs = obs_list[-1][:80]
                results.append(f'"{act}" x{len(obs_list)} → {last_obs}')
        return results[-n:]

    def _extract_exits_with_llm(self, observation: str) -> list[str]:
        """Use the LLM to extract ACTUAL exits from room description."""
        prompt = f"""From the following room description, list ONLY the directions the player can actually move/walk through.
        Only include directions that are described as exits, doorways, tunnels, paths, or passages the player can go through.
        Do NOT include directions that just describe walls, pictures, or objects.

        Room description:
        {observation}

        Reply with ONLY a comma-separated list of directions (e.g. "north, southeast, west") or "none" if no exits are mentioned.
        Use only these words: north, south, east, west, up, down, northeast, northwest, southeast, southwest, enter, exit, in, out"""

        result = call_llm(prompt, "You extract movement exits from text adventure room descriptions. Be precise.", seed=99)
        result = result.lower().strip().strip(".")

        if result == "none" or not result:
            return []

        exits = []
        for word in re.split(r'[,\s]+', result):
            word = word.strip()
            if word in DIRECTION_SET:
                exits.append(word)
        return exits

    def _get_unexplored_exits(self, location: str) -> list[str]:
        """Exits extracted by LLM minus exits already taken or failed."""
        known_exits = self.room_exits.get(location, [])
        taken = self.exits_taken.get(location, set())
        failed = self.exits_failed.get(location, set())
        return [d for d in known_exits if d not in taken and d not in failed]

    def _get_arrival_direction(self, location: str) -> Optional[str]:
        info = self.arrived_via.get(location)
        if info:
            return info[0]
        return None

    def _get_frontier(self) -> list[tuple[str, list[str]]]:
        frontier = []
        for loc in self.visited_locations:
            unexplored = self._get_unexplored_exits(loc)
            if unexplored:
                frontier.append((loc, unexplored))
        return frontier

    def _is_disambiguation(self, observation: str) -> bool:
        obs_lower = observation.lower()
        for pattern in DISAMBIGUATION_PATTERNS:
            if re.search(pattern, obs_lower):
                return True
        return False

    async def _update_inventory(self, client):
        result = await client.call_tool("inventory", {})
        if result:
            inv_text = result.content[0].text.lower()
            items = re.findall(r'\b[a-zA-Z]+\b', inv_text)
            self.inventory = set(items)

    async def _update_room_summary(self, client, observation):
        prompt = f"""Summarize this room in 2 lines. Focus on: takeable items, creatures, sounds, puzzles, interactive objects.

        TEXT:
        {observation}"""
        summary = call_llm(prompt, "You summarize game rooms precisely.", seed=42)
        self.location_notes[self.current_location] = summary.strip()

    async def _update_global_summary(self):
        text = ""
        for loc, note in self.location_notes.items():
            text += f"{loc}: {note}\n"
        prompt = f"""Summarize game progress in 4 lines.
        Focus on: rooms with unexplored exits, unsolved puzzles, key items, objectives.

        TEXT:
        {text}"""
        self.global_summary = call_llm(
            prompt, "You are an expert adventure strategist.", seed=123
        ).strip()

    async def run(self, client, game, max_steps, seed, verbose=False):
        self.__init__()

        moves = 0
        game_completed = False

        # Initial look
        result = await client.call_tool("play_action", {"action": "look"})
        observation = result.content[0].text
        moves += 1

        loc = self._extract_location(observation)
        if loc:
            self.current_location = loc
            self.visited_locations.add(loc)
            self.room_descriptions[loc] = observation
            self.exits_taken[loc] = set()
            self.exits_failed[loc] = set()
            self.room_exits[loc] = self._extract_exits_with_llm(observation)
            self.just_entered_new_room = True

        await self._update_inventory(client)
        await self._update_room_summary(client, observation)

        if verbose:
            print(f"{'=' * 20} Starting {game} {'=' * 20}")
            print(f"[OBSERVATION]: {observation}")

        for step in range(max_steps):
            if verbose:
                print(f"\n----- Step {step + 1} / {max_steps} -----")

            self.step_counter += 1

            if self.current_location == self.last_location:
                self.consecutive_same_loc += 1
            else:
                self.consecutive_same_loc = 0
            self.last_location = self.current_location

            if step > 0 and step % 12 == 0:
                await self._update_global_summary()

            prompt = self._build_prompt(observation)

            effective_seed = seed + step
            if self.consecutive_same_loc > 8:
                effective_seed += self.consecutive_same_loc * 13

            response = call_llm(prompt, SYSTEM_PROMPT, effective_seed)
            thought, tool_name, args = self._parse_response(response)

            if tool_name not in ("play_action", "memory"):
                tool_name = "play_action"
                args = {"action": "look"}

            # Avoid using memory twice in a row
            if tool_name == "memory":
                self.consecutive_memory += 1
                if self.consecutive_memory >= 2:
                    nudge_prompt = prompt + (
                        "\n\nYou already checked memory. You MUST now play_action. "
                        "Try an unexplored exit, take an item, search something, or listen."
                    )
                    response = call_llm(nudge_prompt, SYSTEM_PROMPT, effective_seed + 77)
                    thought, tool_name, args = self._parse_response(response)
                    if tool_name != "play_action":
                        tool_name = "play_action"
                        args = {"action": "listen"}
                    self.consecutive_memory = 0
            else:
                self.consecutive_memory = 0

            # Avoid repeating the same failed action at the same location more than 3 times in a row
            if tool_name == "play_action":
                action = args.get("action", "look")
                count = self._action_count_here(action)
                if count >= 3:
                    unexplored = self._get_unexplored_exits(self.current_location)
                    retry_prompt = prompt + (
                        f'\n\nYou tried "{action}" here {count} times. Choose something DIFFERENT.'
                    )
                    if unexplored:
                        retry_prompt += f' Unexplored exits: {", ".join(unexplored)}.'
                    response = call_llm(retry_prompt, SYSTEM_PROMPT, effective_seed + 53)
                    thought, tool_name, args = self._parse_response(response)
                    if tool_name not in ("play_action", "memory"):
                        tool_name = "play_action"
                        args = {"action": "listen"}

            if verbose:
                print(f"[THOUGHT]: {thought}")
                print(f"[TOOL]: {tool_name}")
                print(f"[ARGS]: {json.dumps(args)}")

            result = await client.call_tool(tool_name, args)
            observation = result.content[0].text
            self.history.append((thought, tool_name, args, observation))

            if verbose:
                print(f"[OBSERVATION]: {observation}")

            if tool_name == "play_action":
                moves += 1
                action_str = args.get("action", "look")
                action_norm = self._norm(action_str)
                old_loc = self.current_location
                self._record_action(old_loc, action_str, observation)

                new_loc = self._extract_location(observation)
                moved = new_loc is not None and new_loc != old_loc

                if moved:
                    self.current_location = new_loc
                    self.visited_locations.add(new_loc)

                    if old_loc not in self.world_map:
                        self.world_map[old_loc] = {}
                    self.world_map[old_loc][action_norm] = new_loc

                    if old_loc not in self.exits_taken:
                        self.exits_taken[old_loc] = set()
                    self.exits_taken[old_loc].add(action_norm)

                    self.arrived_via[new_loc] = (action_norm, old_loc)

                    if new_loc not in self.room_descriptions:
                        self.room_descriptions[new_loc] = observation
                        self.exits_taken[new_loc] = set()
                        self.exits_failed[new_loc] = set()
                        self.room_exits[new_loc] = self._extract_exits_with_llm(observation)
                        self.just_entered_new_room = True
                        await self._update_room_summary(client, observation)
                    else:
                        self.just_entered_new_room = False

                    reverse = REVERSE_DIR.get(action_norm)
                    if reverse:
                        if new_loc not in self.world_map:
                            self.world_map[new_loc] = {}
                        self.world_map[new_loc][reverse] = old_loc
                        if new_loc not in self.exits_taken:
                            self.exits_taken[new_loc] = set()
                        self.exits_taken[new_loc].add(reverse)

                else:
                    self.just_entered_new_room = False
                    if self._is_direction(action_str):
                        if old_loc not in self.exits_failed:
                            self.exits_failed[old_loc] = set()
                        self.exits_failed[old_loc].add(action_norm)

                await self._update_inventory(client)

            obs_lower = observation.lower()
            if any(w in obs_lower for w in ["you have won", "victory", "you win"]):
                game_completed = True
                break
            if any(w in obs_lower for w in ["you have died", "game over"]):
                break

        score_result = await client.call_tool("get_score", {})
        final_score = int(score_result.content[0].text) if score_result else 0

        return RunResult(
            final_score=final_score,
            max_score=350,
            moves=moves,
            locations_visited=self.visited_locations,
            game_completed=game_completed,
            history=[(th, t, o) for th, t, a, o in self.history],
        )

    def _build_prompt(self, last_observation):
        parts = []

        # Disambiguation warning
        if self._is_disambiguation(last_observation):
            parts.append("THE GAME IS ASKING YOU A QUESTION. Answer it directly.\n")

        # Last observation and current location
        parts.append("=== LAST OBSERVATION ===")
        parts.append(last_observation)
        parts.append(f"\nYou are at: {self.current_location}")

        # Guide for new rooms
        if self.just_entered_new_room:
            unexplored = self._get_unexplored_exits(self.current_location)
            arrival_dir = self._get_arrival_direction(self.current_location)
            back_dir = REVERSE_DIR.get(arrival_dir) if arrival_dir else None

            parts.append(f"\nYOU JUST ENTERED A NEW ROOM!")
            parts.append("Priority: 1) Take items on the ground. 2) React to sounds/events. 3) Try unexplored exits (not the way you came).")
            if unexplored:
                forward_exits = [d for d in unexplored if d != back_dir]
                if forward_exits:
                    parts.append(f"Exits to explore: {', '.join(forward_exits)}")
                if back_dir:
                    parts.append(f"(You came from {back_dir} — explore other exits first)")

        # Inventory
        parts.append(f"\n=== INVENTORY ===")
        parts.append(", ".join(sorted(self.inventory)) or "Empty")

        # Unexplored exits at current location
        unexplored = self._get_unexplored_exits(self.current_location)
        taken = self.exits_taken.get(self.current_location, set())
        failed = self.exits_failed.get(self.current_location, set())

        if taken or failed or unexplored:
            parts.append(f"\n=== EXITS AT '{self.current_location}' ===")
            if unexplored:
                parts.append(f"⚡ UNEXPLORED: {', '.join(unexplored)}")
            if taken:
                details = []
                for d in sorted(taken):
                    dest = self.world_map.get(self.current_location, {}).get(d, "?")
                    details.append(f"{d}→{dest}")
                parts.append(f"Working: {', '.join(details)}")
            if failed:
                parts.append(f"Blocked: {', '.join(sorted(failed))}")
            if not unexplored:
                parts.append("All exits explored. Now interact with objects: search, open, examine, push, pull.")

        # Build and exploit frontier
        frontier = self._get_frontier()
        frontier_other = [(loc, exits) for loc, exits in frontier if loc != self.current_location]
        if frontier_other and not unexplored:
            parts.append(f"\n=== FRONTIER: rooms with unexplored exits ===")
            for loc, exits in frontier_other:
                # Find path hint
                path_hint = ""
                for this_dir, dest in self.world_map.get(self.current_location, {}).items():
                    if dest == loc:
                        path_hint = f" (go {this_dir} from here)"
                        break
                parts.append(f"  {loc}: unexplored = {', '.join(exits)}{path_hint}")
            parts.append("→ Navigate to one of these rooms to continue exploring.")

        # Actions tried at current location
        tried_here = self._get_recent_actions_at(self.current_location)
        if tried_here:
            parts.append(f"\n=== ACTIONS TRIED AT '{self.current_location}' ===")
            for a in tried_here:
                parts.append(f"  {a}")

        # Map of visited locations
        if len(self.visited_locations) > 1:
            parts.append(f"\n=== MAP ({len(self.visited_locations)} rooms) ===")
            for loc in sorted(self.visited_locations):
                marker = " [HERE]" if loc == self.current_location else ""
                conns = self.world_map.get(loc, {})
                conn_str = ", ".join(f"{d}→{dest}" for d, dest in conns.items()) if conns else "?"
                unexpl = self._get_unexplored_exits(loc)
                status = f" | unexplored: {', '.join(unexpl)}" if unexpl else " | explored"
                parts.append(f"  {loc}{marker}: {conn_str}{status}")

        # Room notes
        other_notes = {loc: note for loc, note in self.location_notes.items() if loc != self.current_location}
        if other_notes:
            parts.append(f"\n=== ROOM NOTES ===")
            for loc, note in other_notes.items():
                parts.append(f"  {loc}: {note}")

        # Global summary
        if self.global_summary:
            parts.append(f"\n=== STRATEGY ===")
            parts.append(self.global_summary)

        # Recent history
        if self.history:
            recent = self.history[-8:]
            parts.append(f"\n=== RECENT ACTIONS ===")
            for th, tool, args, obs in recent:
                if tool == "play_action":
                    act = args.get("action", "?")
                    obs_short = obs[:120].replace("\n", " ")
                    parts.append(f"  > {act} => {obs_short}")
                else:
                    parts.append(f"  > [memory]")

        if self.consecutive_same_loc >= 6 and not self.just_entered_new_room:
            parts.append(f"""
You've been at '{self.current_location}' for {self.consecutive_same_loc} turns.
Try: listen, search <object>, look in/under/behind <thing>, push/pull/move <thing>, wait.
Or navigate to a frontier room with unexplored exits.
Remember: "use X" is NOT valid — use specific verbs like put, push, pull, light, give.""")

        return "\n".join(parts)

    def _parse_response(self, response):
        thought_m = re.search(r"THOUGHT:\s*(.*?)(?=\nTOOL:|\Z)", response, re.DOTALL)
        tool_m = re.search(r"TOOL:\s*(\S+)", response)
        args_m = re.search(r"ARGS:\s*(\{.*?\})", response, re.DOTALL)

        thought = thought_m.group(1).strip() if thought_m else ""
        tool = tool_m.group(1).strip() if tool_m else ""

        try:
            args = json.loads(args_m.group(1)) if args_m else {}
        except Exception:
            args = {}

        return thought, tool, args

# =============================================================================
# For local testing
# =============================================================================

async def test_agent():
    """Test the agent locally."""
    from fastmcp import Client
    
    # Path to your MCP server
    server_path = "mcp_server.py"
    
    agent = StudentAgent()
    
    async with Client(server_path) as client:
        result = await agent.run(
            client=client,
            game="zork1",
            max_steps=10,
            seed=42,
            verbose=True,
        )
        
        print(f"\nFinal Score: {result.final_score}")
        print(f"Moves: {result.moves}")
        print(f"Locations: {result.locations_visited}")


if __name__ == "__main__":
    import asyncio
    asyncio.run(test_agent())