#!/usr/bin/env python3
"""Control a DroidDeck build over ADB (docs/agent-control.md)."""

from __future__ import annotations

import argparse
import base64
import gzip
import json
import os
import re
import shutil
import subprocess
import struct
import sys
import tempfile
import time
import uuid
import zlib
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from adb_device import DeviceResolutionError, resolve_serial


# tools/release/variants.txt: the same app also ships as other packages (e.g. com.antutu.benchmark.full).
PACKAGE = os.environ.get("DROIDDECK_PACKAGE", "com.droiddeck.launcher")
AUTHORITY = f"content://{PACKAGE}.agent"
STEAM_CLIENT_APP = 769
REQUEST_CAP = 256 * 1024


class ControlError(RuntimeError):
    def __init__(self, message: str, code: str = "COMMAND_FAILED") -> None:
        super().__init__(message)
        self.code = code


def adb_path(configured: str | None) -> str:
    if configured:
        return configured
    environment_adb = os.environ.get("ADB")
    if environment_adb:
        return environment_adb
    found = shutil.which("adb")
    if found:
        return found
    sdk = os.environ.get("ANDROID_HOME") or os.environ.get("ANDROID_SDK_ROOT")
    if sdk:
        candidate = Path(sdk) / "platform-tools" / "adb"
        if candidate.is_file():
            return str(candidate)
    home_dir = Path.home()
    candidate = home_dir / "Library" / "Android" / "sdk" / "platform-tools" / "adb"
    if candidate.is_file():
        return str(candidate)
    raise ControlError("adb not found; set ADB, ANDROID_HOME, or ANDROID_SDK_ROOT", "ADB_NOT_FOUND")


def run_adb(adb: str, serial: str, *args: str, timeout: float = 30) -> subprocess.CompletedProcess[bytes]:
    try:
        return subprocess.run(
            [adb, "-s", serial, *args], capture_output=True, check=False, timeout=timeout
        )
    except (OSError, subprocess.TimeoutExpired) as exc:
        raise ControlError(str(exc), "ADB_FAILED") from exc


def _content_result(output: str) -> dict:
    marker = "json="
    start = output.find(marker)
    if start < 0:
        detail = output.strip() or "The agent provider returned no result"
        raise ControlError(detail, "INVALID_PROVIDER_RESPONSE")
    raw = output[start + len(marker):].lstrip()
    try:
        response, _ = json.JSONDecoder().raw_decode(raw)
        if not isinstance(response, dict):
            raise ValueError("expected an object")
        return response
    except (json.JSONDecodeError, ValueError) as exc:
        raise ControlError("Could not decode the agent provider response", "INVALID_PROVIDER_RESPONSE") from exc


def encode_request(request: dict) -> str:
    encoded = base64.b64encode(json.dumps(request, separators=(",", ":")).encode()).decode()
    if len(encoded) > REQUEST_CAP:
        raise ControlError("The encoded request exceeds 256 KiB", "INVALID_REQUEST")
    return encoded


def provider_call(adb: str, serial: str, method: str, arg: str | None = None, request: dict | None = None,
                  timeout: float = 30) -> dict:
    command = ["shell", "content", "call", "--uri", AUTHORITY, "--method", method]
    if arg is not None:
        command.extend(["--arg", arg])
    if request is not None:
        encoded = encode_request(request)
        command.extend(["--extra", f"request:s:{encoded}"])
    result = run_adb(adb, serial, *command, timeout=timeout)
    output = (result.stdout + result.stderr).decode(errors="replace")
    response = _content_result(output)
    if result.returncode != 0 and response.get("ok") is not False:
        raise ControlError(output.strip() or "ADB content call failed", "ADB_FAILED")
    return response


def checked(response: dict, default_code: str) -> dict:
    """A provider or guest answer, raised as a ControlError when it carries ok: false."""
    if response.get("ok") is False:
        error = response.get("error", {})
        raise ControlError(error.get("message", default_code), error.get("code", default_code))
    return response


def start_agent_activity(adb: str, serial: str, request: dict) -> None:
    encoded = encode_request(request)
    result = run_adb(
        adb,
        serial,
        "shell",
        "am",
        "start",
        "-W",
        "-n",
        f"{PACKAGE}/com.droiddeck.launcher.agent.AgentStartActivity",
        "--es",
        "request",
        encoded,
        timeout=90,
    )
    output = (result.stdout + result.stderr).decode(errors="replace")
    if result.returncode != 0 or "status: ok" not in output.lower():
        raise ControlError(output.strip() or "The agent start Activity did not launch", "ACTIVITY_START_FAILED")


def confirm_start(adb: str, serial: str, request: dict, previous: dict) -> dict:
    deadline = time.monotonic() + 95
    expected_mode = {"steam": "steam", "desktop": "lxqt", "run": "run"}[request["mode"]]
    while True:
        state = get_state(adb, serial)
        session = state.get("session", {})
        phase = session.get("phase")
        if session.get("mode") == expected_mode and session.get("id") != previous.get("id"):
            return state
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise ControlError(
                f"The agent Activity launched but no new session appeared; phase is {phase}",
                "ACTIVITY_START_FAILED",
            )
        time.sleep(min(0.25, remaining))


def get_state(adb: str, serial: str) -> dict:
    return checked(provider_call(adb, serial, "state"), "STATE_FAILED")


