ollie/tools/gui_screenshot

#!/usr/bin/env python3
# ollie:prompt
# ## gui_screenshot
#
# Capture screenshots via KWin (Wayland-safe).
#
# **Args**: `[mode]` or `[mode, "--output=path"]`
# - mode: `active`, `screen`, `area`, `window` (default: active)
#
# ```
# call_tool: calls=[{tool: "gui_screenshot", args: ["active"]}]
# call_tool: calls=[{tool: "gui_screenshot", args: ["screen", "--output=/tmp/shot.png"]}]
# ```
# ollie:end
"""KDE screenshot tool using org.kde.KWin.ScreenShot2.

Modes:
  active   - capture the currently active window (default)
  screen   - capture entire screen (or specific with --screen=N)
  area     - capture a region (interactive picker)
  window   - capture window by title substring match

Output defaults to /tmp/ollie_screenshot_<timestamp>.png
"""
import sys
import os
import time
import argparse
import subprocess
import tempfile

SERVICE = "org.kde.KWin"
OBJ_PATH = "/org/kde/KWin/ScreenShot2"
IFACE = "org.kde.KWin.ScreenShot2"


def parse_args():
    p = argparse.ArgumentParser(description="KDE screenshot via KWin D-Bus")
    p.add_argument("mode", nargs="?", default="active",
                   choices=["active", "screen", "area", "window"],
                   help="Capture mode (default: active)")
    p.add_argument("--output", "-o", default="",
                   help="Output file path (default: /tmp/ollie_screenshot_<ts>.png)")
    p.add_argument("--screen", type=int, default=0,
                   help="Screen index for 'screen' mode")
    p.add_argument("--title", default="",
                   help="Window title substring for 'window' mode")
    p.add_argument("--include-cursor", action="store_true", default=False)
    p.add_argument("--include-decoration", action="store_true", default=True)
    return p.parse_args()


def get_output_path(args):
    if args.output:
        return args.output
    ts = int(time.time())
    return f"/tmp/ollie_screenshot_{ts}.png"


def capture_with_spectacle(mode, output, title=""):
    """Fallback: use spectacle CLI if ScreenShot2 is unavailable."""
    cmd = ["spectacle", "-b", "-n", "-o", output]
    if mode == "active":
        cmd.append("-a")
    elif mode == "screen":
        cmd.append("-f")
    elif mode == "area":
        cmd.append("-r")
    elif mode == "window":
        cmd.append("-a")  # fall back to active window

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
    if result.returncode != 0:
        return None
    return output


def find_window_id_by_title(title):
    """Use KWin scripting to find window UUID by title."""
    script = f"""
    const clients = workspace.windowList();
    for (let i = 0; i < clients.length; i++) {{
        if (clients[i].caption.toLowerCase().includes('{title.lower()}')) {{
            print(clients[i].internalId);
            break;
        }}
    }}
    """
    # Load script via D-Bus
    load_cmd = [
        "gdbus", "call", "--session",
        "--dest", "org.kde.KWin",
        "--object-path", "/Scripting",
        "--method", "org.kde.kwin.Scripting.loadScript", "", script
    ]
    # This is complex; for now fall back to xdotool/wmctrl
    try:
        result = subprocess.run(
            ["xdotool", "search", "--name", title],
            capture_output=True, text=True, timeout=5
        )
        if result.stdout.strip():
            return result.stdout.strip().split("\n")[0]
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    return None


def capture_via_dbus(mode, output, args):
    """Use KWin ScreenShot2 portal-style capture with pipe fd."""
    hints = {
        "include-cursor": args.include_cursor,
        "include-decoration": args.include_decoration,
    }

    # ScreenShot2 methods use file descriptors which gdbus can't handle directly.
    # Use the portal interface instead, or fall back to spectacle.
    # Try the portal first:
    portal_cmd = [
        "gdbus", "call", "--session",
        "--dest", "org.freedesktop.portal.Desktop",
        "--object-path", "/org/freedesktop/portal/desktop",
        "--method", "org.freedesktop.portal.Screenshot.Screenshot",
        "",  # parent_window
        "{'interactive': <false>}",
    ]

    # Spectacle is more reliable for non-interactive capture
    result = capture_with_spectacle(mode, output, args.title)
    if result and os.path.isfile(result):
        return result

    # Last resort: try import (ImageMagick) for X11
    if os.environ.get("DISPLAY"):
        if mode == "active":
            wid = subprocess.run(
                ["xdotool", "getactivewindow"],
                capture_output=True, text=True
            ).stdout.strip()
            if wid:
                r = subprocess.run(["import", "-window", wid, output],
                                   capture_output=True, timeout=10)
                if r.returncode == 0:
                    return output
        elif mode == "screen":
            r = subprocess.run(["import", "-window", "root", output],
                               capture_output=True, timeout=10)
            if r.returncode == 0:
                return output

    return None


def main():
    args = parse_args()
    output = get_output_path(args)

    result = capture_via_dbus(args.mode, output, args)

    if result and os.path.isfile(result):
        size = os.path.getsize(result)
        print(f"{result} ({size} bytes)")
    else:
        print("ERROR: screenshot capture failed", file=sys.stderr)
        print("Ensure spectacle or imagemagick is installed", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()