action_mapping.py

"""Safe translation from OpenEnv browser actions to BrowserGym action strings."""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Sequence

try:
    from .models import BrowserAction, BrowserElement
except ImportError:  # pragma: no cover - direct repo execution
    from models import BrowserAction, BrowserElement


class ActionTranslationError(ValueError):
    """Raised when a structured action cannot be safely translated."""


CLICKABLE_ROLES = {"button", "link", "menuitem", "checkbox", "radio", "tab", "option"}
CLICKABLE_TAGS = {"button", "a", "option", "summary", "label"}
INPUT_ROLES = {"textbox", "searchbox", "combobox", "textarea"}
INPUT_TAGS = {"input", "textarea"}
SELECT_ROLES = {"combobox", "listbox"}
SELECT_TAGS = {"select"}
SUBMIT_KEYWORDS = ("submit", "send", "ok", "done", "go", "next", "continue", "confirm")
STOPWORDS = {
    "a",
    "an",
    "and",
    "at",
    "be",
    "button",
    "by",
    "for",
    "from",
    "in",
    "into",
    "of",
    "on",
    "or",
    "page",
    "press",
    "the",
    "then",
    "this",
    "to",
    "with",
}


@dataclass
class TranslationResult:
    browsergym_action: str
    resolved_action: BrowserAction
    metadata: Dict[str, Any]


def _quote(value: str) -> str:
    return repr(value)


def find_submit_target(elements: Iterable[BrowserElement]) -> Optional[str]:
    """Find a likely submit control in the compact element list."""

    for element in elements:
        if not element.visible or not element.enabled:
            continue
        if not _is_clickable(element):
            continue
        haystack = _element_blob(element).lower()
        if any(keyword in haystack for keyword in SUBMIT_KEYWORDS):
            return element.id
    return None


def translate_browser_action(
    action: BrowserAction,
    elements: Iterable[BrowserElement] = (),
    *,
    instruction: str = "",
    history: Sequence[Dict[str, Any]] = (),
) -> TranslationResult:
    """Convert a safe OpenEnv action into BrowserGym's action language with target fallback."""

    element_list = [element for element in elements if not element.id.startswith("distractor")]
    metadata: Dict[str, Any] = {
        "mode": "direct",
        "requested_target_id": action.target_id,
        "resolved_target_id": action.target_id,
        "candidate_count": len(element_list),
    }

    if action.action_type == "noop":
        return TranslationResult("noop()", action, metadata)
    if action.action_type in {"click", "type", "clear", "select"}:
        target_id, resolution = _resolve_target(action, element_list, instruction=instruction, history=history)
        resolved_action = _copy_action(action, target_id=target_id)
        metadata.update(resolution)
        metadata["resolved_target_id"] = target_id
        if action.action_type == "click":
            return TranslationResult(f"click({_quote(target_id)})", resolved_action, metadata)
        if action.action_type == "type":
            return TranslationResult(
                f"fill({_quote(target_id)}, {_quote(action.text or '')})",
                resolved_action,
                metadata,
            )
        if action.action_type == "clear":
            return TranslationResult(f"clear({_quote(target_id)})", resolved_action, metadata)
        return TranslationResult(
            f"select_option({_quote(target_id)}, {_quote(action.text or '')})",
            resolved_action,
            metadata,
        )
    if action.action_type == "submit":
        if action.target_id:
            try:
                target, resolution = _resolve_target(action, element_list, instruction=instruction, history=history)
            except ActionTranslationError:
                target = find_submit_target(element_list)
                resolution = {"mode": "submit_fallback", "reason": "requested_submit_target_unresolved"}
            else:
                metadata.update(resolution)
            if target:
                resolved_action = _copy_action(action, target_id=target)
                metadata["resolved_target_id"] = target
                metadata.setdefault("mode", "submit_target")
                return TranslationResult(f"click({_quote(target)})", resolved_action, metadata)
        target = find_submit_target(element_list)
        if target:
            resolved_action = _copy_action(action, target_id=target)
            metadata.update(
                {
                    "mode": "submit_target",
                    "reason": "submit_like_element",
                    "resolved_target_id": target,
                }
            )
            return TranslationResult(f"click({_quote(target)})", resolved_action, metadata)
        metadata.update({"mode": "keyboard_submit", "reason": "no_submit_target"})
        return TranslationResult("keyboard_press('Enter')", action, metadata)
    if action.action_type == "scroll":
        return TranslationResult(
            f"scroll({int(action.scroll_dx)}, {int(action.scroll_dy)})",
            action,
            metadata,
        )
    if action.action_type == "ask_oracle":
        raise ActionTranslationError("ask_oracle must be resolved before BrowserGym execution")
    raise ActionTranslationError(f"unsupported action_type: {action.action_type}")


