from __future__ import annotations

from collections import deque
from typing import Any

from data.datasets.android_control_a11y import label_history_frame
from data.datasets.android_control_a11y import render_action_subtree
from data.datasets.android_control_a11y import _outline_image
from rl_harness.android_world_adapter import _clean_text
from rl_harness.types import HarnessFrame
from rl_harness.types import ParsedAction


def _node_name(node: Any | None) -> str:
    if node is None:
        return ""
    for value in (
        getattr(node, "text", ""),
        getattr(node, "content_desc", ""),
        (getattr(node, "resource_id", "") or "").rsplit("/", 1)[-1],
        (getattr(node, "class_name", "") or "").split(".")[-1],
    ):
        value = (value or "").strip()
        if value:
            return value
    return ""


def _node_for_label(frame: HarnessFrame, action: ParsedAction, label_map: dict[int, str] | None):
    if action.label and label_map:
        label_to_idx = {label: idx for idx, label in label_map.items()}
        idx = label_to_idx.get(action.label)
        if idx is not None and 0 <= idx < len(frame.tappable):
            return frame.tappable[idx]
    if action.x is None or action.y is None:
        return None
    best = None
    best_area = None
    for node in frame.tappable:
        bounds = getattr(node, "bounds", None)
        if not bounds:
            continue
        l, t, r, b = bounds
        if l <= action.x <= r and t <= action.y <= b:
            area = max(0, r - l) * max(0, b - t)
            if best_area is None or area < best_area:
                best = node
                best_area = area
    return best


def _history_action(frame: HarnessFrame, action: ParsedAction, label_map: dict[int, str] | None):
    if action.kind in ("tap", "long_press"):
        node = _node_for_label(frame, action, label_map)
        verb = "long pressed" if action.kind == "long_press" else "tapped"
        semantic_verb = "long-pressed" if action.kind == "long_press" else "tapped"
        if node is None:
            return frame.image, verb, semantic_verb
        subtree = render_action_subtree(node)
        name = _node_name(node)
        semantic = f'{semantic_verb} "{name}"' if name else semantic_verb
        return _outline_image(frame.image, node), f"{verb} outlined element:\n{subtree}", semantic

    if action.kind == "type":
        text = _clean_text(action.text, 120)
        action_text = f"type '{text}'" if text else "type"
        semantic = f'typed "{text}"' if text else "typed"
        return frame.image, action_text, semantic

    if action.kind.startswith("scroll_"):
        direction = action.kind.removeprefix("scroll_")
        return frame.image, f"scroll {direction}", f"scrolled {direction}"
    if action.kind == "back":
        return frame.image, "pressed back", "pressed back"
    if action.kind == "home":
        return frame.image, "pressed home", "pressed home"
    if action.kind == "wait":
        return frame.image, "waited", "waited"
    if action.kind == "done":
        return frame.image, "finished", "finished"
    if action.kind == "impossible":
        return frame.image, "marked impossible", "marked impossible"
    return frame.image, action.kind, action.kind


class HistoryBuffer:
    """Maintains model-facing history records across live environment steps."""

    def __init__(self, maxlen: int = 16):
        self.maxlen = maxlen
        self._items = deque(maxlen=maxlen)

    def clear(self):
        self._items.clear()

    def append(self, frame: HarnessFrame, action: ParsedAction, label_map: dict[int, str] | None = None):
        img, action_text, semantic_action = _history_action(frame, action, label_map)
        self._items.append({
            "img": img,
            "action_text": action_text,
            "semantic_action": semantic_action,
            "compressed_tree": frame.compressed_tree,
            "transition_result": "",
        })

    def annotate_last_result(self, result: str) -> None:
        """Attach the next-frame result to the action that produced it."""
        if not self._items:
            return
        if result not in ("changed", "no change"):
            return
        self._items[-1]["transition_result"] = result

    def as_prompt_history(self) -> list[dict[str, Any]]:
        history = []
        for idx, item in enumerate(self._items, start=1):
            row = dict(item)
            row["img"] = label_history_frame(row["img"], f"Step {idx}")
            result = row.get("transition_result")
            if result and " -> " not in (row.get("action_text") or ""):
                row["action_text"] = f"{row.get('action_text') or ''} -> {result}".strip()
                row["semantic_action"] = f"{row.get('semantic_action') or ''} -> {result}".strip()
            history.append(row)
        return history

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