agent.py

import json
import hashlib
import os
import random
import re
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import Optional

from dotenv import load_dotenv
from huggingface_hub import InferenceClient

load_dotenv()

USE_LOCAL_MODEL = os.getenv("USE_LOCAL_MODEL", "1").strip().lower() in {"1", "true", "yes"}
LOCAL_MODEL_ID = os.getenv("LOCAL_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
LOCAL_MAX_NEW_TOKENS = int(os.getenv("LOCAL_MAX_NEW_TOKENS", "56"))
ENABLE_LLM_PICK = os.getenv("ENABLE_LLM_PICK", "0").strip().lower() in {"1", "true", "yes"}
USE_VALID_ACTIONS = os.getenv("USE_VALID_ACTIONS", "1").strip().lower() in {"1", "true", "yes"}
REMOTE_MODEL = os.getenv("REMOTE_MODEL", "Qwen/Qwen2.5-72B-Instruct")

DISALLOWED_ACTIONS = {"inventory", "i", "quit", "q", "restart", "restore", "save", "script", "unscript"}
MOVE_ACTIONS = {
    "north",
    "south",
    "east",
    "west",
    "northeast",
    "northwest",
    "southeast",
    "southwest",
    "up",
    "down",
    "enter",
    "exit",
}
NEGATIVE_PATTERNS = (
    "you can't",
    "you cannot",
    "that's not",
    "i don't",
    "nothing happens",
    "there is no",
    "not here",
    "not open",
    "not allowed",
    "not see",
)

_local_pipeline = None
_remote_client = None
LLM_READY = False

if ENABLE_LLM_PICK:
    if USE_LOCAL_MODEL:
        try:
            import torch
            from transformers import pipeline as hf_pipeline

            dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
            _local_pipeline = hf_pipeline("text-generation", model=LOCAL_MODEL_ID, device_map="auto", torch_dtype=dtype)
            LLM_READY = True
        except Exception:
            LLM_READY = False
    else:
        token = os.getenv("HF_TOKEN")
        if token:
            _remote_client = InferenceClient(token=token)
            LLM_READY = True


@dataclass
class RunResult:
    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)


@dataclass
class ActionStats:
    tries: int = 0
    success: int = 0
    fail: int = 0
    score_gain: int = 0


@dataclass
class LocationMemory:
    visits: int = 0
    valid_actions: list[str] = field(default_factory=list)
    actions: dict[str, ActionStats] = field(default_factory=dict)
    failed: set[str] = field(default_factory=set)
    promising: set[str] = field(default_factory=set)
    exits: dict[str, str] = field(default_factory=dict)
    move_fail_sigs: dict[str, set[str]] = field(default_factory=dict)


def _extract_generated_text(outputs) -> str:
    if not outputs:
        return ""
    first = outputs[0]
    text = first.get("generated_text", "") if isinstance(first, dict) else str(first)
    if isinstance(text, list) and text:
        tail = text[-1]
        return str(tail.get("content", "")) if isinstance(tail, dict) else str(tail)
    return str(text)


def call_llm(prompt: str, seed: int, max_tokens: int = 56) -> str:
    if not LLM_READY:
        return ""
    if USE_LOCAL_MODEL and _local_pipeline is not None:
        kwargs = {"max_new_tokens": min(max_tokens, LOCAL_MAX_NEW_TOKENS), "do_sample": False, "temperature": 0.0}
        try:
            return _extract_generated_text(_local_pipeline([{"role": "user", "content": prompt}], **kwargs))
        except Exception:
            plain = f"USER:\n{prompt}\n\nASSISTANT:\n"
            return _extract_generated_text(_local_pipeline(plain, **kwargs))
    if _remote_client is None:
        return ""
    try:
        resp = _remote_client.chat.completions.create(
            model=REMOTE_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
            max_tokens=max_tokens,
            seed=seed,
        )
        return resp.choices[0].message.content or ""
    except Exception:
        return ""


