#!/usr/bin/env python3
"""Record what an Android screen does, as text, and ask questions of it later.

Why this exists. `adb shell uiautomator dump` is the usual way to read a
screen's layout, and it costs about two seconds a call -- so an animation, a
scroll or a settling layout happens entirely in the gap between two samples,
and the only thing it can answer is what the screen looked like once it had
stopped. This drives a small resident recorder on the device instead (see
UiTrace.java beside this script), which connects to the accessibility service
once and then samples at whatever rate is asked for; 60Hz is comfortable.

The other half is that a screen has a few hundred nodes and a question is
almost always about two or three of them. So recording and reading are separate
steps: `record` saves everything, and `show` selects. A trace can be
re-interrogated with a different selector without touching the device again,
which matters because the interesting question is usually the one you think of
after seeing the first answer.

    ui-trace record -d 3000 --do "tap 'Session settings'" -o /tmp/t.txt
    ui-trace elements /tmp/t.txt          # what is on screen, to pick from
    ui-trace show /tmp/t.txt              # what moved (the default question)
    ui-trace show /tmp/t.txt -m 'Called|ask'   # a timeline for those

Press things by **name**, not by coordinate. `tap 'Save'` resolves the label
against the tree at the moment of the gesture, so it survives anything that
moves the control and fails loudly when the control is genuinely not there.
`tap X Y` still works and is the exception: a coordinate is a position measured
once by hand, and the first thing that moves the control makes the tap land on
whatever now sits there -- which reads exactly like a result.

Nothing here is specific to any app.
"""

import argparse
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

SHARE = Path.home() / ".local/share/ui-trace"
DEVICE_JAR = "/data/local/tmp/uitrace.jar"


def find_adb():
    """The adb to drive the device with, without assuming anything about PATH.

    The one installed beside this script comes first, and that is the point
    rather than a convenience: it is the wrapper that fills in `-s` from the
    checkout you are standing in, and without it a machine with two emulators
    attached answers `failed to get feature set: more than one device`. PATH
    order cannot be relied on for that -- the SDK's own platform-tools sits
    ahead of ~/.local/bin in this machine's ambient PATH, and only a project's
    android-env.sh puts it back.
    """
    sibling = Path(__file__).resolve().parent / "adb"
    if sibling.is_file() and os.access(sibling, os.X_OK):
        return str(sibling)
    for candidate in [
        os.environ.get("ADB"),
        "adb",
        os.environ.get("ANDROID_HOME", "") + "/platform-tools/adb",
        os.environ.get("ANDROID_SDK_ROOT", "") + "/platform-tools/adb",
        str(Path.home() / "Android/Sdk/platform-tools/adb"),
    ]:
        if not candidate or candidate.startswith("/platform-tools"):
            continue
        which = subprocess.run(["sh", "-c", f"command -v {candidate!r}"],
                               capture_output=True, text=True)
        if which.returncode == 0:
            return which.stdout.strip()
    sys.exit("ui-trace: no adb found; set ANDROID_HOME or ADB")


def ensure_jar():
    """Build the device recorder if it is missing or older than its source."""
    jar, source = SHARE / "uitrace.jar", SHARE / "UiTrace.java"
    if not source.exists():
        sys.exit(f"ui-trace: recorder source missing at {source}")
    if not jar.exists() or jar.stat().st_mtime < source.stat().st_mtime:
        print("ui-trace: building the device recorder...", file=sys.stderr)
        subprocess.run([str(SHARE / "build.sh")], check=True,
                       stdout=subprocess.DEVNULL)
    return jar


# --- recording -------------------------------------------------------------

