#!/usr/bin/env python3
"""Window-targeted screen capture for app b-roll.

Adapted from the technique in github.com/jhead/macos-screen-mcp: enumerate windows
with Quartz CGWindowListCopyWindowInfo, then capture ONE window by its window id.
Only the window's own content is captured, so the menu bar, other windows, the
mouse cursor, notification banners and browser "is being debugged" bars never
appear in the frame.

Why not `ffmpeg -f avfoundation -i "Capture screen N"`: a full-display grab records
whatever the compositor last painted. When a browser tab is being driven over the
CDP/debug bridge macOS/Chrome throttle its repaints, so the capture holds stale
frames and all the motion arrives in a burst at the end (observed 2026-08-24: 8 of
13s frozen). It also drags in the menu bar, toolbar and cursor.

The fix here is stop-motion: drive the UI one step, capture one frame, repeat. Each
frame is pulled on demand, so a frame can never be stale relative to the step that
produced it. Assemble the numbered frames into a clip afterwards.

Pixels come from the built-in `screencapture -l <window_id>` (supported CLI, no
deprecated CGWindowListCreateImage path, no Pillow dependency).

Requires: Screen Recording permission for the invoking terminal app.

  capture_window.py list [--filter chrome]
  capture_window.py shot --window "Iron Trainer" --out work/broll/f_001.png
  capture_window.py burst --window 12345 --out work/broll/seq --frames 90 --interval 0.1
  capture_window.py assemble --frames work/broll/seq --out work/broll/b_app.mp4 --fps 30
"""
import argparse
import subprocess
import sys
import time
from pathlib import Path

def list_windows():
    """All on-screen windows that have a title: id, owner, name, bounds.

    Quartz is imported here rather than at module scope so the rest of this
    file — the pure window-picking and frame-dedupe logic — stays importable
    on a machine without pyobjc, CI included.
    """
    try:
        from Quartz import (
            CGWindowListCopyWindowInfo,
            kCGWindowListOptionOnScreenOnly,
            kCGNullWindowID,
        )
    except ImportError:
        sys.exit("pyobjc Quartz missing. Install it with:\n"
                 "  pip install pyobjc-framework-Quartz")
    info = CGWindowListCopyWindowInfo(kCGWindowListOptionOnScreenOnly, kCGNullWindowID)
    out = []
    for w in info or []:
        d = dict(w)
        name = d.get("kCGWindowName")
        if not name:
            continue
        b = d.get("kCGWindowBounds") or {}
        out.append({
            "id": d.get("kCGWindowNumber"),
            "name": name,
            "owner": d.get("kCGWindowOwnerName"),
            "w": int(b.get("Width", 0)),
            "h": int(b.get("Height", 0)),
        })
    return out


def resolve(identifier, windows=None):
    """Window id from an int-like id, or the largest window matching a title/owner.

    `windows` defaults to the live window list; pass one to test the picking
    rule without a display attached.
    """
    s = str(identifier)
    if s.isdigit():
        return int(s)
    needle = s.lower()
    hits = [w for w in (list_windows() if windows is None else windows)
            if needle in (w["name"] or "").lower() or needle in (w["owner"] or "").lower()]
    if not hits:
        sys.exit(f"No window matching {identifier!r}. Run: capture_window.py list")
    # biggest match = the real content window, not a tooltip/panel
    hits.sort(key=lambda w: w["w"] * w["h"], reverse=True)
    return hits[0]["id"]


def shot(window_id, out_path):
    """Capture one window to PNG. -o drops the drop-shadow, -x mutes the shutter."""
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    r = subprocess.run(
        ["screencapture", "-l", str(window_id), "-o", "-x", str(out_path)],
        capture_output=True, text=True,
    )
    if r.returncode != 0 or not out_path.exists():
        sys.exit(f"capture failed for window {window_id}: {r.stderr.strip()}")
    return out_path