def wait_until(adb: str, serial: str, describe: str, done, timeout: float, poll: float = 1.0) -> dict:
    """Poll state until done(state) is true. A failed session ends the wait at once."""
    deadline = time.monotonic() + timeout
    while True:
        state = get_state(adb, serial)
        session = state.get("session", {})
        if done(state):
            return state
        if session.get("phase") == "FAILED":
            failure = session.get("failure") or {}
            raise ControlError(
                f"Session failed ({failure.get('code', 'FAILED')}): {failure.get('message', 'unknown error')}",
                failure.get("code", "SESSION_FAILED"),
            )
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            focus = session.get("focus") or {}
            raise ControlError(
                f"Timed out waiting for {describe}; phase {session.get('phase')}, focused app {focus.get('focusedApp')}",
                "TIMEOUT",
            )
        time.sleep(min(poll, remaining))


def wait_for_phase(adb: str, serial: str, wanted: str, timeout: float) -> dict:
    if wanted == "FAILED":
        deadline = time.monotonic() + timeout
        while True:
            state = get_state(adb, serial)
            if state.get("session", {}).get("phase") == "FAILED":
                return state
            if time.monotonic() > deadline:
                raise ControlError("Timed out waiting for FAILED", "TIMEOUT")
            time.sleep(1)
    return wait_until(adb, serial, wanted, lambda s: s.get("session", {}).get("phase") == wanted, timeout)


def focused(state: dict, app_id: int | None) -> bool:
    """The game has focus the way a player sees it: gamescope focuses it and it is in the list the
    Steam client orders the base layer from. app_id None means any game (not the client)."""
    focus = state.get("session", {}).get("focus") or {}
    current = focus.get("focusedApp")
    if app_id is None:
        return current not in (None, 0, STEAM_CLIENT_APP)
    return current == app_id and app_id in (focus.get("focusableApps") or [])


def wait_for_artifacts(adb: str, serial: str, state: dict, timeout: float) -> dict:
    deadline = time.monotonic() + timeout
    while True:
        session = state.get("session", {})
        if session.get("artifactsComplete") or not session.get("artifactsAvailable"):
            return state
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise ControlError("Session stopped, but artifact collection did not complete", "ARTIFACT_TIMEOUT")
        time.sleep(min(1.0, remaining))
        state = get_state(adb, serial)


def add_serial(response: dict, serial: str) -> dict:
    response["deviceSerial"] = serial
    return response



def displays(adb: str, serial: str) -> list[dict]:
    """Logical displays with the SurfaceFlinger (physical) id screencap and screenrecord want."""
    result = run_adb(adb, serial, "shell", "dumpsys", "display", timeout=30)
    text = result.stdout.decode(errors="replace")
    devices = dict(re.findall(r'DisplayDeviceInfo\{"([^"]+)": uniqueId="local:(\d+)"', text))
    out = []
    for block in re.split(r"\n  Display (?=\d+:)", text)[1:]:
        logical = int(block.split(":", 1)[0])
        device = re.search(r"mPrimaryDisplayDevice=(.+)", block)
        size = re.search(r"real (\d+) x (\d+)", block)
        name = device.group(1).strip() if device else None
        out.append({
            "display": logical,
            "name": name,
            "physicalId": devices.get(name) if name else None,
            "size": [int(size.group(1)), int(size.group(2))] if size else None,
        })
    return out


def physical_display(adb: str, serial: str, logical: int | None) -> str | None:
    if logical is None:
        return None
    for d in displays(adb, serial):
        if d["display"] == logical:
            if not d["physicalId"]:
                raise ControlError(f"Display {logical} has no physical id to capture", "NO_SUCH_DISPLAY")
            return d["physicalId"]
    raise ControlError(f"No display {logical}; see `droiddeckctl displays`", "NO_SUCH_DISPLAY")


def screenshot(adb: str, serial: str, path: str, display: int | None) -> dict:
    destination = Path(path).expanduser().resolve()
    physical = physical_display(adb, serial, display)
    command = ["exec-out", "screencap", "-p"] + (["-d", physical] if physical else [])
    result = run_adb(adb, serial, *command, timeout=30)
    if result.returncode != 0 or not result.stdout.startswith(b"\x89PNG"):
        detail = result.stderr.decode(errors="replace").strip()
        raise ControlError(detail or "screencap failed", "SCREENSHOT_FAILED")
    try:
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_bytes(result.stdout)
    except OSError as exc:
        raise ControlError(str(exc), "FILE_IO") from exc
    return {"ok": True, "command": "screenshot", "path": str(destination), "display": display}


def record(adb: str, serial: str, path: str, seconds: float, display: int | None) -> dict:
    destination = Path(path).expanduser().resolve()
    remote = "/data/local/tmp/droiddeckctl-record.mp4"
    physical = physical_display(adb, serial, display)
    command = ["shell", "screenrecord", "--time-limit", str(max(1, min(180, int(seconds))))]
    if physical:
        command += ["--display-id", physical]
    result = run_adb(adb, serial, *command, remote, timeout=seconds + 30)
    if result.returncode != 0:
        raise ControlError(result.stderr.decode(errors="replace").strip() or "screenrecord failed", "RECORD_FAILED")
    try:
        destination.parent.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        raise ControlError(str(exc), "FILE_IO") from exc
    pulled = run_adb(adb, serial, "pull", remote, str(destination), timeout=120)
    run_adb(adb, serial, "shell", "rm", "-f", remote)
    if pulled.returncode != 0:
        raise ControlError(pulled.stderr.decode(errors="replace").strip() or "adb pull failed", "RECORD_FAILED")
    return {"ok": True, "command": "record", "path": str(destination), "seconds": seconds, "display": display}


