#!/usr/bin/env python3
"""Phase-0 AT-SPI spike for the Nucleic Linux semantic-control agent.

Proves the AT-SPI2 approach that will back the Linux `ax_*` ops (docs/LINUX_VM_AX_AGENT.md):

  1. Read the accessibility tree and emit it as the SAME `node` JSON the macOS agent produces
     (ref/role/subrole/title/value/enabled/focused/frame/actions/children) — so the host side needs
     no protocol change.
  2. Act on controls BY IDENTITY: do_action (press), set_text_contents (set value), grab_focus.
  3. Report both SCREEN and WINDOW extents, so we can confirm the load-bearing assumption behind the
     Wayland->X11 switch: that SCREEN coordinates are real (needed by the pixel fallback + hit-test).

Usage:
    ax_dump_spike.py --app "Nucleic AX Probe" [--act]

Exits non-zero on failure so run_spike.sh can gate CI on it.
"""
import argparse
import json
import sys

import pyatspi

# AT-SPI role names are lowercase-spaced ("push button"); the host is role-string-agnostic (macOS AX
# uses "AXButton"), so we pass the AT-SPI role name through verbatim rather than inventing a mapping.


def _extents(acc, coord):
    try:
        comp = acc.queryComponent()
    except NotImplementedError:
        return None
    try:
        e = comp.getExtents(coord)
        return {"x": e.x, "y": e.y, "w": e.width, "h": e.height}
    except Exception:
        return None


def _actions(acc):
    try:
        act = acc.queryAction()
    except NotImplementedError:
        return []
    return [act.getName(i) for i in range(act.nActions)]


def _value(acc):
    # Prefer editable/text content, then a numeric Value, else the accessible's description-free value.
    for query, attr in (("queryText", "text"),):
        try:
            t = getattr(acc, query)()
            return t.getText(0, -1)
        except NotImplementedError:
            pass
        except Exception:
            pass
    try:
        return acc.queryValue().currentValue
    except NotImplementedError:
        return None
    except Exception:
        return None


def node(acc, reg, depth, max_depth, budget):
    """Serialize one accessible to the macOS-identical `node` shape, assigning a stable `ref`."""
    if budget[0] <= 0:
        return None
    budget[0] -= 1

    ref = "e%d" % len(reg)
    reg[ref] = acc

    try:
        states = acc.getState()
        enabled = states.contains(pyatspi.STATE_ENABLED) or states.contains(pyatspi.STATE_SENSITIVE)
        focused = states.contains(pyatspi.STATE_FOCUSED)
    except Exception:
        enabled, focused = True, False

    n = {
        "ref": ref,
        "role": acc.getRoleName(),
        "title": acc.name or None,
        "value": _value(acc),
        "enabled": bool(enabled),
        "focused": bool(focused),
        "frame": _extents(acc, pyatspi.DESKTOP_COORDS),
        "frameWindow": _extents(acc, pyatspi.WINDOW_COORDS),  # spike-only: compare to catch bad coords
        "actions": _actions(acc),
        "children": [],
    }

    if depth < max_depth:
        try:
            count = acc.childCount
        except Exception:
            count = 0
        for i in range(count):
            try:
                child = acc.getChildAtIndex(i)
            except Exception:
                child = None
            if child is None:
                continue
            c = node(child, reg, depth + 1, max_depth, budget)
            if c is not None:
                n["children"].append(c)
    return n


def find_app(name):
    desktop = pyatspi.Registry.getDesktop(0)
    for i in range(desktop.childCount):
        app = desktop.getChildAtIndex(i)
        if app is not None and (app.name or "") == name:
            return app
    # Fall back to substring match — app names sometimes carry a suffix.
    for i in range(desktop.childCount):
        app = desktop.getChildAtIndex(i)
        if app is not None and name in (app.name or ""):
            return app
    return None


def find_by_title(root, title):
    if (root.name or "") == title:
        return root
    for i in range(root.childCount):
        c = root.getChildAtIndex(i)
        if c is None:
            continue
        hit = find_by_title(c, title)
        if hit is not None:
            return hit
    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--app", required=True, help="target application accessible name")
    ap.add_argument("--max-depth", type=int, default=25)
    ap.add_argument("--budget", type=int, default=800)  # mirror the macOS 800-node cap
    ap.add_argument("--act", action="store_true", help="also exercise do_action/set_text/grab_focus")
    args = ap.parse_args()

    app = find_app(args.app)
    if app is None:
        desktop = pyatspi.Registry.getDesktop(0)
        names = [desktop.getChildAtIndex(i).name for i in range(desktop.childCount)]
        print("FAIL: app %r not on the AT-SPI bus. Visible apps: %r" % (args.app, names), file=sys.stderr)
        return 2

    reg = {}
    tree = node(app, reg, 0, args.max_depth, [args.budget])
    print(json.dumps(tree, indent=2, ensure_ascii=False))
    print("\n[spike] dumped %d nodes" % len(reg), file=sys.stderr)

    # Validate the X11 decision: every on-screen control should have real (non-zero) SCREEN extents.
    bad = []

    # A top-level window can legitimately sit at the origin (esp. headless X with no WM), so only flag
    # *leaf-ish* controls that claim (0,0) — those indicate the SCREEN coordinate space isn't resolving.
    _containers = ("application", "filler", "frame", "window", "panel", "scroll pane")

    def check(n):
        f = n.get("frame")
        if f and f["w"] > 0 and f["h"] > 0 and f["x"] == 0 and f["y"] == 0 and n["role"] not in _containers:
            bad.append(n["ref"])
        for c in n["children"]:
            check(c)

    check(tree)
    if bad:
        print("[spike] WARN: %d nodes report (0,0) SCREEN origin (coords suspect): %r" % (len(bad), bad[:8]), file=sys.stderr)
    else:
        print("[spike] SCREEN coordinates look real (non-zero origins present) — X11 assumption holds", file=sys.stderr)

    if not args.act:
        return 0

    # --- act by identity: the three primitives ax_press / ax_set_value / ax_focus map to ---
    ok = True

    entry = find_by_title(app, "probe-entry")
    if entry is not None:
        try:
            entry.queryComponent().grabFocus()  # ax_focus
            entry.queryEditableText().setTextContents("hello from at-spi")  # ax_set_value
            got = entry.queryText().getText(0, -1)
            print("[spike] set_text -> %r" % got, file=sys.stderr)
            ok = ok and (got == "hello from at-spi")
        except Exception as e:
            print("[spike] FAIL entry actions: %r" % e, file=sys.stderr)
            ok = False
    else:
        print("[spike] FAIL: probe-entry not found", file=sys.stderr)
        ok = False

    button = find_by_title(app, "probe-button")
    status = find_by_title(app, "probe-status")
    if button is not None and status is not None:
        try:
            before = _value(status)
            act = button.queryAction()  # ax_press
            names = [act.getName(i) for i in range(act.nActions)]
            idx = names.index("click") if "click" in names else 0
            act.doAction(idx)
            new_status = _value(status)  # the label's TEXT value is what the click mutates
            print("[spike] pressed button (action %r) -> status text %r (was %r)" % (names[idx], new_status, before), file=sys.stderr)
            ok = ok and (new_status == "clicked")
        except Exception as e:
            print("[spike] FAIL button action: %r" % e, file=sys.stderr)
            ok = False
    else:
        print("[spike] FAIL: probe-button/probe-status not found", file=sys.stderr)
        ok = False

    print("[spike] ACT RESULT: %s" % ("PASS" if ok else "FAIL"), file=sys.stderr)
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())