def distinct_frames(paths):
    """Drop each frame identical to the one before it, keeping order.

    screencapture manages roughly 3-5 fps, so a stepped animation is
    over-sampled on purpose and most frames repeat the previous state. Keeping
    one frame per rendered state is what turns that into motion; it also trims
    the idle head and tail spent waiting for the UI to move.

    Compares consecutive frames only. Two identical states separated by a
    different one are both kept — a UI that returns to a previous state is
    still moving.
    """
    import hashlib
    keep, last = [], None
    for f in paths:
        h = hashlib.md5(Path(f).read_bytes()).hexdigest()
        if h != last:
            keep.append(f)
        last = h
    return keep


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("list", help="list on-screen windows")
    p.add_argument("--filter", help="substring match on title or owner")

    p = sub.add_parser("shot", help="capture one window to PNG")
    p.add_argument("--window", required=True, help="window id, or title/owner substring")
    p.add_argument("--out", required=True)

    p = sub.add_parser("burst", help="capture N frames on a timer (unattended motion)")
    p.add_argument("--window", required=True)
    p.add_argument("--out", required=True, help="directory for frame_%%05d.png")
    p.add_argument("--frames", type=int, default=90)
    p.add_argument("--interval", type=float, default=0.1, help="seconds between frames")

    p = sub.add_parser("assemble", help="numbered frames -> mp4")
    p.add_argument("--frames", required=True, help="directory of frame_*.png")
    p.add_argument("--out", required=True)
    p.add_argument("--fps", type=int, default=30)
    p.add_argument("--crf", type=int, default=18)
    p.add_argument("--dedupe", action="store_true",
                   help="drop consecutive identical frames first. screencapture only "
                        "manages ~3-5 fps, so the intended pattern is to over-sample a "
                        "stepped animation and keep one frame per rendered state; this "
                        "also trims the idle head/tail while waiting for the UI to move.")
    p.add_argument("--crop", help="ffmpeg crop, e.g. 2992:1660:0:214 to cut browser chrome")

    a = ap.parse_args()

    if a.cmd == "list":
        for w in list_windows():
            line = f"{w['id']:>7}  {w['w']:>5}x{w['h']:<5}  {w['owner']}  |  {w['name']}"
            if not a.filter or a.filter.lower() in line.lower():
                print(line)

    elif a.cmd == "shot":
        print(shot(resolve(a.window), a.out))

    elif a.cmd == "burst":
        wid = resolve(a.window)
        d = Path(a.out)
        d.mkdir(parents=True, exist_ok=True)
        for i in range(a.frames):
            shot(wid, d / f"frame_{i:05d}.png")
            time.sleep(a.interval)
        print(f"{a.frames} frames -> {d}")

    elif a.cmd == "assemble":
        frames = sorted(Path(a.frames).glob("frame_*.png"))
        if not frames:
            sys.exit(f"no frame_*.png in {a.frames}")
        src_dir, n_in = Path(a.frames), len(frames)

        if a.dedupe:
            keep = distinct_frames(frames)
            src_dir = Path(a.frames) / "_dedup"
            if src_dir.exists():
                for old in src_dir.glob("frame_*.png"):
                    old.unlink()
            src_dir.mkdir(parents=True, exist_ok=True)
            for i, f in enumerate(keep):
                (src_dir / f"frame_{i:05d}.png").write_bytes(f.read_bytes())
            print(f"dedupe: {n_in} -> {len(keep)} distinct states")
            frames = keep

        vf = "scale=trunc(iw/2)*2:trunc(ih/2)*2"
        if a.crop:
            vf = f"crop={a.crop},{vf}"
        cmd = ["ffmpeg", "-y", "-v", "error", "-framerate", str(a.fps),
               "-i", str(src_dir / "frame_%05d.png"),
               "-c:v", "libx264", "-crf", str(a.crf), "-pix_fmt", "yuv420p",
               "-vf", vf, a.out]
        subprocess.run(cmd, check=True)
        print(f"{len(frames)} frames -> {a.out}")


if __name__ == "__main__":
    main()