class StudentAgent:
    def __init__(self):
        self.rng = random.Random(0)
        self.tool_names = set()

        self.score = 0
        self.max_score_seen = 350
        self.no_progress_streak = 0
        self.same_location_streak = 0

        self.last_obs_sig = ""
        self.last_loc = "Unknown"
        self.llm_calls = 0

        self.history = []
        self.action_trace = deque(maxlen=16)
        self.location_trace = deque(maxlen=20)
        self.action_queue = deque(maxlen=16)
        self.tabu_pairs = deque(maxlen=28)

        self.location_memory = defaultdict(LocationMemory)
        self.global_stats = defaultdict(ActionStats)
        self.global_failures = defaultdict(int)
        self.global_successes = defaultdict(int)
        self.term_counts = defaultdict(int)

        self.inventory = set()

    async def run(self, client, game: str, max_steps: int, seed: int, verbose: bool = False) -> RunResult:
        self.rng.seed(seed)
        tools = await client.list_tools()
        self.tool_names = {t.name for t in tools}
        if "play_action" not in self.tool_names:
            return RunResult(0, 0, 0, set(), False, error="play_action tool missing")

        visited, out_hist = set(), []

        text = self._extract_tool_text(await client.call_tool("play_action", {"action": "look"}))
        obs, meta = self._split_observation_and_meta(text)
        cur_loc = self._loc(obs, meta)
        self._record(action="look", from_loc=cur_loc, obs=obs, meta=meta)
        visited.add(self._visit_loc(obs, meta))

        if "inventory" in self.tool_names:
            self.inventory.update(await self._fetch_inventory(client))
        if USE_VALID_ACTIONS and "get_valid_actions" in self.tool_names:
            valid = await self._fetch_valid_actions(client)
            if self._valid_actions_usable(valid):
                self.location_memory[cur_loc].valid_actions = valid

        for step in range(1, max_steps + 1):
            if self._done(obs, meta):
                break

            cur_loc = self._loc(obs, meta)
            loc_mem = self.location_memory[cur_loc]

            entered_new = bool(meta.get("changed_location", False)) if isinstance(meta, dict) else False
            if loc_mem.visits <= 1:
                entered_new = True

            if "inventory" in self.tool_names and (step <= 3 or step % 14 == 0 or self.no_progress_streak >= 3):
                self.inventory.update(await self._fetch_inventory(client))

            if USE_VALID_ACTIONS and "get_valid_actions" in self.tool_names and (
                entered_new
                or not loc_mem.valid_actions
                or (self.no_progress_streak >= 4 and step % 4 == 0)
            ):
                valid = await self._fetch_valid_actions(client)
                if self._valid_actions_usable(valid):
                    loc_mem.valid_actions = valid

            action, thought = self._choose_action(cur_loc, obs, step, max_steps, seed + step)
            try:
                text = self._extract_tool_text(await client.call_tool("play_action", {"action": action}))
                obs, meta = self._split_observation_and_meta(text)
            except Exception as exc:
                obs, meta = f"Tool error: {exc}", {}

            self._record(action=action, from_loc=cur_loc, obs=obs, meta=meta)
            visited.add(self._visit_loc(obs, meta))
            out_hist.append((thought, action, obs[:220]))

            if verbose:
                print(f"[step {step}] score={self.score} loc={self._loc(obs,meta)} stall={self.no_progress_streak} action={action}")
            if self._done(obs, meta):
                break

        moves = int(meta.get("moves", 0)) if isinstance(meta, dict) else 0
        return RunResult(self.score, self.max_score_seen, moves, visited, self._done(obs, meta), history=out_hist)

    async def _fetch_valid_actions(self, client) -> list[str]:
        try:
            raw = self._extract_tool_text(await client.call_tool("get_valid_actions", {"timeout_s": 2.5}))
        except Exception:
            return []
        return self._parse_valid_actions(raw)

    async def _fetch_inventory(self, client) -> set[str]:
        try:
            raw = self._extract_tool_text(await client.call_tool("inventory", {}))
        except Exception:
            return set()
        return self._parse_inventory(raw)

    def _record(self, action: str, from_loc: str, obs: str, meta: dict):
        action = self._norm(action)
        to_loc = self._loc(obs, meta)
        self.location_memory[to_loc].visits += 1
        self.location_trace.append(to_loc)
        self.action_trace.append(action)

        prev_score = self.score
        score_now = int(meta.get("score", self._extract_score(obs))) if isinstance(meta, dict) else self._extract_score(obs)
        self.score = max(self.score, score_now)
        score_gain = max(0, self.score - prev_score)

        if isinstance(meta, dict):
            self.max_score_seen = max(self.max_score_seen, int(meta.get("max_score", self.max_score_seen)))

        obs_sig = self._sig(obs)
        prev_sig = self.last_obs_sig
        changed_loc = (to_loc != from_loc) or (bool(meta.get("changed_location", False)) if isinstance(meta, dict) else False)
        negative = self._negative(obs)

        progress = "none"
        if score_gain > 0:
            progress = "score"
        elif changed_loc:
            progress = "move"
        elif obs_sig != self.last_obs_sig and not negative and action not in {"look", "inventory"}:
            progress = "state"

        if progress == "none":
            self.no_progress_streak += 1
            self.same_location_streak += 1
        else:
            self.no_progress_streak = 0
            self.same_location_streak = 0 if changed_loc else max(0, self.same_location_streak - 1)

        from_mem = self.location_memory[from_loc]
        stats = from_mem.actions.setdefault(action, ActionStats())
        global_stats = self.global_stats[action]
        stats.tries += 1
        global_stats.tries += 1

        if progress == "none":
            stats.fail += 1
            global_stats.fail += 1
            if self._is_move(action):
                # Contextual move memory: avoid repeating same failed move in same local context.
                key_sig = prev_sig or obs_sig
                from_mem.move_fail_sigs.setdefault(action, set()).add(key_sig)
            else:
                from_mem.failed.add(action)  # strict for non-move actions
            self.global_failures[action] += 1
        else:
            stats.success += 1
            global_stats.success += 1
            self.global_successes[action] += 1
            if action in from_mem.failed:
                from_mem.failed.remove(action)
            if self._is_move(action) and action in from_mem.move_fail_sigs:
                from_mem.move_fail_sigs.pop(action, None)

        if score_gain > 0:
            stats.score_gain += score_gain
            global_stats.score_gain += score_gain
            from_mem.promising.add(action)

        if changed_loc and self._is_move(action):
            from_mem.exits[action] = to_loc

        for hint in self._extract_promising_actions(obs):
            from_mem.promising.add(hint)
        if progress in {"state", "score"}:
            for follow in self._extract_followup_actions(obs):
                if follow not in from_mem.failed:
                    self.action_queue.append(follow)
        self._update_terms(obs)

        self.history.append({"loc": to_loc, "action": action, "score": self.score, "progress": progress})
        self.tabu_pairs.append((from_loc, action))
        if len(self.history) > 120:
            self.history = self.history[-120:]

        self.last_obs_sig = obs_sig
        self.last_loc = to_loc

    def _choose_action(self, loc: str, obs: str, step: int, max_steps: int, seed: int) -> tuple[str, str]:
        loc_mem = self.location_memory[loc]
        valid = list(loc_mem.valid_actions)
        valid_set = set(valid)
        obs_sig = self._sig(obs)
        low_obs = obs.lower()
        has_audio_hint = any(w in low_obs for w in ("noise", "sound", "hear", "listen"))

        hints = self._extract_promising_actions(obs)
        moves = [a for a in valid if self._is_move(a)] or list(MOVE_ACTIONS)

        verbs = self._candidate_verbs(valid)
        objects = sorted(self._extract_objects(obs))[:8]
        inv_objs = sorted(self.inventory)[:4]
        terms = self._top_terms(8)
        allow_composed = self._composed_mode(valid)

        generated = []
        for obj in objects:
            generated.extend([f"examine {obj}", f"search {obj}"])
            generated.extend(f"{v} {obj}" for v in verbs[:4])
        for obj in inv_objs:
            generated.extend(f"{v} {obj}" for v in verbs[:2])
        composed = self._compose_actions(objects, inv_objs, terms) if allow_composed else []

        pool = hints + sorted(loc_mem.promising) + valid + generated + composed + moves + ["look", "wait", "listen"]

        candidates = []
        seen = set()
        for cand in pool:
            a = self._norm(cand)
            if not a or a in seen or len(a) > 48:
                continue
            if a in DISALLOWED_ACTIONS or self._unsafe_action(a) or self._incomplete_action(a):
                continue
            if a == "listen" and not has_audio_hint and not allow_composed:
                continue
            if self._is_move(a) and self._move_context_blocked(loc, a, obs_sig):
                continue
            seen.add(a)
            candidates.append(a)

        # Heavy context management: failed actions are banned in that location.
        candidates = [a for a in candidates if a not in loc_mem.failed]

        if not candidates:
            backup = [a for a in moves if a not in loc_mem.failed]
            if backup:
                candidates = backup
            else:
                return "look", "fallback"

        while self.action_queue:
            queued = self._norm(self.action_queue.popleft())
            if queued and queued in candidates and queued not in loc_mem.failed:
                return queued, "queued-followup"

        # Prioritize frontier exploration after a small amount of local probing.
        local_non_move_tries = sum(st.tries for a, st in loc_mem.actions.items() if not self._is_move(a))
        fresh_local = [
            a
            for a in candidates
            if a in valid_set
            and not self._is_move(a)
            and a not in {"look", "wait"}
            and self._likely_actionable(a)
            and loc_mem.actions.get(a, ActionStats()).tries == 0
        ]
        if fresh_local and loc_mem.visits <= 2 and local_non_move_tries < 2 and self.no_progress_streak <= 1:
            pick = max(
                fresh_local,
                key=lambda a: (
                    self._contextual_priority(a, obs),
                    self._score_action(loc, a, valid_set, set(hints)),
                ),
            )
            return pick, "local-probe"

        composed_probe = [
            a
            for a in candidates
            if self._is_composed_action(a)
            and loc_mem.actions.get(a, ActionStats()).tries == 0
            and self._likely_actionable(a)
        ]
        if allow_composed and composed_probe and self.no_progress_streak >= 2:
            pick = max(composed_probe[:12], key=lambda a: self._score_action(loc, a, valid_set, set(hints)))
            return pick, "composed-probe"

        frontier_move = self._frontier_move(loc, candidates, obs_sig)
        rich_non_move = [a for a in valid if " " in a and not self._is_move(a)]
        frontier_gate = (self.no_progress_streak >= 1 or local_non_move_tries >= 2 or loc_mem.visits >= 2)
        if len(rich_non_move) >= 2:
            frontier_gate = (self.no_progress_streak >= 3 or local_non_move_tries >= 4 or loc_mem.visits >= 4)
        if frontier_move and frontier_gate:
            return frontier_move, "frontier"

        if self.same_location_streak >= 3 or self.no_progress_streak >= 3 or loc_mem.visits >= 4:
            move_pick = self._least_tried_move(loc, candidates)
            if move_pick:
                return move_pick, "exploration-bias"

        ranked = sorted(
            candidates,
            key=lambda a: self._score_action(loc, a, valid_set, set(hints)),
            reverse=True,
        )
        if not ranked:
            return "look", "fallback"

        if self._should_use_llm(step, max_steps, ranked):
            picked = self._llm_pick(loc, obs, ranked[:8], seed)
            if picked in ranked:
                self.llm_calls += 1
                return picked, "llm-pick"

        return ranked[0], "heuristic"

    def _move_context_blocked(self, loc: str, action: str, obs_sig: str) -> bool:
        fails = self.location_memory[loc].move_fail_sigs.get(action, set())
        return bool(obs_sig and obs_sig in fails)

    def _frontier_move(self, loc: str, candidates: list[str], obs_sig: str) -> str:
        loc_mem = self.location_memory[loc]
        move_candidates = [a for a in candidates if self._is_move(a)]
        if not move_candidates:
            return ""

        # Direct frontier: untried exits from current location first.
        untried = [
            a
            for a in move_candidates
            if loc_mem.actions.get(a, ActionStats()).tries == 0 and not self._move_context_blocked(loc, a, obs_sig)
        ]
        if untried:
            untried.sort(key=lambda a: loc_mem.actions.get(a, ActionStats()).tries)
            return untried[0]

        # Indirect frontier: shortest known path to a location with untried exits.
        q = deque([(loc, "")])
        seen = {loc}
        while q:
            node, first_act = q.popleft()
            if node != loc and self._has_untried_exits(node):
                if first_act and first_act in move_candidates and not self._move_context_blocked(loc, first_act, obs_sig):
                    return first_act
            for act, dst in self.location_memory[node].exits.items():
                if not dst or dst in seen:
                    continue
                seen.add(dst)
                q.append((dst, first_act or act))
        return ""

    def _has_untried_exits(self, loc: str) -> bool:
        mem = self.location_memory[loc]
        moves = [a for a in mem.valid_actions if self._is_move(a)]
        if not moves:
            moves = list(MOVE_ACTIONS)
        return any(mem.actions.get(a, ActionStats()).tries == 0 for a in moves)

    def _score_action(self, loc: str, action: str, valid_set: set[str], hint_set: set[str]) -> float:
        loc_mem = self.location_memory[loc]
        st = loc_mem.actions.get(action, ActionStats())
        gt = self.global_stats.get(action, ActionStats())

        score = 60.0 * st.score_gain + 8.0 * st.success - 12.0 * st.fail - 3.0 * st.tries - 1.0 * gt.tries
        if st.tries == 0:
            score += 10.0
        if gt.tries == 0:
            score += 4.0
        if action in valid_set:
            score += 8.0
        if action in hint_set or action in loc_mem.promising:
            score += 10.0
        if self._is_move(action):
            score += 2.0
            if st.tries == 0:
                score += 6.0
            if self.no_progress_streak >= 1:
                score += 12.0
            if self.same_location_streak >= 2:
                score += 10.0
        else:
            if st.tries == 0 and loc_mem.visits <= 2:
                score += 12.0
            if self.same_location_streak >= 3:
                score -= 10.0
            if self._is_composed_action(action):
                score += 6.0 if self.no_progress_streak >= 2 else 1.0

        recent = list(self.action_trace)[-3:]
        if action in recent:
            score -= 10.0
        if len(recent) >= 2 and recent[-1] == recent[-2] == action:
            score -= 20.0
        if (loc, action) in self.tabu_pairs:
            score -= 14.0

        if self.global_failures[action] >= 3 and self.global_successes[action] == 0:
            score -= 12.0

        opposite = {"north": "south", "south": "north", "east": "west", "west": "east", "up": "down", "down": "up", "enter": "exit", "exit": "enter"}
        last = self.action_trace[-1] if self.action_trace else ""
        if opposite.get(last, "") == action and self.no_progress_streak <= 1:
            score -= 16.0

        return score

    def _least_tried_move(self, loc: str, candidates: list[str]) -> str:
        loc_mem = self.location_memory[loc]
        moves = [a for a in candidates if self._is_move(a) and a not in loc_mem.failed]
        if not moves:
            return ""
        obs_sig = self.last_obs_sig
        moves = [a for a in moves if not self._move_context_blocked(loc, a, obs_sig)] or moves
        moves.sort(key=lambda a: loc_mem.actions.get(a, ActionStats()).tries)
        return moves[0]

    @staticmethod
    def _verb_priority(action: str) -> int:
        head = action.split()[0] if action.split() else ""
        pri = {"open": 6, "unlock": 6, "take": 5, "get": 5, "read": 4, "use": 4, "enter": 3, "examine": 2, "search": 1}
        return pri.get(head, 0)

    def _contextual_priority(self, action: str, obs: str) -> int:
        base = self._verb_priority(action)
        head = action.split()[0] if action.split() else ""
        low = obs.lower()
        if any(w in low for w in ("noise", "sound", "hear")):
            if head in {"search", "examine", "listen"}:
                base += 4
            if head in {"open", "take"}:
                base -= 3
        if head in {"open", "unlock"} and any(k in action for k in ("door", "window", "mailbox", "gate", "chest", "box", "lid")):
            base += 5
        return base

    @staticmethod
    def _is_composed_action(action: str) -> bool:
        head = action.split()[0] if action.split() else ""
        if head in {"ask", "tell", "show", "give", "follow", "catch", "greet", "talk"}:
            return True
        return any(tok in action for tok in (" about ", " with ", " to ", " under ", " inside ", " behind ", " in ", " on "))

    @staticmethod
    def _composed_mode(valid_actions: list[str]) -> bool:
        if not valid_actions:
            return True
        non_move = [a for a in valid_actions if a not in MOVE_ACTIONS and a not in {"look", "wait", "listen"}]
        rich = [a for a in non_move if " " in a]
        return len(rich) <= 1

    @staticmethod
    def _likely_actionable(action: str) -> bool:
        parts = action.split()
        if len(parts) <= 1:
            return parts[0] in {"look", "wait", "listen"} if parts else False
        obj = parts[-1]
        generic = {
            "north",
            "south",
            "east",
            "west",
            "up",
            "down",
            "inside",
            "outside",
            "back",
            "front",
            "forest",
            "field",
            "path",
            "road",
            "area",
            "dark",
            "night",
        }
        return obj not in generic

    def _update_terms(self, obs: str):
        for t in self._extract_keywords(obs):
            self.term_counts[t] += 1

    def _top_terms(self, k: int) -> list[str]:
        ranked = sorted(self.term_counts.items(), key=lambda kv: kv[1], reverse=True)
        return [t for t, _ in ranked[:k]]

    @staticmethod
    def _extract_keywords(obs: str) -> list[str]:
        low = obs.lower()
        stop = {
            "the", "and", "for", "with", "from", "that", "this", "there", "here", "into", "over",
            "you", "your", "are", "was", "were", "have", "has", "had", "not", "but", "all", "any",
            "north", "south", "east", "west", "up", "down", "look", "time", "night", "dark",
            "forest", "field", "path", "road", "place", "some", "thing", "things",
        }
        out = []
        for tok in re.findall(r"[a-z]{3,16}", low):
            if tok in stop:
                continue
            out.append(tok)
        return out

    def _compose_actions(self, objects: list[str], inv_objs: list[str], terms: list[str]) -> list[str]:
        nouns, seen = [], set()
        for item in list(objects) + list(terms):
            n = self._normalize_noun(str(item))
            if not n or n in seen or len(n) > 24:
                continue
            seen.add(n)
            nouns.append(n)
        nouns = nouns[:8]
        actors = [n for n in nouns if len(n.split()) == 1][:5]
        inv_clean = [self._normalize_noun(x) for x in inv_objs]
        inv_clean = [x for x in inv_clean if x]

        out = []
        for n in nouns:
            out.extend([f"look for {n}", f"look under {n}", f"look behind {n}", f"look inside {n}"])
            out.extend([f"follow {n}", f"catch {n}"])
        for a in actors:
            out.extend([f"greet {a}", f"talk to {a}"])
            for t in nouns[:4]:
                if t != a:
                    out.extend([f"ask {a} about {t}", f"tell {a} about {t}"])
        for it in inv_clean[:3]:
            for a in actors[:3]:
                out.extend([f"show {it} to {a}", f"give {it} to {a}"])
            for t in nouns[:3]:
                out.extend([f"use {it} with {t}", f"put {it} in {t}", f"put {it} on {t}"])

        dedup, used = [], set()
        for a in out:
            a = self._norm(a)
            if not a or a in used or len(a) > 48:
                continue
            if a in DISALLOWED_ACTIONS or self._unsafe_action(a) or self._incomplete_action(a):
                continue
            used.add(a)
            dedup.append(a)
            if len(dedup) >= 48:
                break
        return dedup

    @staticmethod
    def _normalize_noun(text: str) -> str:
        low = re.sub(r"[^a-z ]+", " ", (text or "").lower())
        toks = [t for t in low.split() if t]
        drop = {
            "a", "an", "the", "of", "to", "in", "on", "at", "from", "with", "and",
            "all", "everything", "anything", "something", "thing", "things",
            "north", "south", "east", "west", "up", "down",
        }
        toks = [t for t in toks if t not in drop]
        if not toks:
            return ""
        out = " ".join(toks[:2])
        return out if len(out) >= 3 else ""

    def _should_use_llm(self, step: int, max_steps: int, ranked: list[str]) -> bool:
        if not ENABLE_LLM_PICK or not LLM_READY:
            return False
        if len(ranked) < 3:
            return False
        if self.llm_calls >= min(10, max_steps // 8 + 2):
            return False
        return self.no_progress_streak >= 3 or step <= 2 or step % 16 == 0

    def _llm_pick(self, loc: str, obs: str, candidates: list[str], seed: int) -> str:
        opts = "\n".join(f"- {c}" for c in candidates)
        prompt = (
            f"Location: {loc}\n"
            f"Score: {self.score}/{self.max_score_seen}\n"
            f"No-progress streak: {self.no_progress_streak}\n"
            f"Recent actions: {list(self.action_trace)[-6:]}\n\n"
            f"Observation:\n{obs[:1200]}\n\n"
            f"Choose one exact action from this list:\n{opts}\n"
            "Return JSON: {\"action\":\"...\"}."
        )
        text = call_llm(prompt, seed=seed, max_tokens=40)
        if not text:
            return ""

        try:
            payload = json.loads(text)
            pick = self._norm(str(payload.get("action", "")))
            return pick if pick in candidates else ""
        except Exception:
            pass

        low = text.lower()
        for cand in candidates:
            if cand in low:
                return cand
        return ""

    @staticmethod
    def _extract_tool_text(result) -> str:
        if hasattr(result, "content") and result.content:
            chunk = result.content[0]
            return chunk.text if hasattr(chunk, "text") else str(chunk)
        if isinstance(result, list) and result:
            item = result[0]
            return item.text if hasattr(item, "text") else str(item)
        return str(result)

    def _split_observation_and_meta(self, text: str) -> tuple[str, dict]:
        marker = "\n[META]"
        if marker not in text:
            return text.strip(), {}
        obs, meta_blob = text.rsplit(marker, 1)
        try:
            meta = json.loads(meta_blob.strip())
            if isinstance(meta, dict):
                return obs.strip(), meta
        except Exception:
            pass
        return text.strip(), {}

    @staticmethod
    def _norm(action: str) -> str:
        a = " ".join((action or "").strip().lower().split())
        alias = {
            "n": "north",
            "s": "south",
            "e": "east",
            "w": "west",
            "u": "up",
            "d": "down",
            "i": "inventory",
            "l": "look",
            "ne": "northeast",
            "nw": "northwest",
            "se": "southeast",
            "sw": "southwest",
        }
        if a in alias:
            return alias[a]
        if a.startswith("go "):
            return alias.get(a.split(" ", 1)[1].strip(), a.split(" ", 1)[1].strip())
        return a

    @staticmethod
    def _is_move(action: str) -> bool:
        return action in MOVE_ACTIONS

    @staticmethod
    def _unsafe_action(action: str) -> bool:
        parts = action.split()
        if not parts:
            return True
        if " all " in f" {action} ":
            return True
        if parts[0] in {"take", "drop", "put", "throw", "eat"} and len(parts) >= 2:
            if parts[-1] in {"thing", "things", "something", "anything", "everything", "all"}:
                return True
        return False

    @staticmethod
    def _incomplete_action(action: str) -> bool:
        parts = action.split()
        if len(parts) != 1:
            return False
        return parts[0] in {"examine", "search", "take", "drop", "open", "close", "read", "use", "insert", "put", "throw", "attack"}

    @staticmethod
    def _parse_inventory(raw: str) -> set[str]:
        text = str(raw or "").strip()
        low = text.lower()
        if not text or "empty" in low or "not initialized" in low:
            return set()
        payload = text.split(":", 1)[1] if ":" in text else text
        out = set()
        for part in payload.lower().split(","):
            item = " ".join(part.strip().split())
            if item and len(item) <= 30 and re.fullmatch(r"[a-z][a-z -]*[a-z]", item):
                out.add(item)
        return out

    @staticmethod
    def _parse_valid_actions(raw: str) -> list[str]:
        try:
            payload = json.loads(raw)
            actions = payload.get("valid_actions", []) if isinstance(payload, dict) else []
        except Exception:
            return []
        if not isinstance(actions, list):
            return []

        out, seen = [], set()
        for item in actions:
            a = " ".join(str(item).strip().lower().split())
            if not a or a in seen or len(a) > 48:
                continue
            if a in DISALLOWED_ACTIONS or StudentAgent._unsafe_action(a) or StudentAgent._incomplete_action(a):
                continue
            seen.add(a)
            out.append(a)
            if len(out) >= 36:
                break
        return out

    @staticmethod
    def _valid_actions_usable(actions: list[str]) -> bool:
        return bool(actions) and any(a in MOVE_ACTIONS for a in actions)

    @staticmethod
    def _extract_score(text: str) -> int:
        for pat in (r"\[score:\s*(\d+)", r"score[:\s]+(\d+)"):
            m = re.search(pat, text, flags=re.IGNORECASE)
            if m:
                return int(m.group(1))
        return 0

    @staticmethod
    def _sig(text: str) -> str:
        return re.sub(r"\s+", " ", (text or "").strip().lower())[:320]

    def _loc(self, obs: str, meta: dict) -> str:
        if isinstance(meta, dict):
            loc = str(meta.get("location", "")).strip()
            if loc:
                return self._canon_loc(loc)
        lines = [ln.strip() for ln in obs.splitlines() if ln.strip()]
        if not lines:
            return "Unknown"
        first = lines[0]
        if len(first) <= 70 and not first.endswith((".", "!", "?")):
            return self._canon_loc(first)
        sent = re.split(r"[.!?]", first, maxsplit=1)[0].strip()
        return self._canon_loc(sent[:64] if sent else "Unknown")

    def _visit_loc(self, obs: str, meta: dict) -> str:
        base = self._loc(obs, meta)
        sig = self._sig(obs)
        if not sig:
            return base
        h = hashlib.sha1(sig.encode("utf-8")).hexdigest()[:8]
        action = ""
        if isinstance(meta, dict):
            action = self._norm(str(meta.get("action", "")))
        head = action.split()[0] if action else ""
        streak = min(self.no_progress_streak, 9)
        linger = min(self.same_location_streak, 9)
        return f"{base}@{h}:{head}:{streak}:{linger}" if head else f"{base}@{h}:{streak}:{linger}"

    @staticmethod
    def _canon_loc(loc: str) -> str:
        text = " ".join((loc or "Unknown").strip().split())
        parts = text.split("#")
        if len(parts) >= 3:
            tail = parts[-1]
            if tail.isdigit():
                return f"{parts[0]}#{tail}"
        return text

    @staticmethod
    def _done(obs: str, meta: dict) -> bool:
        if isinstance(meta, dict) and bool(meta.get("done", False)):
            return True
        low = obs.lower()
        return any(x in low for x in ("game over", "you have died", "you are dead", "the end"))

    @staticmethod
    def _negative(text: str) -> bool:
        low = text.lower()
        return any(p in low for p in NEGATIVE_PATTERNS)

    @staticmethod
    def _extract_promising_actions(obs: str) -> list[str]:
        low = obs.lower()
        out = []
        for m in re.finditer(r'"([a-z][a-z ]{1,36})"', low):
            out.append(" ".join(m.group(1).split()))
        for pat in (r"(?:try|perhaps|maybe|you could|you can)\s+([a-z]+(?:\s+[a-z]+){0,3})",):
            for m in re.finditer(pat, low):
                out.append(" ".join(m.group(1).split()))

        dedup, seen = [], set()
        for a in out:
            a = StudentAgent._norm(a)
            if not a or a in seen or a in DISALLOWED_ACTIONS or len(a) > 48:
                continue
            seen.add(a)
            dedup.append(a)
        return dedup[:10]

    @staticmethod
    def _extract_followup_actions(obs: str) -> list[str]:
        low = obs.lower()
        out = []

        for d in MOVE_ACTIONS:
            if re.search(rf"\b{d}\b", low):
                out.append(d)

        patterns = (
            r"(?:allow|allows|allowed|can|could|may|might)\s+(?:you\s+)?([a-z]+(?:\s+[a-z]+){0,2})",
            r"(?:lets?|enable|enables)\s+(?:you\s+)?(?:to\s+)?([a-z]+(?:\s+[a-z]+){0,2})",
        )
        for pat in patterns:
            for m in re.finditer(pat, low):
                cand = " ".join(m.group(1).split())
                if cand == "entry":
                    cand = "enter"
                out.append(cand)

        dedup, seen = [], set()
        for a in out:
            a = StudentAgent._norm(a)
            if not a or a in seen or len(a) > 48:
                continue
            if a in DISALLOWED_ACTIONS or StudentAgent._unsafe_action(a) or StudentAgent._incomplete_action(a):
                continue
            seen.add(a)
            dedup.append(a)
        return dedup[:10]

    @staticmethod
    def _extract_objects(obs: str) -> set[str]:
        stop = {
            "you",
            "your",
            "there",
            "here",
            "north",
            "south",
            "east",
            "west",
            "up",
            "down",
            "room",
            "path",
            "road",
            "forest",
            "field",
            "farm",
            "night",
            "dark",
            "sky",
            "moon",
            "place",
            "some",
        }
        low = obs.lower()
        out = set()
        patterns = (
            r"(?:there is|there are|you see|you can see)\s+(?:a|an|the|some)?\s*([a-z][a-z\- ]{1,30})",
            r"\b(?:a|an|the)\s+([a-z][a-z\- ]{1,22})\s+(?:is|are|lies|sits|stands)\b",
            r"\bin\s+([a-z]{4,20})\b",
        )
        for pat in patterns:
            for m in re.finditer(pat, low):
                cand = " ".join(m.group(1).split())
                cand = re.split(r"\b(?:here|there|that|which|who|with|near|on|in|from|to)\b", cand, maxsplit=1)[0].strip()
                if cand and cand not in stop and len(cand) <= 30:
                    out.add(cand)
        return out

    def _candidate_verbs(self, valid_actions: list[str]) -> list[str]:
        out, seen = [], set()
        for action in valid_actions:
            head = action.split()[0] if action.split() else ""
            if not head or head in MOVE_ACTIONS or head in DISALLOWED_ACTIONS:
                continue
            if head in seen:
                continue
            seen.add(head)
            out.append(head)
            if len(out) >= 8:
                break

        ranked_global = sorted(
            self.global_stats.items(),
            key=lambda kv: (kv[1].score_gain, kv[1].success - kv[1].fail),
            reverse=True,
        )
        for action, _ in ranked_global:
            head = action.split()[0] if action.split() else ""
            if head and head not in MOVE_ACTIONS and head not in DISALLOWED_ACTIONS and head not in seen:
                seen.add(head)
                out.append(head)
            if len(out) >= 10:
                break

        for fallback in ("examine", "search", "open", "take", "read", "use", "listen"):
            if fallback not in seen:
                out.append(fallback)
        return out[:10]