agentic-mobile-control/perception/ui_parser.py

from __future__ import annotations

import re
import xml.etree.ElementTree as ET
from dataclasses import replace
from typing import Any

from core.models import Bounds, SceneElement

_ANDROID_BOUNDS = re.compile(r"\[(?P<x1>-?\d+),(?P<y1>-?\d+)\]\[(?P<x2>-?\d+),(?P<y2>-?\d+)\]")

# Accessibility interaction-state attributes, as literally emitted by
# WDA/XCUITest page_source (enabled) and Appium UiAutomator2 page_source
# (enabled, clickable, selected, checked, focused). iOS page_source does not
# emit selected/checked/focused/clickable, so those stay None for iOS trees.
_STATE_ATTR_KEYS = ("enabled", "clickable", "selected", "checked", "focused")


def parse_ui_tree(raw_tree: Any) -> list[SceneElement]:
    if raw_tree is None:
        return []
    if isinstance(raw_tree, str):
        return _parse_xml(raw_tree)
    if isinstance(raw_tree, dict):
        elements: list[SceneElement] = []
        _parse_dict_node(raw_tree, elements)
        return _with_stable_ids(elements, "ui")
    if isinstance(raw_tree, list):
        elements = []
        for node in raw_tree:
            if isinstance(node, dict):
                _parse_dict_node(node, elements)
        return _with_stable_ids(elements, "ui")
    return []


def _parse_xml(raw_tree: str) -> list[SceneElement]:
    if not raw_tree.strip():
        return []
    root = ET.fromstring(raw_tree)
    elements = []
    for node in root.iter():
        bounds = _bounds_from_attrs(node.attrib)
        if bounds is None or bounds.width <= 0 or bounds.height <= 0:
            continue
        elements.append(
            SceneElement(
                id="",
                type=_normalize_type(_local_name(node.tag), node.attrib),
                text=_text_from_attrs(node.attrib),
                bounds=bounds,
                confidence=1.0,
                source="ui",
                **_state_from_attrs(node.attrib),
            )
        )
    return _with_stable_ids(elements, "ui")


def _parse_dict_node(node: dict[str, Any], elements: list[SceneElement]) -> None:
    attrs = dict(node)
    children = attrs.pop("children", None) or attrs.pop("nodes", None) or []
    bounds = _bounds_from_attrs(attrs)
    if bounds and bounds.width > 0 and bounds.height > 0:
        elements.append(
            SceneElement(
                id="",
                type=_normalize_type(
                    str(attrs.get("type") or attrs.get("class") or "unknown"),
                    attrs,
                ),
                text=_text_from_attrs(attrs),
                bounds=bounds,
                confidence=1.0,
                source="ui",
                **_state_from_attrs(attrs),
            )
        )
    for child in children:
        if isinstance(child, dict):
            _parse_dict_node(child, elements)


def _with_stable_ids(elements: list[SceneElement], prefix: str) -> list[SceneElement]:
    return [
        element if element.id else replace(element, id=f"{prefix}-{index:03d}")
        for index, element in enumerate(elements)
    ]


def _bounds_from_attrs(attrs: dict[str, Any]) -> Bounds | None:
    if all(key in attrs for key in ("x", "y", "width", "height")):
        return Bounds(
            float(attrs["x"]),
            float(attrs["y"]),
            float(attrs["width"]),
            float(attrs["height"]),
        )
    if all(key in attrs for key in ("left", "top", "right", "bottom")):
        left = float(attrs["left"])
        top = float(attrs["top"])
        right = float(attrs["right"])
        bottom = float(attrs["bottom"])
        return Bounds(left, top, right - left, bottom - top)
    raw_bounds = attrs.get("bounds") or attrs.get("rect")
    if isinstance(raw_bounds, dict):
        return _bounds_from_attrs(raw_bounds)
    if isinstance(raw_bounds, str):
        match = _ANDROID_BOUNDS.fullmatch(raw_bounds.strip())
        if match:
            x1 = float(match.group("x1"))
            y1 = float(match.group("y1"))
            x2 = float(match.group("x2"))
            y2 = float(match.group("y2"))
            return Bounds(x1, y1, x2 - x1, y2 - y1)
    return None


def _text_from_attrs(attrs: dict[str, Any]) -> str | None:
    for key in ("label", "name", "text", "value", "placeholder"):
        value = attrs.get(key)
        if value not in (None, ""):
            return str(value)
    return None


def _state_from_attrs(attrs: dict[str, Any]) -> dict[str, bool | None]:
    return {key: _parse_bool(attrs.get(key)) for key in _STATE_ATTR_KEYS}


def _parse_bool(value: Any) -> bool | None:
    if isinstance(value, bool):
        return value
    if value is None:
        return None
    text = str(value).strip().lower()
    if text in ("true", "1", "yes"):
        return True
    if text in ("false", "0", "no"):
        return False
    return None


def _normalize_type(raw_type: str, attrs: dict[str, Any]) -> str:
    candidate = (
        attrs.get("role")
        or attrs.get("type")
        or attrs.get("class")
        or attrs.get("className")
        or raw_type
    )
    normalized = str(candidate).split(".")[-1].replace("XCUIElementType", "").lower()
    if "button" in normalized:
        return "button"
    if "textfield" in normalized or "textarea" in normalized or "input" in normalized:
        return "input"
    if "statictext" in normalized or normalized in {"text", "label"}:
        return "text"
    if "image" in normalized or "icon" in normalized:
        return "image"
    if "cell" in normalized or "row" in normalized:
        return "cell"
    if "window" in normalized:
        return "window"
    return normalized or "unknown"


def _local_name(tag: str) -> str:
    return tag.rsplit("}", 1)[-1]