def record(args):
    jar = ensure_jar()
    adb = find_adb()
    target = ["-s", args.serial] if args.serial else []
    pushed = subprocess.run([adb, *target, "push", str(jar), DEVICE_JAR],
                            capture_output=True, text=True)
    if pushed.returncode != 0:
        # Two emulators running is the usual cause, and adb says so plainly; there is
        # nothing to add to its own message except which flag fixes it.
        sys.exit(f"ui-trace: {pushed.stderr.strip() or 'adb push failed'}"
                 f"{'' if target else ' (use -s SERIAL to pick a device)'}")
    remote = f"CLASSPATH={DEVICE_JAR} app_process /system/bin UiTrace"
    remote += f" --duration {args.duration} --interval {args.interval}"
    for step in args.do:
        remote += f" --do {step!r}"
    out = Path(args.out) if args.out else Path(
        tempfile.mkstemp(prefix="ui-trace-", suffix=".txt")[1])
    with out.open("w") as handle:
        result = subprocess.run([adb, *target, "shell", remote], stdout=handle,
                                stderr=subprocess.STDOUT)
    text = out.read_text()
    if result.returncode != 0 or "# error" in text:
        # The error lines rather than the head of the file. An action that failed -- a
        # `tap` whose label is not on screen -- reports itself part way through a trace
        # that may be thousands of frames long, so printing the beginning showed a
        # perfectly ordinary first frame and said nothing about what went wrong.
        problems = [line for line in text.splitlines() if line.startswith(("# error", "# at"))]
        print("\n".join(problems[:40]) if problems else text[:2000], file=sys.stderr)
        sys.exit("ui-trace: the device recorder failed")
    frames, _ = parse(out)
    print(f"ui-trace: {len(frames)} frames over {frames[-1][0] if frames else 0}ms "
          f"-> {out}", file=sys.stderr)
    return out


# --- reading ---------------------------------------------------------------

class Node:
    __slots__ = ("depth", "box", "flags", "cls", "rid", "text", "desc", "label")

    def __init__(self, parts):
        self.depth = int(parts[1])
        self.box = tuple(int(v) for v in parts[2:6])
        self.flags, self.cls, self.rid, self.text, self.desc = parts[6:11]
        self.label = (self.text or self.desc or self.rid.split("/")[-1]
                      or self.cls or "?")

    def haystack(self):
        return "\t".join((self.label, self.rid, self.cls, self.text, self.desc))


def parse(path):
    """A trace file into (frames, actions). A frame is (t_ms, [Node]).

    A missing or unreadable trace exits with one line rather than a stack trace: the
    caller is usually reading this in a terminal or a context window, and a traceback
    says nothing a sentence cannot.
    """
    frames, actions, current = [], [], None
    try:
        text = Path(path).read_text()
    except OSError as problem:
        sys.exit(f"ui-trace: cannot read {path}: {problem.strerror}")
    for line in text.splitlines():
        if line.startswith("F "):
            _, _, at = line.split(" ", 2)
            current = (int(at), [])
            frames.append(current)
        elif line.startswith("N\t"):
            parts = line.split("\t")
            # A label may itself be empty, so pad rather than trusting the count.
            parts += [""] * (11 - len(parts))
            if current:
                current[1].append(Node(parts))
        elif line.startswith("A "):
            at, what = line[2:].split("\t", 1)
            actions.append((int(at), what))
    return frames, actions


def keyed(nodes, pattern):
    """Selected nodes by a stable key: the label, plus an index when repeated."""
    seen, out = {}, {}
    for node in nodes:
        if pattern and not pattern.search(node.haystack()):
            continue
        count = seen.get(node.label, 0)
        seen[node.label] = count + 1
        out[node.label if count == 0 else f"{node.label} #{count + 1}"] = node
    return out


def sizes(frames, pattern):
    """{key: set of (width, height)} -- see [summarise] for why this is tracked."""
    seen = {}
    for _, nodes in frames:
        for key, node in keyed(nodes, pattern).items():
            box = node.box
            seen.setdefault(key, set()).add((box[2] - box[0], box[3] - box[1]))
    return seen


def series(frames, pattern, field):
    """{key: [(t, value or None)]} over every frame, for the chosen field."""
    pick = {"top": lambda b: b[1], "bottom": lambda b: b[3],
            "left": lambda b: b[0], "right": lambda b: b[2],
            "box": lambda b: b}[field]
    tracks, order = {}, []
    for index, (at, nodes) in enumerate(frames):
        present = keyed(nodes, pattern)
        for key, node in present.items():
            if key not in tracks:
                # Backfilled as absent, so every track is one value per frame and
                # a row can be read across by index. An element that appears
                # part-way through is the common case, not the exotic one.
                tracks[key] = [(frames[i][0], None) for i in range(index)]
                order.append(key)
            tracks[key].append((at, pick(node.box)))
        for key in tracks:
            if key not in present:
                tracks[key].append((at, None))
    return order, tracks