def to_browsergym_action(action: BrowserAction, elements: Iterable[BrowserElement] = ()) -> str:
    """Backward-compatible action translation helper."""

    return translate_browser_action(action, elements).browsergym_action


def _resolve_target(
    action: BrowserAction,
    elements: Sequence[BrowserElement],
    *,
    instruction: str,
    history: Sequence[Dict[str, Any]],
) -> tuple[str, Dict[str, Any]]:
    element_ids = {element.id for element in elements}
    if action.target_id and action.target_id in element_ids:
        return action.target_id, {"mode": "exact", "reason": "target_id_match"}

    candidates = _candidate_pool(action, elements)
    if not candidates:
        raise ActionTranslationError(f"{action.action_type} requires a valid target candidate")

    context_text = " ".join(
        bit
        for bit in (
            instruction,
            action.target_id or "",
            action.text or "",
            action.reasoning or "",
        )
        if bit
    )
    context_tokens = _tokenize(context_text)
    quoted_phrases = _quoted_phrases(context_text)
    scored: List[tuple[float, BrowserElement, List[str]]] = []
    for element in candidates:
        score, reasons = _score_candidate(
            action,
            element,
            context_tokens=context_tokens,
            quoted_phrases=quoted_phrases,
            history=history,
        )
        scored.append((score, element, reasons))

    scored.sort(key=lambda item: item[0], reverse=True)
    best_score, best_element, reasons = scored[0]
    runner_up_score = scored[1][0] if len(scored) > 1 else float("-inf")
    if best_score < 0.9:
        raise ActionTranslationError(f"{action.action_type} target could not be resolved safely")
    if len(scored) > 1 and best_score - runner_up_score < 0.2:
        raise ActionTranslationError(f"{action.action_type} target resolution ambiguous")
    return best_element.id, {
        "mode": "semantic",
        "reason": reasons[:4],
        "top_score": round(best_score, 4),
        "runner_up_score": round(runner_up_score, 4) if runner_up_score != float("-inf") else None,
    }


def _candidate_pool(action: BrowserAction, elements: Sequence[BrowserElement]) -> List[BrowserElement]:
    visible_enabled = [element for element in elements if element.visible and element.enabled]

    if action.action_type in {"type", "clear"}:
        preferred = [element for element in visible_enabled if _is_input_like(element)]
        if preferred:
            return preferred
        fallback = [element for element in elements if _is_input_like(element)]
        return fallback or visible_enabled or list(elements)
    if action.action_type == "select":
        preferred = [element for element in visible_enabled if _is_select_like(element)]
        if preferred:
            return preferred
        fallback = [element for element in visible_enabled if _is_clickable(element)]
        if fallback:
            return fallback
        select_like_all = [element for element in elements if _is_select_like(element)]
        clickable_all = [element for element in elements if _is_clickable(element)]
        return select_like_all or clickable_all or visible_enabled or list(elements)
    if action.action_type in {"click", "submit"}:
        preferred = [element for element in visible_enabled if _is_clickable(element)]
        if preferred:
            return preferred
        fallback = [element for element in visible_enabled if _is_interactable(element)]
        if fallback:
            return fallback
        clickable_all = [element for element in elements if _is_clickable(element)]
        interactable_all = [element for element in elements if _is_interactable(element)]
        return clickable_all or interactable_all or visible_enabled or list(elements)
    return visible_enabled or list(elements)


