#!/usr/bin/env python3
"""
kpxc-search: put a query into the search box of an already running KeePassXC
via AT-SPI (the Linux accessibility bus). Works on X11 and Wayland.

Usage:  kpxc-search 'u:alice url:ssh://db01.example.com' [--timeout 60] [--debug]

Requirements:
  - PyGObject + the Atspi typelib (at-spi2-core)
  - KeePassXC started with QT_LINUX_ACCESSIBILITY_ALWAYS_ON=1

Speed: every AT-SPI property read is a D-Bus round trip, so walking the widget
tree is the slow part. After the first successful run the script remembers the
search box's position in the tree (in $XDG_RUNTIME_DIR/kpxc-search.json) and
goes straight there next time, verifying it is still the right widget. It falls
back to the full walk only if that check fails (layout changed, new version...).

Exit codes: 0 = query set, 1 = KeePassXC not found, 2 = search box not found
(e.g. database still locked and --timeout expired).
"""
import argparse
import json
import os
import sys
import time
from pathlib import Path

import gi

gi.require_version("Atspi", "2.0")
from gi.repository import Atspi, GLib  # noqa: E402

# Never descend into these: the entry/group views expose one accessible object
# per cell, and walking them over D-Bus is what makes a naive search crawl.
PRUNE_ROLES = {
    Atspi.Role.TABLE,
    Atspi.Role.TREE_TABLE,
    Atspi.Role.TREE,
    Atspi.Role.LIST,
    Atspi.Role.LIST_BOX,
    Atspi.Role.TABLE_CELL,
    Atspi.Role.TREE_ITEM,
    Atspi.Role.LIST_ITEM,
    Atspi.Role.MENU_BAR,
    Atspi.Role.MENU,
}

# Qt reports a QLineEdit as ENTRY on some versions and as TEXT on others.
TEXT_ROLES = {Atspi.Role.ENTRY, Atspi.Role.TEXT}

CACHE_FILE = (
    Path(os.environ.get("XDG_RUNTIME_DIR") or f"/tmp/kpxc-search-{os.getuid()}")
    / "kpxc-search.json"
)

DEBUG = False


def log(msg_fn):
    # Takes a callable so the D-Bus calls inside debug messages only happen
    # when --debug is on.
    if DEBUG:
        print(msg_fn(), file=sys.stderr)


def safe(fn, default=None):
    try:
        return fn()
    except GLib.Error:
        return default


def child_at(node, i):
    return safe(lambda: node.get_child_at_index(i))


def children(node):
    n = safe(node.get_child_count, 0) or 0
    for i in range(n):
        child = child_at(node, i)
        if child is not None:
            yield i, child


def is_line_edit(role, states):
    return (
        role in TEXT_ROLES
        and states is not None
        and states.contains(Atspi.StateType.SHOWING)
        and states.contains(Atspi.StateType.EDITABLE)
        # A disabled (greyed-out) field still reports EDITABLE, e.g. the
        # search box while the database is locked.
        and states.contains(Atspi.StateType.ENABLED)
        and states.contains(Atspi.StateType.SENSITIVE)
        and not states.contains(Atspi.StateType.MULTI_LINE)
    )


def walk(node, ancestors=(), path=(), depth=0, max_depth=30):
    """Yield (node, role, states, ancestor_roles, index_path) over the visible,
    non-pruned part of the tree. Each node costs two D-Bus calls."""
    if node is None or depth > max_depth:
        return
    role = safe(node.get_role)
    states = None
    if depth > 0:
        states = safe(node.get_state_set)
        if states is None or not states.contains(Atspi.StateType.SHOWING):
            return
    log(lambda: f"{'  ' * depth}{safe(node.get_role_name, '?')} "
                f"{(safe(node.get_name, '') or '')!r}")
    yield node, role, states, ancestors, path
    if role in PRUNE_ROLES:
        return
    for i, child in children(node):
        yield from walk(child, ancestors + (role,), path + (i,), depth + 1,
                        max_depth)


def find_keepassxc():
    desktop = Atspi.get_desktop(0)
    for _, app in children(desktop):
        if "keepassxc" in (safe(app.get_name, "") or "").lower():
            return app
    return None


def full_search(app):
    """Slow path: walk the tree. Returns (node, cache_record) or (None, None)."""
    candidates = []
    for node, role, states, ancestors, path in walk(app):
        if not is_line_edit(role, states):
            continue
        name = safe(node.get_name, "") or ""
        desc = safe(node.get_description, "") or ""
        in_toolbar = Atspi.Role.TOOL_BAR in ancestors
        log(lambda: f">> candidate: name={name!r} desc={desc!r} "
                    f"in_toolbar={in_toolbar} path={list(path)}")
        record = {
            "path": list(path),
            "roles": [int(r) for r in ancestors] + [int(role)],
            "name": name,
        }
        if "search" in f"{name} {desc}".lower():
            return node, record
        if in_toolbar:
            candidates.append((node, record))
    # The search field is the only editable line edit in KeePassXC's toolbar.
    if len(candidates) == 1:
        return candidates[0]
    return None, None


def load_cache():
    try:
        return json.loads(CACHE_FILE.read_text())
    except (OSError, ValueError):
        return None


def save_cache(record):
    try:
        CACHE_FILE.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
        CACHE_FILE.write_text(json.dumps(record))
    except OSError:
        pass


def cached_search(app, record):
    """Fast path: follow the remembered child indices, then verify the target.
    Costs one D-Bus call per level plus a few for the check."""
    if not record:
        return None
    path, roles, name = record.get("path"), record.get("roles"), record.get("name")
    if not isinstance(path, list) or not isinstance(roles, list) \
            or len(roles) != len(path) + 1:
        return None
    # Without a name to compare against, also check every level's role.
    check_each_level = not name
    node = app
    for depth, i in enumerate(path):
        node = child_at(node, i)
        if node is None:
            return None
        if check_each_level and int(safe(node.get_role, -1)) != roles[depth + 1]:
            return None
    role = safe(node.get_role)
    if role is None or int(role) != roles[-1]:
        return None
    if not is_line_edit(role, safe(node.get_state_set)):
        return None
    if (safe(node.get_name, "") or "") != name:
        return None
    return node


def main():
    global DEBUG
    ap = argparse.ArgumentParser()
    ap.add_argument("query")
    ap.add_argument("--timeout", type=float, default=0,
                    help="seconds to keep waiting for the search box "
                         "(e.g. while you unlock the database)")
    ap.add_argument("--debug", action="store_true")
    args = ap.parse_args()
    DEBUG = args.debug

    record = load_cache()
    deadline = time.monotonic() + args.timeout
    while True:
        t0 = time.monotonic()
        app = find_keepassxc()
        box = None
        if app is not None:
            box = cached_search(app, record)
            how = "cache"
            if box is None:
                how = "walk"
                box, new_record = full_search(app)
                if box is not None:
                    save_cache(new_record)
                    record = new_record
        log(lambda: f"pass took {time.monotonic() - t0:.3f}s "
                    f"(app={app is not None}, box={box is not None}, via {how if app else '-'})")
        if box is not None:
            break
        if time.monotonic() >= deadline:
            if app is None:
                print("KeePassXC not found on the accessibility bus",
                      file=sys.stderr)
                return 1
            print("search box not found (database locked?); "
                  "rerun with --debug to see the widget tree", file=sys.stderr)
            return 2
        time.sleep(0.2)

    Atspi.EditableText.set_text_contents(box, args.query)
    safe(lambda: Atspi.Component.grab_focus(box))
    return 0


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