def pacing_summary(text: str, package: str = PACKAGE) -> dict:
    """What reached the panel, from a `dumpsys SurfaceFlinger --timestats -dump`.

    The game is on the compositor's own layer (droiddeck_wayland_game) when it has one, otherwise on
    the session's SurfaceView. SurfaceFlinger keeps one record per layer AND refresh rate, so a panel
    that switches between 60 and 120 Hz mid-capture splits the same layer into several records: they
    are summed. A game's own fps counter does not show any of this - it counts frames handed over, not
    frames shown, and 60 handed over can be 48 on screen with one in four held for two refreshes.
    """
    records = {"game": [], "surface": []}
    for part in re.split(r"^layerName = ", text, flags=re.M)[1:]:
        name, _, body = part.partition("\n")
        kind = ("game" if "droiddeck_wayland_game" in name
                else "surface" if f"SurfaceView[{package}/" in name else None)
        frames = re.search(r"^totalFrames = (\d+)", body, re.M)
        hist = re.search(r"^present2present histogram is as below:\n(.*)", body, re.M)
        if kind and frames and hist:
            records[kind].append((re.sub(r"#\d+$", "", name.strip()), int(frames.group(1)), hist.group(1)))
    chosen = records["game"] or records["surface"]
    histogram: dict[int, int] = {}
    for _, _, hist in chosen:
        for ms, count in re.findall(r"(\d+)ms=(\d+)", hist):
            if int(count):
                histogram[int(ms)] = histogram.get(int(ms), 0) + int(count)
    intervals = sum(histogram.values())
    cadence = max(histogram, key=histogram.get) if histogram else None
    off = intervals - histogram.get(cadence, 0) if histogram else 0
    return {
        "layers": sorted({name for name, _, _ in chosen}),
        "frames": sum(frames for _, frames, _ in chosen),
        "refreshRates": sorted({int(r) for r in re.findall(r"^displayRefreshRate = (\d+)", text, re.M)}),
        "presentToPresentMs": {str(ms): histogram[ms] for ms in sorted(histogram)},
        "cadenceMs": cadence,
        "offCadence": off,
        "offCadencePercent": round(100.0 * off / intervals, 1) if intervals else None,
    }


def pacing(adb: str, serial: str, seconds: float) -> dict:
    """Frames on screen and their intervals over `seconds`, from SurfaceFlinger's timestats."""
    seconds = max(1.0, min(600.0, seconds))
    timestats = ("shell", "dumpsys", "SurfaceFlinger", "--timestats")
    if run_adb(adb, serial, *timestats, "-clear", "-enable").returncode != 0:
        raise ControlError("could not enable SurfaceFlinger timestats", "PACING_FAILED")
    started = time.monotonic()
    time.sleep(seconds)
    dump = run_adb(adb, serial, *timestats, "-dump", timeout=60)
    elapsed = time.monotonic() - started
    run_adb(adb, serial, *timestats, "-disable")
    if dump.returncode != 0:
        raise ControlError(dump.stderr.decode(errors="replace").strip() or "timestats dump failed", "PACING_FAILED")
    summary = pacing_summary(dump.stdout.decode(errors="replace"))
    if not summary["layers"]:
        raise ControlError("No DroidDeck layer drew during the capture: is a session on screen?", "NO_FRAMES")
    summary["fpsOnScreen"] = round(summary["frames"] / elapsed, 1)
    return {"ok": True, "command": "pacing", "seconds": round(elapsed, 1), **summary}



def start_session(adb: str, serial: str, request: dict) -> dict:
    if request.get("mode") == "run":
        require_commands(adb, serial)
    current = get_state(adb, serial)
    if current["session"].get("phase") not in {"IDLE", "FAILED"} or current["session"].get("running"):
        raise ControlError("A session is already starting or running", "SESSION_ACTIVE")
    start_agent_activity(adb, serial, request)
    return confirm_start(adb, serial, request, current["session"])


def require_commands(adb: str, serial: str) -> dict:
    access = checked(provider_call(adb, serial, "access"), "ACCESS_FAILED")
    if access.get("commands") is not True:
        raise ControlError("Debugging tools are off: turn on Setup > Diagnostics > Debugging tools on the device",
                           "AGENT_COMMANDS_DISABLED")
    return access


def input_request(args: argparse.Namespace) -> dict:
    kind = args.input_type
    if kind == "key":
        return {"type": "key", "key": args.key, "action": args.action, "modifiers": args.mod or [], "holdMs": args.hold}
    if kind == "text":
        return {"type": "text", "text": args.text, "commit": args.commit}
    if kind == "tap":
        return {"type": "tap", "x": args.x, "y": args.y, "space": args.space, "holdMs": args.hold}
    if kind == "swipe":
        return {"type": "swipe", "x": args.x, "y": args.y, "x2": args.x2, "y2": args.y2,
                "space": args.space, "durationMs": args.duration}
    if kind == "pad":
        return {"type": "pad", "button": args.button, "action": args.action, "holdMs": args.hold}
    if kind == "stick":
        return {"type": "stick", "stick": args.stick, "x": args.x, "y": args.y, "holdMs": args.hold}
    if kind == "trigger":
        return {"type": "trigger", "trigger": args.trigger, "value": args.value, "holdMs": args.hold}
    raise ControlError(f"Unknown input type {kind}", "INVALID_INPUT")


