"""Decision making: menu-mediated, model-agnostic.

The design lesson from testing tiny models on real screenshots: a 256M VLM can
*read* a button label from pixels but cannot reliably reason, and a 0.5B text model
writes perfect JSON but picks the wrong widget when handed noisy free-form text.
Both do much better when the harness hands them a **short numbered menu** and asks
for a number. That is what this module builds, and it keeps the model's output space
finite no matter how small the model is.

Brains shipped here:

* :class:`MockBrain`    - replays a scripted action list (tests, demos, replays)
* :class:`PlannerBrain` - rule-based, no model: clicks the element matching the goal
* :class:`ModelBrain`   - base for real models: menu, prompt, lenient parsing
  (see :mod:`tiny_computer_agent.models` for the text and vision backends)
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any, Iterator, Mapping, Protocol, Sequence

from tiny_computer_agent.actions import ACTION_KINDS, Action, ActionError, parse_action
from tiny_computer_agent.observe import Element, Observation, normalise_label


class BrainError(RuntimeError):
    """Raised when a brain cannot produce a valid action."""


class GoalUnreachable(BrainError):
    """Raised by a brain that can decide nothing further, without claiming success.

    Deliberately distinct from a crash: the harness records the episode as a
    *failure*, because a brain admitting "I give up" must never be logged as
    "goal met" - and the harness, not the brain, owns that judgement.
    """


def extract_json_objects(raw: str) -> list[dict[str, Any]]:
    """Every flat JSON object in a reply, in the order they appear.

    Regex-based and non-nesting: replies from tiny models are one-level objects.
    Callers should prefer the **last** object, because a model that echoes its own
    prompt before answering puts the prompt's example first and its answer last.
    """
    objects: list[dict[str, Any]] = []
    for match in re.finditer(r"\{[^{}]*\}", raw or "", re.DOTALL):
        try:
            payload = json.loads(match.group(0))
        except json.JSONDecodeError:
            continue
        if isinstance(payload, dict):
            objects.append(payload)
    return objects


def extract_json_object(raw: str) -> dict[str, Any]:
    """The decision out of a reply: the *last* JSON object, not the first."""
    objects = extract_json_objects(raw)
    if not objects:
        raise BrainError(f"no JSON object in the model reply: {(raw or '')[:200]!r}")
    return objects[-1]


@dataclass(frozen=True)
class MenuItem:
    """One numbered choice shown to the model."""

    index: int
    element: Element

    @property
    def label(self) -> str:
        return self.element.label


class DecisionMenu:
    """A short, numbered, closed choice set built from the element index.

    Elements are presented in reading order (top to bottom, then left to right),
    and capped: a 20-item menu is as good as free-form text to a small model.
    """

    def __init__(self, observation: Observation, *, max_items: int = 12) -> None:
        ordered = sorted(observation.elements, key=lambda e: (e.y, e.x))[:max_items]
        self.items: tuple[MenuItem, ...] = tuple(
            MenuItem(index=i + 1, element=e) for i, e in enumerate(ordered)
        )

    def __len__(self) -> int:
        return len(self.items)

    def __bool__(self) -> bool:
        return bool(self.items)

    def __iter__(self) -> Iterator[MenuItem]:
        return iter(self.items)

    def item(self, value: Any) -> MenuItem | None:
        """Look up a 1-based index, tolerating the strings models actually send."""
        try:
            number = int(float(str(value).strip()))
        except (TypeError, ValueError):
            return None
        if 1 <= number <= len(self.items):
            return self.items[number - 1]
        return None

    def find(self, label: str) -> Element | None:
        """Resolve a label against the menu (normalised, like Observation.find)."""
        wanted = normalise_label(label)
        if not wanted:
            return None
        for item in self.items:
            if normalise_label(item.label) == wanted:
                return item.element
        for item in self.items:
            if wanted in normalise_label(item.label):
                return item.element
        return None

    def mentioned_in(self, text: str) -> tuple[MenuItem, ...]:
        """Menu entries whose label appears in the text, in order of first mention."""
        haystack = normalise_label(text)
        found: list[tuple[int, MenuItem]] = []
        for item in self.items:
            needle = normalise_label(item.label)
            if needle and (position := haystack.find(needle)) != -1:
                found.append((position, item))
        return tuple(item for _, item in sorted(found, key=lambda pair: pair[0]))

    def render(self) -> str:
        return "\n".join(
            f"  {item.index}. {item.label}  (at {item.element.center[0]},{item.element.center[1]})"
            for item in self.items
        )


def _recover_from_text(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None:
    """Salvage a decision from a payload that named the right thing in the wrong field.

    Observed for real: a 0.5B planner replied ``{"kind": "click", "index": 0,
    "reason": "Save"}`` - the correct answer, in the wrong place. If the index is
    unusable but exactly one menu label is mentioned anywhere in the payload, take
    the label. It is still only ever a label the detector saw in this frame, so the
    deny list, the dry run and the step budget all continue to apply.
    """
    text = json.dumps(payload, ensure_ascii=False)
    mentioned = menu.mentioned_in(text)
    if len({item.label for item in mentioned}) == 1:
        item = mentioned[0]
        return Action(
            "click", label=item.label,
            reason=f"recovered {item.label!r} from the model's reply: {text[:80]}",
        )
    return None


def _action_from_payload(payload: Mapping[str, Any], menu: DecisionMenu) -> Action | None:
    """Interpret one JSON object, tolerating the ways small models mangle the shape."""
    kind = payload.get("kind")
    if isinstance(kind, str) and kind.strip().lower() not in ACTION_KINDS:
        kind = None  # e.g. {"kind": "file"} - a garbled click; fall back to index/label

    if payload.get("index") is not None:
        item = menu.item(payload["index"])
        if item is not None:
            return Action("click", label=item.label, reason=f"menu index {payload['index']}")
        recovered = _recover_from_text(payload, menu)
        if recovered is not None:
            return recovered
        raise BrainError(
            f"the model chose index {payload['index']!r}, which is not in the "
            f"menu of {len(menu)} items"
        )

    if kind in (None, "click") and "label" in payload:
        element = menu.find(str(payload["label"]))
        if element is None:
            raise BrainError(f"the model named {payload['label']!r}, which is not on screen")
        return Action("click", label=element.label, reason=f"the model named {element.label!r}")

    if kind is not None:
        try:
            return parse_action(dict(payload, kind=kind))
        except ActionError as exc:
            # a valid-looking verb with no usable target: maybe the label is in the prose
            return _recover_from_text(payload, menu) or (_raise_brain_error(exc))

    return _recover_from_text(payload, menu)


def _raise_brain_error(exc: Exception) -> Action:
    raise BrainError(f"the model proposed an invalid action: {exc}") from exc


_DONE = re.compile(r"\b(done|finished|complete[d]?)\b", re.IGNORECASE)


def parse_decision(raw: str, observation: Observation, *, max_menu_items: int = 12) -> Action:
    """Turn a small model's reply into a validated Action - leniently, but safely.

    Tried in order:

    1. the **last** JSON object in the reply (``index``, ``label`` or a full action)
    2. a label the reply names (most specific evidence of intent)
    3. a bare menu number
    4. a "done" word

    A reply that mentions three or more menu entries is rejected as an echo of the
    prompt rather than a decision - the failure mode of the weakest models, and the
    one most likely to make an agent click something arbitrary.
    """
    text = " ".join((raw or "").split())
    if not text:
        raise BrainError("the model returned an empty reply")
    menu = DecisionMenu(observation, max_items=max_menu_items)

    for payload in reversed(extract_json_objects(text)):
        action = _action_from_payload(payload, menu)
        if action is not None:
            return action

    mentioned = menu.mentioned_in(text)
    if len(mentioned) >= 3:
        raise BrainError(
            f"the reply mentions {len(mentioned)} menu entries, so it looks like an echo "
            f"of the prompt rather than a decision: {text[:120]!r}"
        )
    if mentioned:
        item = mentioned[0]
        note = f"the reply names {item.label!r}"
        if len(mentioned) > 1:
            note += f" (ambiguous: {len(mentioned)} labels mentioned; took the first)"
        return Action("click", label=item.label, reason=note)

    number = re.search(r"\b(\d{1,2})\b", text)
    if number is not None:
        item = menu.item(number.group(1))
        if item is not None:
            return Action(
                "click", label=item.label,
                reason=f"menu index {number.group(1)} from reply {text[:40]!r}",
            )

    if _DONE.search(text):
        return Action("done", reason=_DONE.search(text).group(0).lower())

    raise BrainError(f"could not read a decision out of the reply: {text[:120]!r}")


class Brain(Protocol):
    """Anything that can turn (goal, observation) into one validated action."""

    name: str

    def decide(self, goal: str, observation: Observation) -> Action: ...


class MockBrain:
    """Replays a fixed script of actions, exactly once. Deterministic and offline."""

    name = "mock"

    def __init__(self, script: Sequence[Action | Mapping[str, Any]]) -> None:
        self._script: list[Action] = [
            a if isinstance(a, Action) else parse_action(a) for a in script
        ]
        self._cursor = 0

    @property
    def remaining(self) -> int:
        return max(0, len(self._script) - self._cursor)

    def decide(self, goal: str, observation: Observation) -> Action:
        if self._cursor >= len(self._script):
            return Action("done", reason="script exhausted")
        action = self._script[self._cursor]
        self._cursor += 1
        return action


class PlannerBrain:
    """Rule-based baseline: click the element matching the goal, then stop.

    Its purpose is threefold - it proves the harness works end to end with no model
    at all, it is a fallback when a model download is undesirable, and it shows how
    much of a "computer use" task is really an indexing problem. It waits when the
    target is not on screen yet, refuses to click the same label twice, and gives up
    rather than flailing. It will not generalise the way a model can.
    """

    name = "planner"

    def __init__(self, target: str | None = None, *,
                 done_when_text_contains: str | None = None,
                 max_waits: int = 3, wait_seconds: float = 0.5) -> None:
        self.target = target
        self.done_when_text_contains = done_when_text_contains
        self.max_waits = max_waits
        self.wait_seconds = wait_seconds
        self._clicked: set[str] = set()
        self._waits = 0

    def decide(self, goal: str, observation: Observation) -> Action:
        if self.done_when_text_contains:
            needle = self.done_when_text_contains.lower()
            if needle in observation.text.lower():
                return Action("done", reason=f"page text contains {self.done_when_text_contains!r}")

        target = self.target or self._target_from(goal)
        if target:
            element = observation.find(target)
            if element is None:
                if self._waits < self.max_waits:
                    self._waits += 1
                    return Action("wait", seconds=self.wait_seconds,
                                  reason=f"{target!r} is not on screen yet")
                raise GoalUnreachable(f"{target!r} never appeared on screen")
            key = element.label.strip().lower()
            if key in self._clicked:
                return Action("done", reason=f"already clicked {element.label!r}; goal assumed met")
            self._clicked.add(key)
            return Action("click", label=element.label,
                          reason=f"matched {target!r} in the element index")
        raise GoalUnreachable(f"no element name found in the goal {goal!r}")

    @staticmethod
    def _target_from(goal: str) -> str | None:
        """Best effort: quoted text wins, else the last capitalised word."""
        quoted = re.search(r"[\"\']([^\"\']{2,40})[\"\']", goal)
        if quoted:
            return quoted.group(1)
        words = [w for w in goal.split() if w[:1].isupper()]
        return words[-1].strip(".,!?") if words else None


class ModelBrain:
    """Shared plumbing for model-backed brains: menu, prompt, lenient parsing.

    Subclasses implement :meth:`_generate`. ``requires_index`` marks brains that
    cannot decide without an element index (a text-only planner): for those, an
    empty screen is a hard stop rather than a prompt.
    """

    name = "model"
    requires_index = False

    def __init__(self, model_id: str, *, max_menu_items: int = 12,
                 max_new_tokens: int = 64) -> None:
        self.model_id = model_id
        self.max_menu_items = max_menu_items
        self.max_new_tokens = max_new_tokens
        self.last_prompt: str | None = None
        self.last_reply: str | None = None
        # "what I already tried", newest last; each entry's outcome is filled in at
        # the next decision, when the new observation can be compared to the old one
        self.history: list[str] = []
        self._last_fingerprint: tuple | None = None

    @staticmethod
    def _fingerprint(observation: Observation) -> tuple:
        return (tuple(element.label for element in observation.elements),
                observation.text[:200])

    def build_prompt(self, goal: str, menu: DecisionMenu, observation: Observation,
                     *, history: Sequence[str] = (), screen_changed: bool = True) -> str:
        """A short prompt with a strict output shape.

        Deliberately contains placeholders (``<number>``) rather than worked
        examples: the weakest models echo the prompt back, and an echoed
        placeholder fails validation instead of being mistaken for a decision.
        """
        lines = [
            "You are controlling a computer. Choose exactly one action.",
            "",
            "Reply with one JSON object and nothing else, in this shape:",
            '  {"kind": "click", "index": <number from the list>, "reason": "<brief>"}',
            '  {"kind": "done", "reason": "<brief>"}',
            "Use done only when the goal is already met.",
            "",
            f"Goal: {goal}",
        ]
        if observation.text:
            lines += ["", f"Text on screen: {observation.text[:300]}"]
        lines += ["", f"Clickable items ({len(menu)}):"]
        lines.append(menu.render() if menu else "  (none detected)")
        if history:
            lines += ["", "Already attempted - do NOT repeat these, they did not work:"]
            lines += [f"  - {entry}" for entry in history]
        if not screen_changed:
            lines += ["", "The screen did not change after your last action. "
                          "Choose a DIFFERENT action."]
        return "\n".join(lines)

    def decide(self, goal: str, observation: Observation) -> Action:
        menu = DecisionMenu(observation, max_items=self.max_menu_items)
        if self.requires_index and not menu:
            raise GoalUnreachable(
                "no clickable elements on screen - is screen capture and OCR working?"
            )
        fingerprint = self._fingerprint(observation)
        screen_changed = (self._last_fingerprint is None
                          or fingerprint != self._last_fingerprint)

        # Finalise the previous attempt first: only now do we know whether it changed
        # anything, and "you already tried this and nothing happened" is the single most
        # useful thing to tell a small model.
        if self.history:
            previous = self.history[-1].split(" -> ")[0]
            self.history[-1] = (
                f"{previous} -> "
                + ("the screen changed" if screen_changed else "the screen did NOT change")
            )

        self.last_prompt = self.build_prompt(
            goal, menu, observation, history=tuple(self.history),
            screen_changed=screen_changed,
        )
        self.last_reply = self._generate(self.last_prompt, observation)
        action = parse_decision(self.last_reply, observation,
                                max_menu_items=self.max_menu_items)
        self.history.append(action.describe())
        del self.history[:-3]
        self._last_fingerprint = fingerprint
        return action

    def _generate(self, prompt: str, observation: Observation) -> str:
        raise NotImplementedError(
            "subclass ModelBrain and implement _generate(prompt, observation)"
        )