宜博 / pixelphone

"""
UI control MCP tools — generic on-screen control primitives for the Android MCP server.

Independent module (E2 / AgentPhone parity). Registered into the existing
FastMCP instance by android_mcp._build_server() via a single line:

    from ui_control_tools import register_ui_control_tools
    register_ui_control_tools(mcp)

These are the generic execution-layer fallbacks for phoneuse.md's plan P3/P4:
when there is no app-specific connector / Accessibility path, the agent can
still drive any screen via raw input injection + screen capture + view dump.

  tap            — input tap x y
  swipe          — input swipe x1 y1 x2 y2 [duration_ms]
  long_press     — input swipe x y x y (long duration) — press-and-hold
  input_text     — input text <text> (into the focused field)
  screenshot     — screencap -p <path>  (returns device path + byte size)
  find_element   — uiautomator dump + match by text / resource-id, return center coords
  click_by_text  — find_element(text=...) then tap its center

Permission model — mirrors get_ui_tree in android_mcp.py:
  * On-device (Termux): commands run via `su -c` (Termux uid lacks INJECT_EVENTS).
  * Off-device (e.g. a Mac driving the phone over USB): commands run via
    `adb shell`, where the shell uid already has the needed permissions, so
    `su` is unnecessary. Detection is automatic via the Termux path check, so
    the exact same module is both deployable on-device and testable from a host.

uiautomator / AccessibilityService conflict (find_element / click_by_text):
  `uiautomator dump` uses a UiAutomation accessibility connection. On this
  device an always-on AccessibilityService (the phoneuse AgentAccessibilityService)
  is bound, and while it is bound UiAutomation.getRootInActiveWindow() returns
  null ("ERROR: null root node returned by UiTestAutomationBridge"). The legacy
  get_ui_tree tool hits the same wall. When _dump_ui detects this, and the env
  var PHONEUSE_UI_AUTODISABLE_A11Y is truthy, it briefly disables enabled
  accessibility services, dumps, and always restores them in a finally block.
  Default is OFF (no auto-toggle) so the event bus is never disturbed silently;
  without it, find_element returns an actionable error naming the conflict.

This module only injects input / captures the screen / reads the view tree. It
holds no allowlist of its own: callers are responsible for not driving payment,
deletion, or sensitive-toggle screens. Keep it to safe surfaces.
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import time
from pathlib import Path
from typing import Any, Optional

# ── shell routing (on-device su  /  off-device adb) ──────────────────────────

_TIMEOUT = 15
_ON_DEVICE = Path("/data/data/com.termux/files/usr").exists()


def _adb_prefix() -> list[str]:
    """`adb shell` invocation, preferring the repo-bundled platform-tools."""
    repo_adb = Path(__file__).resolve().parent / "platform-tools" / "adb"
    adb = str(repo_adb) if repo_adb.exists() else "adb"
    cmd = [adb]
    serial = os.environ.get("ANDROID_SERIAL", "")
    if serial:
        cmd += ["-s", serial]
    cmd.append("shell")
    return cmd


def _run(cmd: str, timeout: int = _TIMEOUT) -> subprocess.CompletedProcess:
    """Run a single shell command on the device.

    On-device (Termux) we need root for input/screencap/uiautomator, so we go
    through `su -c`. Off-device we drive the phone through `adb shell`, whose
    uid already holds those permissions — no `su` needed (and it avoids the
    nested-quoting trap of `adb shell su -c '...'`).
    """
    if _ON_DEVICE:
        argv = ["su", "-c", cmd]
    else:
        argv = _adb_prefix() + [cmd]
    return subprocess.run(argv, capture_output=True, text=True, timeout=timeout)


# ── view-tree dump + parsing ─────────────────────────────────────────────────

_UI_DUMP_PATH = "/sdcard/agent_ui_control_dump.xml"
_NODE_RE = re.compile(r"<node\b[^>]*?/?>")
_ATTR_RE = re.compile(r'(\S+?)="(.*?)"')
_BOUNDS_RE = re.compile(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]")

_A11Y_AUTODISABLE = os.environ.get("PHONEUSE_UI_AUTODISABLE_A11Y", "").lower() in (
    "1", "true", "yes", "on")


def _raw_dump() -> tuple[str, str, bool]:
    """Run uiautomator dump once. Returns (xml, stderr, null_root_seen).

    uiautomator prints "UI hierchary dumped to: <path>" to stdout on success
    (then we cat the file), or "ERROR: null root node ..." to stderr on failure
    (exit code is 0 either way, so we inspect the text, not returncode).
    """
    r = _run(f"uiautomator dump {_UI_DUMP_PATH}; cat {_UI_DUMP_PATH} 2>/dev/null", timeout=12)
    stdout = r.stdout or ""
    stderr = r.stderr or ""
    null_root = "null root" in (stdout + stderr).lower()
    # Drop the "UI hierch(i)ary dumped to: ..." preamble line if present.
    xml = "\n".join(l for l in stdout.splitlines() if "dumped to:" not in l).strip()
    return xml, stderr.strip(), null_root


def _a11y_get_services() -> str:
    r = _run("settings get secure enabled_accessibility_services", timeout=6)
    return (r.stdout or "").strip()


def _a11y_restore(saved: str) -> None:
    if saved and saved.lower() != "null":
        _run(f'settings put secure enabled_accessibility_services "{saved}"', timeout=6)
        _run("settings put secure accessibility_enabled 1", timeout=6)


def _dump_ui() -> tuple[str, str]:
    """Dump the foreground view tree. Returns (xml, error).

    Primary path is uiautomator. If it returns a null root because an
    AccessibilityService is bound (see module docstring), and auto-disable is
    opted into via PHONEUSE_UI_AUTODISABLE_A11Y, briefly unbind the services,
    re-dump, and always restore them.
    """
    try:
        xml, stderr, null_root = _raw_dump()
        if "<hierarchy" in xml:
            return xml, ""
        if null_root and _A11Y_AUTODISABLE:
            saved = _a11y_get_services()
            try:
                _run("settings delete secure enabled_accessibility_services", timeout=6)
                _run("settings put secure accessibility_enabled 0", timeout=6)
                time.sleep(2.0)
                xml2, stderr2, _ = _raw_dump()
            finally:
                _a11y_restore(saved)
            if "<hierarchy" in xml2:
                return xml2, ""
            return "", (stderr2 or "uiautomator dump still failed after a11y auto-disable")
        if null_root:
            return "", ("uiautomator returned null root: an AccessibilityService is bound "
                        "and conflicts with UiAutomation. Set PHONEUSE_UI_AUTODISABLE_A11Y=1 "
                        "to auto-toggle accessibility services around the dump.")
        return "", (stderr[:200] or "uiautomator dump returned empty output")
    except subprocess.TimeoutExpired:
        return "", "uiautomator dump timed out"


def _center(bounds: str) -> Optional[tuple[int, int]]:
    m = _BOUNDS_RE.search(bounds or "")
    if not m:
        return None
    x1, y1, x2, y2 = (int(g) for g in m.groups())
    return (x1 + x2) // 2, (y1 + y2) // 2


def _parse_nodes(xml: str) -> list[dict]:
    """Parse uiautomator XML into a flat list of node attribute dicts."""
    nodes: list[dict] = []
    for tag in _NODE_RE.finditer(xml):
        attrs = dict(_ATTR_RE.findall(tag.group(0)))
        nodes.append(attrs)
    return nodes


def _match_nodes(text: str = "", resource_id: str = "",
                 partial: bool = True) -> tuple[list[dict], str]:
    """Find nodes matching text and/or resource-id. Returns (matches, error).

    Each match dict carries: text, resource_id, content_desc, class, bounds,
    and center {x, y}.
    """
    xml, err = _dump_ui()
    if err:
        return [], err
    matches: list[dict] = []
    for attrs in _parse_nodes(xml):
        node_text = attrs.get("text", "")
        node_desc = attrs.get("content-desc", "")
        node_rid = attrs.get("resource-id", "")
        ok = True
        if text:
            hay = f"{node_text}\x00{node_desc}"
            ok = ok and ((text in hay) if partial else (text in (node_text, node_desc)))
        if resource_id:
            ok = ok and ((resource_id in node_rid) if partial else (node_rid == resource_id))
        if not ok or (not node_text and not node_desc and not text and not resource_id):
            continue
        c = _center(attrs.get("bounds", ""))
        matches.append({
            "text": node_text,
            "resource_id": node_rid,
            "content_desc": node_desc,
            "class": attrs.get("class", ""),
            "clickable": attrs.get("clickable", ""),
            "bounds": attrs.get("bounds", ""),
            "center": {"x": c[0], "y": c[1]} if c else None,
        })
    return matches, ""


# ── core primitives (testable without an MCP instance) ───────────────────────

def _do_tap(x: int, y: int) -> dict:
    try:
        r = _run(f"input tap {int(x)} {int(y)}")
        return {"ok": r.returncode == 0, "action": "tap", "x": int(x), "y": int(y),
                "stderr": (r.stderr or "").strip()[:160]}
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": "tap timed out"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def _do_swipe(x1: int, y1: int, x2: int, y2: int, duration_ms: int = 300) -> dict:
    try:
        r = _run(f"input swipe {int(x1)} {int(y1)} {int(x2)} {int(y2)} {int(duration_ms)}")
        return {"ok": r.returncode == 0, "action": "swipe",
                "from": {"x": int(x1), "y": int(y1)}, "to": {"x": int(x2), "y": int(y2)},
                "duration_ms": int(duration_ms), "stderr": (r.stderr or "").strip()[:160]}
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": "swipe timed out"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def _do_long_press(x: int, y: int, duration_ms: int = 800) -> dict:
    # A press-and-hold is a zero-distance swipe with a long duration.
    try:
        d = max(int(duration_ms), 500)
        r = _run(f"input swipe {int(x)} {int(y)} {int(x)} {int(y)} {d}")
        return {"ok": r.returncode == 0, "action": "long_press", "x": int(x), "y": int(y),
                "duration_ms": d, "stderr": (r.stderr or "").strip()[:160]}
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": "long_press timed out"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def _escape_input_text(text: str) -> str:
    """Escape a string for Android `input text`.

    `input text` treats spaces as token separators, so they must be sent as
    %s; a handful of shell-significant characters are backslash-escaped.
    """
    out = text
    for ch in ("\\", '"', "'", "&", "<", ">", "(", ")", ";", "|", "$", "`"):
        out = out.replace(ch, "\\" + ch)
    out = out.replace(" ", "%s")
    return out


def _do_input_text(text: str) -> dict:
    try:
        safe = _escape_input_text(text)
        r = _run(f"input text {safe}")
        return {"ok": r.returncode == 0, "action": "input_text",
                "chars": len(text), "stderr": (r.stderr or "").strip()[:160]}
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": "input_text timed out"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def _do_screenshot(path: str = "") -> dict:
    try:
        if not path:
            path = f"/sdcard/agent_screenshot_{int(time.time())}.png"
        r = _run(f"screencap -p {path}", timeout=20)
        if r.returncode != 0:
            return {"ok": False, "path": path, "error": (r.stderr or "screencap failed").strip()[:200]}
        # Read back the byte size on-device (toybox stat -c %s).
        s = _run(f"stat -c %s {path} 2>/dev/null || wc -c < {path}", timeout=8)
        size_txt = (s.stdout or "").strip().split()[-1] if s.stdout.strip() else ""
        try:
            size_bytes = int(size_txt)
        except ValueError:
            size_bytes = None
        return {"ok": size_bytes not in (None, 0), "path": path, "size_bytes": size_bytes}
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": "screenshot timed out"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def _do_find_element(text: str = "", resource_id: str = "", partial: bool = True,
                     limit: int = 20) -> dict:
    if not text and not resource_id:
        return {"ok": False, "error": "provide text and/or resource_id to match"}
    matches, err = _match_nodes(text=text, resource_id=resource_id, partial=partial)
    if err:
        return {"ok": False, "error": err}
    matches = matches[:max(1, limit)]
    return {"ok": bool(matches), "query": {"text": text, "resource_id": resource_id,
            "partial": partial}, "count": len(matches), "matches": matches}


def _do_click_by_text(text: str, partial: bool = True, index: int = 0) -> dict:
    matches, err = _match_nodes(text=text, partial=partial)
    if err:
        return {"ok": False, "error": err}
    if not matches:
        return {"ok": False, "error": f"no element matched text {text!r}"}
    if index < 0 or index >= len(matches):
        index = 0
    target = matches[index]
    c = target.get("center")
    if not c:
        return {"ok": False, "error": "matched element has no parseable bounds",
                "match": target}
    tap = _do_tap(c["x"], c["y"])
    return {"ok": tap.get("ok", False), "action": "click_by_text", "text": text,
            "tapped": c, "match": target, "matched_total": len(matches)}


# ── tool registration ────────────────────────────────────────────────────────

def register_ui_control_tools(mcp) -> None:
    """Register the generic UI-control primitives onto an existing FastMCP."""

    @mcp.tool()
    def tap(x: int, y: int) -> str:
        """Tap the screen at absolute pixel coordinates (x, y).

        Generic execution-layer fallback (phoneuse.md plan P3/P4). Use
        find_element / get_ui_tree first to locate the coordinates. Do not use
        on payment, deletion, or sensitive-setting screens.

        Args:
            x: horizontal pixel coordinate.
            y: vertical pixel coordinate.
        """
        return json.dumps(_do_tap(x, y), ensure_ascii=False)

    @mcp.tool()
    def swipe(x1: int, y1: int, x2: int, y2: int, duration_ms: int = 300) -> str:
        """Swipe/drag from (x1, y1) to (x2, y2). Good for scrolling and paging.

        Args:
            x1: start x. y1: start y. x2: end x. y2: end y.
            duration_ms: gesture duration in milliseconds (default 300).
        """
        return json.dumps(_do_swipe(x1, y1, x2, y2, duration_ms), ensure_ascii=False)

    @mcp.tool()
    def long_press(x: int, y: int, duration_ms: int = 800) -> str:
        """Press and hold at (x, y) for duration_ms (min 500ms).

        Args:
            x: horizontal pixel coordinate.
            y: vertical pixel coordinate.
            duration_ms: hold duration in milliseconds (default 800, min 500).
        """
        return json.dumps(_do_long_press(x, y, duration_ms), ensure_ascii=False)

    @mcp.tool()
    def input_text(text: str) -> str:
        """Type text into the currently focused input field.

        Tap/find the target field first so it has focus. Spaces and shell-
        significant characters are escaped automatically.

        Args:
            text: the text to type.
        """
        return json.dumps(_do_input_text(text), ensure_ascii=False)

    @mcp.tool()
    def screenshot(path: str = "") -> str:
        """Capture the current screen to a PNG on the device and report its size.

        Args:
            path: device path to write the PNG to. Defaults to
                  /sdcard/agent_screenshot_<unixts>.png.
        Returns JSON with the saved path and the file's byte size.
        """
        return json.dumps(_do_screenshot(path), ensure_ascii=False)

    @mcp.tool()
    def find_element(text: str = "", resource_id: str = "", partial: bool = True,
                     limit: int = 20) -> str:
        """Find on-screen elements by visible text and/or resource-id.

        Dumps the foreground view tree (uiautomator) and returns matching nodes
        with their bounds and tap-ready center coordinates.

        Args:
            text: match against a node's text or content-desc.
            resource_id: match against a node's resource-id.
            partial: substring match (default True) vs exact match.
            limit: max matches to return (default 20).
        """
        return json.dumps(_do_find_element(text, resource_id, partial, limit),
                          ensure_ascii=False)

    @mcp.tool()
    def click_by_text(text: str, partial: bool = True, index: int = 0) -> str:
        """Find an element by its visible text and tap its center.

        Convenience combiner over find_element + tap. Do not use on payment,
        deletion, or sensitive-setting screens.

        Args:
            text: visible text (or content-desc) to locate and tap.
            partial: substring match (default True) vs exact match.
            index: which match to tap if several share the text (default 0).
        """
        return json.dumps(_do_click_by_text(text, partial, index), ensure_ascii=False)