def xr_ppm_to_png(data: bytes, linear: bool = False) -> bytes:
    header = re.match(rb"P6\n([0-9]+) ([0-9]+)\n255\n", data)
    if not header:
        raise ControlError("Invalid XR image header", "XR_CAPTURE_FAILED")
    width, height = map(int, header.groups())
    pixels = data[header.end():]
    if not (0 < width <= 16384 and 0 < height <= 16384) or len(pixels) != width * height * 3:
        raise ControlError("Invalid XR image dimensions or incomplete pixels", "XR_CAPTURE_FAILED")
    def chunk(kind: bytes, value: bytes) -> bytes:
        return struct.pack('>I', len(value)) + kind + value + struct.pack('>I', zlib.crc32(kind + value))
    stride = width * 3
    rows = b''.join(b'\0' + pixels[start:start + stride] for start in range(0, len(pixels), stride))
    transfer = chunk(b'gAMA', struct.pack('>I', 100000)) if linear else chunk(b'sRGB', b'\0')
    return (b'\x89PNG\r\n\x1a\n' + chunk(b'IHDR', struct.pack('>IIBBBBB', width, height, 8, 2, 0, 0, 0)) + transfer +
            chunk(b'IDAT', zlib.compress(rows)) + chunk(b'IEND', b''))


def xr_capture(adb: str, serial: str, destination: Path, timeout: float) -> dict:
    if not 0 < timeout <= 120:
        raise ControlError("XR capture timeout must be between 0 and 120 seconds", "INVALID_TIMEOUT")
    remote = f'/sdcard/Android/data/{PACKAGE}/files/xr-capture'
    token = uuid.uuid4().hex
    with tempfile.TemporaryDirectory() as temporary:
        request = Path(temporary) / 'request'
        request.write_text(token, encoding='ascii')
        pushed = run_adb(adb, serial, 'push', str(request), remote + '.request.tmp')
        if pushed.returncode:
            raise ControlError(pushed.stderr.decode(errors='replace'), 'XR_CAPTURE_FAILED')
        moved = run_adb(adb, serial, 'shell', 'mv', remote + '.request.tmp', remote + '.request')
        if moved.returncode:
            raise ControlError(moved.stderr.decode(errors='replace'), 'XR_CAPTURE_FAILED')
        deadline = time.monotonic() + timeout
        metadata = None
        while time.monotonic() < deadline:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            try:
                result = run_adb(adb, serial, 'shell', 'cat', remote + '.json', timeout=remaining)
            except ControlError as error:
                if not isinstance(error.__cause__, subprocess.TimeoutExpired):
                    raise
                break
            try:
                answer = json.loads(result.stdout)
                if answer.get('id') == token:
                    metadata = answer
                    break
            except (ValueError, AttributeError):
                pass
            time.sleep(0.2)
        if metadata is None:
            try:
                run_adb(adb, serial, 'shell', 'rm', '-f', remote + '.request', timeout=1)
            except ControlError:
                pass
            raise ControlError('No fresh XR capture. Enter SteamVR using a build with XR capture support.', 'TIMEOUT')
        if not metadata.get('ok'):
            detail = 'No valid XR projection layer is being submitted' if not metadata.get('validLayer') else 'GPU readback or image write failed'
            raise ControlError(detail, 'XR_CAPTURE_FAILED')
        images = {}
        for eye in (('left', 'right', 'source') if metadata.get('sourceImage') else ('left', 'right')):
            pulled = run_adb(adb, serial, 'exec-out', 'gzip', '-c', remote + '-' + eye + '.ppm', timeout=120)
            if pulled.returncode:
                raise ControlError(pulled.stderr.decode(errors='replace'), 'XR_CAPTURE_FAILED')
            try:
                pixels = gzip.decompress(pulled.stdout)
            except (OSError, EOFError, zlib.error) as error:
                raise ControlError('Incomplete compressed XR image', 'XR_CAPTURE_FAILED') from error
            images[eye] = xr_ppm_to_png(pixels, linear=eye == 'source' and metadata.get('sourceColorSpace') == 'linear')
        # Detect overlapping captures before publishing images with another request's metadata.
        check = run_adb(adb, serial, 'shell', 'cat', remote + '.json')
        try:
            unchanged = json.loads(check.stdout).get('id') == token
        except (ValueError, AttributeError):
            unchanged = False
        if not unchanged:
            raise ControlError('XR capture was replaced by another request; retry', 'XR_CAPTURE_FAILED')
    try:
        destination.mkdir(parents=True, exist_ok=True)
        for eye, image in images.items():
            (destination / (eye + '.png')).write_bytes(image)
        (destination / 'capture.json').write_text(json.dumps(metadata, indent=2), encoding='utf-8')
    except OSError as error:
        raise ControlError(str(error), 'FILE_IO') from error
    return {'ok': True, 'command': 'xr-frame', 'deviceSerial': serial, 'capture': metadata,
            'left': str((destination / 'left.png').resolve()), 'right': str((destination / 'right.png').resolve()),
            **({'source': str((destination / 'source.png').resolve())} if 'source' in images else {}),
            'scope': 'App eye textures before headset compositor reprojection; excludes Meta system overlays'}