def render(value):
    if value is None:
        return "-"
    if isinstance(value, tuple):
        return "{},{}..{},{}".format(*value)
    return str(value)


def cap(rows, limit):
    """The first [limit] rows, saying how many were left out.

    Every listing here shortens the same way, because the caller is usually an
    agent with a context window and a screen has hundreds of nodes.
    """
    if limit <= 0 or len(rows) <= limit:
        return rows
    print(f"  (showing {limit} of {len(rows)}; --limit 0 for all)")
    return rows[:limit]


def cmd_elements(args):
    frames, _ = parse(args.trace)
    counts = {}
    for _, nodes in frames:
        for node in nodes:
            # A node whose only name is its class name names nothing -- a screen has
            # hundreds of those and listing them is the noise this tool exists to cut.
            if not args.all and not (node.text or node.desc or node.rid):
                continue
            entry = counts.setdefault(node.label, [0, node.cls, node.rid, node.flags])
            entry[0] += 1
    ranked = sorted(counts.items(), key=lambda kv: -kv[1][0])
    print(f"{len(counts)} named labels over {len(frames)} frames"
          + ("" if args.all else " (--all also lists unnamed structure)"))
    for label, (count, cls, rid, flags) in cap(ranked, args.limit):
        mark = f" [{flags}]" if flags != "-" else ""
        print(f"  {count:5d}x  {label[:44]:<44} {cls}{mark}"
              + (f"  {rid.split('/')[-1]}" if rid else ""))


def cmd_show(args):
    frames, actions = parse(args.trace)
    if not frames:
        sys.exit("ui-trace: no frames in that trace")
    pattern = re.compile(args.match) if args.match else None
    order, tracks = series(frames, pattern, args.field)
    if not pattern:
        summarise(order, tracks, sizes(frames, pattern), actions, args)
        return
    timeline(frames, order, tracks, actions, args)


def summarise(order, tracks, shapes, actions, args):
    """What moved -- the question a trace is usually recorded to answer.

    Everything that held still is dropped, because a screen is mostly things
    that held still and reading them is the cost this tool exists to avoid.
    """
    for at, what in actions:
        print(f"  action  t={at}ms  {what}")
    moved = []
    for key in order:
        points = [(t, v) for t, v in tracks[key] if v is not None]
        if not points:
            continue
        values = [v for _, v in points]
        if len(set(values)) > 1:
            moved.append((key, points))
    total = len(order)
    print(f"{len(moved)} of {total} elements changed {args.field} during the trace")
    for key, points in cap(moved, args.limit):
        values = [v for _, v in points]
        changes = [(t, v) for i, (t, v) in enumerate(points)
                   if i == 0 or v != points[i - 1][1]]
        span = f"t={changes[1][0]}..{changes[-1][0]}ms" if len(changes) > 1 else ""
        net = values[-1] - values[0] if not isinstance(values[0], tuple) else None
        arrow = f"{render(values[0])} -> {render(values[-1])}"
        delta = f" ({net:+d})" if net is not None else ""
        # A label is only a name, and two different nodes can share one -- a text
        # and the padded control wrapping it both read as their text. When the
        # size changes too, the two readings are probably not the same node, and
        # the move is partly the difference between them rather than travel.
        shape = shapes.get(key, set())
        note = f"  [size changed {len(shape)}x -- may be different nodes]" if len(shape) > 1 else ""
        print(f"  {key[:40]:<40} {arrow}{delta}  {len(changes) - 1} steps  {span}{note}")
    if not moved:
        print("  (nothing moved)")


