#!/usr/bin/env python3
"""Capture a screenshot for visual verification (Wayland / Hyprland).

Claude can read PNG files, so this turns "does the UI look right?" into an
inspectable artifact: capture the screen or one window, then read the path it
prints.

Modes (pick one):
  --full              capture everything (all outputs)            [default]
  --active            capture the currently focused window
  --window REGEX      capture the window whose class or title matches REGEX
                      (case-insensitive). Errors if zero or many match.
  --list              list open windows (class / title / workspace) and exit

Options:
  --output PATH       where to write the PNG
                      (default: /tmp/claude-screenshot-<timestamp>.png)
  --delay SECONDS     wait before capturing (let a window settle / animate)

On success the saved path is printed on its own line; nothing else goes to
stdout, so a caller can capture it directly.

Requires grim and hyprctl. A window must be visible on the active workspace
for its region to capture meaningfully — grim reads the compositor's rendered
output, so a window on another workspace yields whatever is drawn at that
region, not the hidden window.
"""

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time
from datetime import datetime


def die(msg):
    print(f"screenshot: {msg}", file=sys.stderr)
    sys.exit(1)


def require_tools(*tools):
    missing = [t for t in tools if shutil.which(t) is None]
    if missing:
        die(f"missing required tool(s): {', '.join(missing)} "
            f"(this script targets Wayland/Hyprland)")


def hypr(*args):
    """Run `hyprctl -j <args>` and return parsed JSON."""
    out = subprocess.run(["hyprctl", "-j", *args],
                         capture_output=True, text=True)
    if out.returncode != 0:
        die(f"hyprctl {' '.join(args)} failed: {out.stderr.strip()}")
    return json.loads(out.stdout)


def geometry_str(client):
    """Build a grim -g geometry string from a Hyprland client's at/size."""
    (x, y), (w, h) = client["at"], client["size"]
    return f"{x},{y} {w}x{h}"


def list_windows():
    clients = hypr("clients")
    rows = sorted(clients, key=lambda c: (c.get("workspace", {}).get("id", 0),
                                          c.get("class", "")))
    if not rows:
        print("(no open windows)")
        return
    for c in rows:
        ws = c.get("workspace", {}).get("name", "?")
        print(f"  ws:{ws:<8} class={c.get('class','')!r:30} title={c.get('title','')!r}")


def find_window(regex):
    pat = re.compile(regex, re.IGNORECASE)
    matches = [c for c in hypr("clients")
               if pat.search(c.get("class", "") or "")
               or pat.search(c.get("title", "") or "")]
    if not matches:
        die(f"no window matches {regex!r} — run --list to see open windows")
    if len(matches) > 1:
        labels = "\n".join(f"  class={c.get('class','')!r} title={c.get('title','')!r}"
                           for c in matches)
        die(f"{len(matches)} windows match {regex!r}; narrow it:\n{labels}")
    return matches[0]


def grim(output, geometry=None):
    cmd = ["grim"]
    if geometry:
        cmd += ["-g", geometry]
    cmd.append(output)
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode != 0:
        die(f"grim failed: {res.stderr.strip()}")


def main():
    p = argparse.ArgumentParser(add_help=True, description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    mode = p.add_mutually_exclusive_group()
    mode.add_argument("--full", action="store_true", help="capture all outputs (default)")
    mode.add_argument("--active", action="store_true", help="capture the focused window")
    mode.add_argument("--window", metavar="REGEX", help="capture window matching REGEX")
    mode.add_argument("--list", action="store_true", help="list windows and exit")
    p.add_argument("--output", "-o", help="PNG output path")
    p.add_argument("--delay", type=float, default=0.0, help="seconds to wait before capture")
    args = p.parse_args()

    require_tools("grim", "hyprctl")

    if args.list:
        list_windows()
        return

    if args.delay:
        time.sleep(args.delay)

    geometry = None
    if args.window:
        geometry = geometry_str(find_window(args.window))
    elif args.active:
        win = hypr("activewindow")
        if not win or "at" not in win:
            die("no active window")
        geometry = geometry_str(win)
    # else: --full / default → whole screen

    out = args.output or f"/tmp/claude-screenshot-{datetime.now():%Y%m%d-%H%M%S}.png"
    parent = os.path.dirname(os.path.abspath(out))
    os.makedirs(parent, exist_ok=True)
    grim(out, geometry)
    print(out)


if __name__ == "__main__":
    main()