def execute(args: argparse.Namespace, adb: str, serial: str) -> dict:
    command = args.command
    if command == 'xr-frame':
        return xr_capture(adb, serial, Path(args.path), args.timeout)

    if command == "state":
        return add_serial(get_state(adb, serial), serial)

    if command == "start":
        request = {"mode": args.mode}
        if args.mode == "steam":
            if args.ui == "desktop":
                request["steamUi"] = "desktop"
            if args.url:
                if not args.url.startswith("steam://"):
                    raise ControlError("Steam URL must start with steam://", "INVALID_STEAM_URL")
                request["steamUrl"] = args.url
        if args.reuse:
            current = get_state(adb, serial)["session"]
            wanted = {"steam": "steam", "desktop": "lxqt"}[args.mode]
            if current.get("running") and current.get("mode") == wanted and current.get("phase") not in {"STOPPING", "FAILED"}:
                state = wait_for_phase(adb, serial, "READY", args.timeout) if args.wait else get_state(adb, serial)
                return add_serial(dict(state, reused=True), serial)
        started = start_session(adb, serial, request)
        if args.wait:
            return add_serial(wait_for_phase(adb, serial, "READY", args.timeout), serial)
        session = started["session"]
        return add_serial({"ok": True, "command": "start", "mode": args.mode,
                           "sessionId": session.get("id"), "phase": session.get("phase")}, serial)

    if command == "run":
        program_args = args.arguments[1:] if args.arguments[:1] == ["--"] else args.arguments
        started = start_session(adb, serial, {"mode": "run", "program": args.program, "programArgs": program_args})
        session = started["session"]
        return add_serial({"ok": True, "command": "start", "mode": "run",
                           "sessionId": session.get("id"), "phase": session.get("phase")}, serial)

    if command == "wait":
        if args.phase == "focused":
            if args.app_id is None:
                raise ControlError("wait focused needs an app id", "INVALID_ARGUMENTS")
            return add_serial(wait_until(adb, serial, f"app {args.app_id} to have focus",
                                         lambda s: focused(s, args.app_id), args.timeout, 0.5), serial)
        if args.phase == "game":
            return add_serial(wait_until(adb, serial, "a game to have focus",
                                         lambda s: focused(s, None), args.timeout, 0.5), serial)
        phases = {"ready": "READY", "idle": "IDLE", "failed": "FAILED"}
        return add_serial(wait_for_phase(adb, serial, phases[args.phase], args.timeout), serial)

    if command == "stop":
        current = get_state(adb, serial)
        if current["session"].get("phase") in {"IDLE", "FAILED"}:
            return add_serial(wait_for_artifacts(adb, serial, current, args.artifact_timeout), serial)
        checked(provider_call(adb, serial, "stop"), "STOP_FAILED")
        state = wait_for_phase(adb, serial, "IDLE", args.timeout)
        return add_serial(wait_for_artifacts(adb, serial, state, args.artifact_timeout), serial)

    if command == "resume":
        return add_serial(checked(provider_call(adb, serial, "resume"), "RESUME_FAILED"), serial)

    if command == "launch":
        response = checked(provider_call(adb, serial, "launch", request={"appId": str(args.app_id)}), "LAUNCH_FAILED")
        if args.wait:
            state = wait_until(adb, serial, f"app {args.app_id} to have focus",
                               lambda s: focused(s, args.app_id), args.timeout, 0.5)
            response["focus"] = state["session"].get("focus")
        return add_serial(response, serial)

    if command == "quit":
        response = provider_call(adb, serial, "quit", request={"appId": args.app_id, "timeout": args.timeout},
                                 timeout=args.timeout + 40)
        return add_serial(checked(response, "QUIT_FAILED"), serial)

    if command == "focus":
        return add_serial(checked(provider_call(adb, serial, "focus", timeout=40), "FOCUS_FAILED"), serial)

    if command == "guest":
        argv = args.argv[1:] if args.argv[:1] == ["--"] else args.argv
        if not argv:
            raise ControlError("guest needs a command after --", "INVALID_ARGUMENTS")
        env = {}
        for item in args.env or []:
            key, sep, value = item.partition("=")
            if not sep:
                raise ControlError(f"--env takes KEY=VALUE, not {item}", "INVALID_ARGUMENTS")
            env[key] = value
        request = {"argv": argv, "timeout": args.timeout, "env": env}
        if args.cwd:
            request["cwd"] = args.cwd
        if args.stdin is not None:
            request["stdin"] = args.stdin
        response = checked(provider_call(adb, serial, "guest", request=request, timeout=args.timeout + 40), "GUEST_FAILED")
        if args.check and response.get("exitCode") != 0:
            raise ControlError(f"{argv[0]} exited with {response.get('exitCode')}: "
                               f"{(response.get('stderr') or {}).get('text', '')[-400:]}", "GUEST_EXIT")
        return add_serial(response, serial)

    if command == "cdp":
        if args.cdp_op == "targets":
            request = {"op": "targets"}
        else:
            request = {"op": "eval", "expression": args.expression, "awaitPromise": not args.no_await,
                       "timeout": args.timeout}
            if args.target:
                request["target"] = args.target
        timeout = getattr(args, "timeout", 15)
        return add_serial(checked(provider_call(adb, serial, "cdp", request=request, timeout=timeout + 40),
                                  "CDP_FAILED"), serial)

    if command == "input":
        return add_serial(checked(provider_call(adb, serial, "input", request=input_request(args), timeout=60),
                                  "INPUT_FAILED"), serial)

    if command == "ui":
        if args.ui_op == "dump":
            request = {"op": "dump", "all": args.all}
        else:
            request = {"op": args.ui_op, "index": args.index}
            if args.match:
                request["match"] = args.match
            else:
                request["tag"] = args.tag
            if args.ui_op == "set-text":
                request["text"] = args.text
        return add_serial(checked(provider_call(adb, serial, "ui", request=request), "UI_FAILED"), serial)

    if command == "env":
        if args.env_op == "set":
            request = {"op": "set", "lines": args.lines, "persistent": args.persistent}
        elif args.env_op == "clear":
            request = {"op": "clear", "scope": args.scope}
        else:
            request = {"op": "list"}
        return add_serial(checked(provider_call(adb, serial, "env", request=request), "ENV_FAILED"), serial)

    if command == "override":
        access = require_commands(adb, serial)
        if not re.fullmatch(r"[A-Za-z0-9._-]{1,64}", args.name) or args.name in {".", ".."}:
            raise ControlError("Override names must be simple filenames", "INVALID_NAME")
        source = Path(args.path).expanduser()
        if not source.is_file():
            raise ControlError(f"No such file: {source}", "FILE_IO")
        inbox = access["inbox"]
        run_adb(adb, serial, "shell", "mkdir", "-p", inbox)
        pushed = run_adb(adb, serial, "push", str(source), f"{inbox}/{args.name}", timeout=300)
        if pushed.returncode != 0:
            raise ControlError(pushed.stderr.decode(errors="replace").strip() or "adb push failed", "PUSH_FAILED")
        return add_serial(checked(provider_call(adb, serial, "override", request={"name": args.name}),
                                  "OVERRIDE_FAILED"), serial)

    if command == "prefs":
        if args.prefs_op == "set":
            request = {"op": "set", "file": args.file, "key": args.key, "type": args.type, "value": args.value}
        else:
            request = {"op": "get", "file": args.file}
        return add_serial(checked(provider_call(adb, serial, "prefs", request=request), "PREFS_FAILED"), serial)

    if command == "access":
        return add_serial(checked(provider_call(adb, serial, "access"), "ACCESS_FAILED"), serial)

    if command == "logs":
        state = get_state(adb, serial)
        session = state.get("session", {})
        if not session.get("running") and session.get("artifactsAvailable") and not session.get("artifactsComplete"):
            state = wait_for_artifacts(adb, serial, state, args.timeout)
            session = state.get("session", {})
        remote_dir = session.get("logDir")
        if not remote_dir or not session.get("id") or not session.get("artifactsAvailable"):
            raise ControlError("No session artifact directory is available", "NO_SESSION_LOGS")
        destination = Path(args.destination or session["id"]).expanduser().resolve()
        result = run_adb(adb, serial, "pull", remote_dir, str(destination), timeout=120)
        if result.returncode != 0:
            detail = (result.stderr or result.stdout).decode(errors="replace").strip()
            raise ControlError(detail or "adb pull failed", "LOG_PULL_FAILED")
        return add_serial({"ok": True, "command": "logs", "sessionId": session["id"], "path": str(destination)}, serial)

    if command == "screenshot":
        return add_serial(screenshot(adb, serial, args.path, args.display), serial)

    if command == "record":
        return add_serial(record(adb, serial, args.path, args.seconds, args.display), serial)

    if command == "pacing":
        return add_serial(pacing(adb, serial, args.seconds), serial)

    if command == "wake":
        run_adb(adb, serial, "shell", "input", "keyevent", "KEYCODE_WAKEUP")
        run_adb(adb, serial, "shell", "wm", "dismiss-keyguard")
        power = run_adb(adb, serial, "shell", "dumpsys", "power").stdout.decode(errors="replace")
        awake = re.search(r"mWakefulness=(\w+)", power)
        return add_serial({"ok": True, "command": "wake", "wakefulness": awake.group(1) if awake else None}, serial)

    if command == "displays":
        return add_serial({"ok": True, "command": "displays", "displays": displays(adb, serial)}, serial)

    raise ControlError(f"Unknown command: {command}", "UNKNOWN_COMMAND")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--adb", help="adb executable (defaults to ADB, PATH, or the Android SDK)")
    parser.add_argument("--serial", help="ADB serial; defaults to ADB_SERIAL or ANDROID_SERIAL")
    commands = parser.add_subparsers(dest="command", required=True)

    state = commands.add_parser("state", help="print machine-readable session state")
    state.add_argument("--json", action="store_true", help="accepted for explicit machine-readable usage")

    start = commands.add_parser("start", help="start Steam or the desktop")
    start.add_argument("mode", choices=["steam", "desktop"])
    start.add_argument("--ui", choices=["bigpicture", "desktop"], default="bigpicture")
    start.add_argument("--url", help="steam:// URL to pass to the client")
    start.add_argument("--wait", action="store_true", help="wait for READY after starting")
    start.add_argument("--reuse", action="store_true", help="keep a running session of this mode instead of failing")
    start.add_argument("--timeout", type=float, default=90)

    run = commands.add_parser("run", help="run a guest program under Gamescope")
    run.add_argument("program")
    run.add_argument("arguments", nargs=argparse.REMAINDER)

    wait = commands.add_parser("wait", help="wait for a session phase, or for a game to have focus")
    wait.add_argument("phase", choices=["ready", "idle", "failed", "focused", "game"])
    wait.add_argument("app_id", nargs="?", type=int, help="with focused: the Steam app id")
    wait.add_argument("--timeout", type=float, default=90)

    stop = commands.add_parser("stop", help="stop the session and wait for IDLE")
    stop.add_argument("--timeout", type=float, default=90)
    stop.add_argument("--artifact-timeout", type=float, default=300)
    commands.add_parser("resume", help="resume a suspended session")

    launch = commands.add_parser("launch", help="launch a Steam game in the running Steam session")
    launch.add_argument("app_id", type=int)
    launch.add_argument("--wait", action="store_true", help="wait until the game has focus")
    launch.add_argument("--timeout", type=float, default=180)

    quit_ = commands.add_parser("quit", help="quit a running game (through the client, then signals)")
    quit_.add_argument("app_id", type=int)
    quit_.add_argument("--timeout", type=float, default=15)

    commands.add_parser("focus", help="gamescope's focus and every Xwayland server's windows, fresh")

    guest = commands.add_parser("guest", help="run a command inside the session (debugging tools)")
    guest.add_argument("--timeout", type=float, default=30)
    guest.add_argument("--env", action="append", metavar="KEY=VALUE")
    guest.add_argument("--cwd")
    guest.add_argument("--stdin")
    guest.add_argument("--check", action="store_true", help="fail unless the command exits 0")
    guest.add_argument("argv", nargs=argparse.REMAINDER)

    cdp = commands.add_parser("cdp", help="the Steam client's DevTools (debugging tools)")
    cdp_ops = cdp.add_subparsers(dest="cdp_op", required=True)
    cdp_ops.add_parser("targets", help="list DevTools targets")
    evaluate = cdp_ops.add_parser("eval", help="evaluate JavaScript (default target SharedJSContext)")
    evaluate.add_argument("expression")
    evaluate.add_argument("--target", help="title or URL substring of the target")
    evaluate.add_argument("--no-await", action="store_true", help="do not await a returned promise")
    evaluate.add_argument("--timeout", type=float, default=15)

    inp = commands.add_parser("input", help="keys, text, touch and pad input into the session")
    kinds = inp.add_subparsers(dest="input_type", required=True)
    key = kinds.add_parser("key", help="a key by name (ENTER, ESC, A, F1...) or evdev code")
    key.add_argument("key")
    key.add_argument("--action", choices=["tap", "down", "up"], default="tap")
    key.add_argument("--mod", action="append", help="modifier held around a tap (CTRL, SHIFT, ALT, META)")
    key.add_argument("--hold", type=int, default=40, help="ms between down and up")
    text = kinds.add_parser("text", help="type text (US layout keys; --commit for text-input)")
    text.add_argument("text")
    text.add_argument("--commit", action="store_true")
    tap = kinds.add_parser("tap", help="touch at a point of the session's output")
    tap.add_argument("x", type=float)
    tap.add_argument("y", type=float)
    tap.add_argument("--space", choices=["output", "fraction"], default="output")
    tap.add_argument("--hold", type=int, default=60)
    swipe = kinds.add_parser("swipe", help="drag a touch from one point to another")
    for name in ("x", "y", "x2", "y2"):
        swipe.add_argument(name, type=float)
    swipe.add_argument("--space", choices=["output", "fraction"], default="output")
    swipe.add_argument("--duration", type=int, default=300)
    pad = kinds.add_parser("pad", help="a virtual pad button (a, b, x, y, lb, rb, start, select, guide, qam, up...)")
    pad.add_argument("button")
    pad.add_argument("--action", choices=["tap", "down", "up"], default="tap")
    pad.add_argument("--hold", type=int, default=120)
    stick = kinds.add_parser("stick", help="hold a stick at x, y (-1..1)")
    stick.add_argument("stick", choices=["left", "right"])
    stick.add_argument("x", type=float)
    stick.add_argument("y", type=float)
    stick.add_argument("--hold", type=int, default=300)
    trigger = kinds.add_parser("trigger", help="hold a trigger (0..1)")
    trigger.add_argument("trigger", choices=["left", "right"])
    trigger.add_argument("value", type=float, nargs="?", default=1.0)
    trigger.add_argument("--hold", type=int, default=200)

    ui = commands.add_parser("ui", help="the app's own UI: dump nodes, click or type by test tag")
    ui_ops = ui.add_subparsers(dest="ui_op", required=True)
    dump = ui_ops.add_parser("dump", help="tagged, labelled and clickable nodes of every app window")
    dump.add_argument("--all", action="store_true", help="every node")
    for name in ("click", "set-text"):
        op = ui_ops.add_parser(name)
        op.add_argument("tag", nargs="?", help="test tag, e.g. rail-steam")
        op.add_argument("--match", help="instead of a tag: text or description containing this")
        op.add_argument("--index", type=int, default=0, help="which match, in tree order")
        if name == "set-text":
            op.add_argument("--text", required=True)

    env = commands.add_parser("env", help="session environment lines (debugging tools)")
    env_ops = env.add_subparsers(dest="env_op", required=True)
    env_ops.add_parser("list")
    env_set = env_ops.add_parser("set", help="KEY=VALUE lines for the next session")
    env_set.add_argument("lines", nargs="+")
    env_set.add_argument("--persistent", action="store_true", help="every session until cleared")
    env_clear = env_ops.add_parser("clear")
    env_clear.add_argument("scope", nargs="?", choices=["next", "persistent", "all"], default="next")

    override = commands.add_parser("override", help="put a binary first on the next session's PATH (debugging tools)")
    override.add_argument("name", help="the program name, e.g. gamescope")
    override.add_argument("path", help="the local file to push")

    prefs = commands.add_parser("prefs", help="the app's preferences (debugging tools)")
    prefs_ops = prefs.add_subparsers(dest="prefs_op", required=True)
    prefs_get = prefs_ops.add_parser("get")
    prefs_get.add_argument("file", help="preferences file, e.g. session")
    prefs_set = prefs_ops.add_parser("set")
    prefs_set.add_argument("file")
    prefs_set.add_argument("key")
    prefs_set.add_argument("value", nargs="?")
    prefs_set.add_argument("--type", choices=["bool", "int", "long", "float", "string", "remove"], default="string")

    commands.add_parser("access", help="whether debugging tools are enabled on the device")

    logs = commands.add_parser("logs", help="pull the latest session artifact directory")
    logs.add_argument("which", choices=["latest"])
    logs.add_argument("destination", nargs="?")
    logs.add_argument("--timeout", type=float, default=90)

    shot = commands.add_parser("screenshot", help="save a device screenshot")
    shot.add_argument("path")
    shot.add_argument("--display", type=int, help="logical display id (see displays); default the default display")

    xr = commands.add_parser('xr-frame', help='capture both submitted XR eyes and pose metadata')
    xr.add_argument('path', help='output directory for left.png, right.png and capture.json')
    xr.add_argument('--timeout', type=float, default=15)

    rec = commands.add_parser("record", help="record the screen to an mp4")
    rec.add_argument("path")
    rec.add_argument("--seconds", type=float, default=10)
    rec.add_argument("--display", type=int)

    pace = commands.add_parser("pacing", help="frames that reached the screen and their intervals (SurfaceFlinger)")
    pace.add_argument("--seconds", type=float, default=30)

    commands.add_parser("displays", help="the device's displays and their capture ids")
    commands.add_parser("wake", help="wake the device and dismiss the keyguard")
    return parser