def _score_candidate(
    action: BrowserAction,
    element: BrowserElement,
    *,
    context_tokens: Sequence[str],
    quoted_phrases: Sequence[str],
    history: Sequence[Dict[str, Any]],
) -> tuple[float, List[str]]:
    blob = _element_blob(element).lower()
    score = 0.0
    reasons: List[str] = []

    overlap = 0
    for token in context_tokens:
        if token in blob:
            overlap += 1
    if overlap:
        score += float(overlap)
        reasons.append(f"overlap:{overlap}")

    for phrase in quoted_phrases:
        if phrase and phrase in blob:
            score += 2.0
            reasons.append("quoted_phrase")

    semantic_hints = _semantic_hints(element)
    if action.action_type in {"click", "submit"} and _is_clickable(element):
        score += 1.2
        reasons.append("clickable")
    if action.action_type in {"type", "clear"} and _is_input_like(element):
        score += 1.4
        reasons.append("input_like")
    if action.action_type == "select" and _is_select_like(element):
        score += 1.4
        reasons.append("select_like")
    if "submit_like" in semantic_hints and action.action_type in {"click", "submit"}:
        score += 1.0
        reasons.append("submit_like")

    rank_score = element.attributes.get("rank_score")
    try:
        score += max(0.0, min(float(rank_score), 8.0)) * 0.1
        if rank_score is not None:
            reasons.append("observation_rank")
    except (TypeError, ValueError):
        pass

    if element.visible:
        score += 0.3
    else:
        score -= 1.0
        reasons.append("not_visible")
    if element.enabled:
        score += 0.3
    else:
        score -= 1.0
        reasons.append("not_enabled")

    history_penalty = _history_penalty(element, history)
    if history_penalty:
        score -= history_penalty
        reasons.append("history_penalty")

    if not (element.text or "").strip() and not _is_interactable(element):
        score -= 0.5
        reasons.append("empty_noninteractive")

    return score, reasons


def _history_penalty(element: BrowserElement, history: Sequence[Dict[str, Any]]) -> float:
    text_lc = (element.text or "").strip().lower()
    for item in history[-4:]:
        target_id = str(item.get("target_id") or "")
        if target_id and target_id == element.id:
            return 0.8
        history_text = " ".join(
            str(item.get(key) or "")
            for key in ("action_type", "target_id", "text", "status", "browsergym_action")
        ).lower()
        if text_lc and len(text_lc) >= 3 and text_lc in history_text:
            return 0.5
    return 0.0


def _copy_action(action: BrowserAction, **updates: Any) -> BrowserAction:
    if hasattr(action, "model_copy"):
        return action.model_copy(update=updates)
    return action.copy(update=updates)


def _element_blob(element: BrowserElement) -> str:
    attrs = element.attributes or {}
    attr_bits = []
    for key in ("name", "aria-label", "placeholder", "value", "title", "alt"):
        value = attrs.get(key)
        if value not in (None, ""):
            attr_bits.append(str(value))
    return " ".join(
        bit
        for bit in (
            element.text,
            element.role,
            element.tag,
            element.type,
            " ".join(attr_bits),
        )
        if bit
    )


def _semantic_hints(element: BrowserElement) -> List[str]:
    hints: List[str] = []
    haystack = _element_blob(element).lower()
    if _is_clickable(element):
        hints.append("clickable")
    if _is_input_like(element):
        hints.append("input_like")
    if _is_select_like(element):
        hints.append("select_like")
    if any(keyword in haystack for keyword in SUBMIT_KEYWORDS):
        hints.append("submit_like")
    return hints


def _tokenize(text: str) -> List[str]:
    tokens = []
    for token in re.findall(r"[a-z0-9]+", (text or "").lower()):
        if len(token) < 2 or token in STOPWORDS:
            continue
        tokens.append(token)
    return tokens


def _quoted_phrases(text: str) -> List[str]:
    phrases = []
    for pair in re.findall(r'"([^"]+)"|\'([^\']+)\'', text or ""):
        phrase = (pair[0] or pair[1] or "").strip().lower()
        if phrase:
            phrases.append(phrase)
    return phrases


def _is_interactable(element: BrowserElement) -> bool:
    return _is_clickable(element) or _is_input_like(element) or _is_select_like(element)


def _is_clickable(element: BrowserElement) -> bool:
    if bool((element.attributes or {}).get("clickable")):
        return True
    role = (element.role or "").lower()
    tag = (element.tag or "").lower()
    return role in CLICKABLE_ROLES or tag in CLICKABLE_TAGS


def _is_input_like(element: BrowserElement) -> bool:
    role = (element.role or "").lower()
    tag = (element.tag or "").lower()
    type_name = (element.type or "").lower()
    return role in INPUT_ROLES or tag in INPUT_TAGS or type_name in {"text", "search", "email", "password"}


def _is_select_like(element: BrowserElement) -> bool:
    role = (element.role or "").lower()
    tag = (element.tag or "").lower()
    return role in SELECT_ROLES or tag in SELECT_TAGS