#!/usr/bin/env python3
"""qs-shot — screenshot capture with a QuickShell toast preview.

Subcommands:
  pick    show a wofi chooser (select area / full screen) then capture
  area    interactively select a region (grimblast)
  screen  capture the current output (grimblast)
  open    <path> open a saved screenshot in the default viewer (toast action)
  copy    <path> re-push a saved screenshot to the clipboard (toast action)

The capture is copied to the clipboard AND saved to a temp file. A
notification toast is posted with the image preview and Open / Copy
actions; the shell intercepts those in calNotifAction (identifier
payload carries the saved file path).

All binary paths are injected via $QS_* env vars by the .local/bin/qs-shot
wrapper so execution works regardless of PATH. When no DISPLAY/WAYLAND_DISPLAY
is present, wofi can't open, so the picker exits silently with code 1.
"""

import os
import subprocess
import sys
import time

TARGETS = {
    "area": "Select area",
    "screen": "Full screen",
}


def run(command, **kwargs):
    return subprocess.run(command, capture_output=True, text=True, **kwargs)


def wofi(entries, prompt):
    cmd = [os.environ.get("QS_WOFI", "wofi"), "--dmenu", "--prompt", prompt]
    result = run(cmd, input=entries)
    if result.returncode != 0:
        return None
    return result.stdout.rstrip("\n")


def _output_path():
    return os.path.join(
        os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")),
        "qs-shot",
        "shot-%s.png" % time.strftime("%Y%m%d-%H%M%S"),
    )


def _grab(target, out):
    """grimblast save <target> <out> — file only, no clipboard."""
    grimblast = os.environ.get("QS_GRIMBLAST", "grimblast")
    return subprocess.run(
        [grimblast, "save", target, out],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    ).returncode == 0


def _copy_to_clip(out):
    """wl-copy < out — fire-and-forget, mirrors the `copy` subcommand."""
    wl_copy = os.environ.get("QS_WLCOPY", "wl-copy")
    subprocess.Popen(["sh", "-c",
                      'exec "$1" --type image/png < "$2"', "--",
                      wl_copy, out],
                     stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def _notify(out, target):
    summary = "Screenshot" if target == "screen" else "Screenshot (area)"
    body = "\u2702 + \U0001f4be  " + out
    notify = os.environ.get("QS_NOTIFY", "notify-send")
    # fire-and-forget: notify-send with --action stays alive as the D-Bus action
    # sender waiting for ActionInvoked; quickshell handles the shot-* actions by
    # identifier in QML instead, so the sender can exit immediately.
    subprocess.Popen([
        notify,
        "-a", "Screenshot",
        "-t", "8000",
        "--hint=string:image-path:" + out,
        "--action=shot-open:%s=OPEN" % out,
        "--action=shot-copy:%s=COPY" % out,
        summary,
        body,
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def pick_mode():
    if "WAYLAND_DISPLAY" not in os.environ and "DISPLAY" not in os.environ:
        return 1
    out = _output_path()
    os.makedirs(os.path.dirname(out), exist_ok=True)
    _prune_old_shots(os.path.dirname(out), keep=40)
    # Capture BEFORE the chooser opens. wofi's layer surface lingers as a
    # composited zombie after it exits (Hyprland keeps it in hyprctl layers even
    # with a dead pid), so any post-exit wait races teardown and can bake the
    # menu into the shot. A pre-capture makes that impossible for full screen.
    if not _grab("screen", out):
        return 1
    label = wofi("\n".join(TARGETS.values()), "\U0001f4f7 Screenshot")
    if not label:
        try:
            os.unlink(out)
        except OSError:
            pass
        return 0
    target = None
    for t, entry in TARGETS.items():
        if label == entry:
            target = t
            break
    if target is None:
        try:
            os.unlink(out)
        except OSError:
            pass
        return 0
    if target == "area":
        if not _grab("area", out):
            return 1
    _copy_to_clip(out)
    _notify(out, target)
    return 0


def _prune_old_shots(directory, keep=40):
    try:
        files = sorted(
            os.path.join(directory, f) for f in os.listdir(directory)
            if f.startswith("shot-") and f.endswith(".png")
        )
        for old in files[:-keep]:
            os.unlink(old)
    except OSError:
        pass


def capture(target):
    out = _output_path()
    os.makedirs(os.path.dirname(out), exist_ok=True)
    _prune_old_shots(os.path.dirname(out), keep=40)
    if not _grab(target, out):
        return 1
    _copy_to_clip(out)
    _notify(out, target)
    return 0


def main():
    sub = sys.argv[1] if len(sys.argv) > 1 else "pick"
    if sub == "area" or sub == "screen":
        return capture(sub)
    if sub == "pick":
        return pick_mode()
    if sub == "open" and len(sys.argv) > 2:
        xdg_open = os.environ.get("QS_XDGO", "xdg-open")
        subprocess.Popen([xdg_open, sys.argv[2]])
        return 0
    if sub == "copy" and len(sys.argv) > 2:
        wl_copy = os.environ.get("QS_WLCOPY", "wl-copy")
        subprocess.Popen(["sh", "-c",
                          'exec "$1" --type image/png < "$2"', "--",
                          wl_copy, sys.argv[2]])
        return 0
    print(__doc__, file=sys.stderr)
    return 2


if __name__ == "__main__":
    raise SystemExit(main())