def timeline(frames, order, tracks, actions, args):
    """A row per change, not a row per frame.

    Consecutive frames that say the same thing are one row with a time range:
    at 60Hz a settled second is sixty identical rows, and the reader wants the
    moment it stopped being settled.
    """
    # "Nothing matched" and "matched but never moved" are different answers, and
    # an empty table says neither of them.
    if not order:
        print(f"  nothing matched /{args.match}/ in any of {len(frames)} frames."
              f" Try `ui-trace elements` to see what is there.")
        return
    if len(order) > args.max_columns:
        print(f"  {len(order)} elements match, which is too many to read across."
              f" Narrowing to the first {args.max_columns}; use a tighter -m,"
              f" or --max-columns to widen.")
        order = order[:args.max_columns]
    width = max(12, max((len(k) for k in order), default=12))
    width = min(width, 24)
    columns = [k[:width] for k in order]
    print("  " + "t (ms)".rjust(13) + " | " + " | ".join(c.ljust(width) for c in columns))
    print("  " + "-" * 13 + "-+-" + "-+-".join("-" * width for _ in columns))
    pending = list(actions)
    previous, start, last = None, None, None
    rows = [[t for t, _ in tracks[order[0]]]] if order else [[]]

    def flush():
        if previous is None:
            return
        label = f"{start}" if start == last else f"{start}-{last}"
        print("  " + label.rjust(13) + " | "
              + " | ".join(render(v).ljust(width) for v in previous))

    rows_out = 0
    for index, (at, _) in enumerate(frames):
        values = tuple(tracks[k][index][1] for k in order)
        while pending and pending[0][0] <= at:
            flush()
            previous = None
            print(f"  >>> {pending[0][1]}  @ t={pending[0][0]}ms")
            pending.pop(0)
        if values != previous:
            flush()
            rows_out += 1
            if args.limit > 0 and rows_out > args.limit:
                print(f"  ... more changes after t={at}ms (--limit 0 for all)")
                return
            previous, start = values, at
        last = at
    flush()
    for at, what in pending:
        print(f"  >>> {what}  @ t={at}ms")


def main():
    parser = argparse.ArgumentParser(
        prog="ui-trace", description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    subs = parser.add_subparsers(dest="command", required=True)

    rec = subs.add_parser("record", help="record a trace from a device")
    rec.add_argument("-s", "--serial", help="adb device serial")
    rec.add_argument("-d", "--duration", type=int, default=3000, help="ms to record")
    rec.add_argument("-i", "--interval", type=int, default=16,
                     help="ms between samples (0 = as fast as possible)")
    rec.add_argument("--do", action="append", default=[], metavar="ACTION",
                     help="'wait MS' | 'tap X Y' | 'swipe X1 Y1 X2 Y2 [MS]', repeatable")
    rec.add_argument("-o", "--out", help="where to save the trace")
    rec.add_argument("-m", "--match", help="show a timeline for these afterwards")
    rec.add_argument("--field", default="top",
                     choices=["top", "bottom", "left", "right", "box"])
    rec.add_argument("--max-columns", type=int, default=6)
    rec.add_argument("--limit", type=int, default=10, help="max rows printed; 0 for all")

    for name, help_text in [("show", "what moved, or a timeline with -m"),
                            ("elements", "every label in a trace, to pick from")]:
        sub = subs.add_parser(name, help=help_text)
        sub.add_argument("trace")
        if name == "show":
            sub.add_argument("-m", "--match", help="regex over label, id, class, text")
            sub.add_argument("--field", default="top",
                             choices=["top", "bottom", "left", "right", "box"])
            sub.add_argument("--max-columns", type=int, default=6)
            sub.add_argument("--limit", type=int, default=10,
                             help="max rows printed; 0 for all")
        else:
            sub.add_argument("--limit", type=int, default=20,
                             help="max rows printed; 0 for all")
            sub.add_argument("--all", action="store_true",
                             help="include nodes with no name of their own")

    args = parser.parse_args()
    if args.command == "record":
        out = record(args)
        args.trace = out
        cmd_show(args)
    elif args.command == "show":
        cmd_show(args)
    else:
        cmd_elements(args)


if __name__ == "__main__":
    main()