EXIT_ADB = {"ADB_NOT_FOUND", "ADB_FAILED", "INVALID_PROVIDER_RESPONSE"}
EXIT_SESSION = {
    "SESSION_ACTIVE", "NO_ACTIVE_SESSION", "SESSION_NOT_SUSPENDED", "SESSION_FAILED", "GUEST_EXIT",
    "GUEST_START_FAILED", "RUNTIME_INSTALL_FAILED", "DESKTOP_INSTALL_FAILED", "ACTIVITY_START_FAILED",
    "NO_STEAM_SESSION", "SESSION_NOT_READY", "AGENT_UNAVAILABLE", "LAUNCH_REJECTED", "NO_PAD", "NO_FRAMES",
}
EXIT_TIMEOUT = {"TIMEOUT", "ARTIFACT_TIMEOUT"}
EXIT_FILE = {"NO_SESSION_LOGS", "LOG_PULL_FAILED", "SCREENSHOT_FAILED", "FILE_IO", "RECORD_FAILED", "PUSH_FAILED",
             "PACING_FAILED", "XR_CAPTURE_FAILED"}
EXIT_DENIED = {"AGENT_COMMANDS_DISABLED", "PROTECTED_PREFS", "PROTECTED_CONTROL"}


def exit_code(exc: Exception) -> int:
    if isinstance(exc, DeviceResolutionError):
        return 3
    code = getattr(exc, "code", "")
    if code in EXIT_ADB:
        return 3
    if code in EXIT_SESSION:
        return 4
    if code in EXIT_TIMEOUT:
        return 5
    if code in EXIT_FILE:
        return 6
    if code in EXIT_DENIED:
        return 7
    return 2


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    serial = None
    try:
        adb = adb_path(args.adb)
        serial = resolve_serial(adb, args.serial, log=True)
        result = execute(args, adb, serial)
        print(json.dumps(result, indent=2, sort_keys=True))
        return 0
    except (ControlError, DeviceResolutionError) as exc:
        error = {"code": getattr(exc, "code", "DEVICE_RESOLUTION_FAILED"), "message": str(exc)}
        code = exit_code(exc)
    except Exception as exc:
        error = {"code": "INTERNAL_ERROR", "message": f"{type(exc).__name__}: {exc}"}
        code = 2
    response = {"ok": False, "error": error}
    if serial:
        response["deviceSerial"] = serial
    print(json.dumps(response, indent=2, sort_keys=True))
    return code


if __name__ == "__main__":
    raise SystemExit(main())
