#!/usr/bin/env -S uv run --quiet --with pillow python3
"""vmshot -- one-shot screenshot contact sheet for the SOTS lab Windows guests.

Why this exists
---------------
The lab went from one Windows guest (VM 140) to five (140, 141, 144, 145, 146).
There was no way to see what they were all doing without issuing a QEMU
``screendump`` per guest by hand.  This grabs every guest in parallel and lays
them out on a single labelled contact sheet.

Mechanism
---------
``qm monitor <id> <<< "screendump <file> -f png"`` on the Proxmox host.  This
reads the guest framebuffer straight out of QEMU:

* no guest agent, no guest network, nothing installed inside the guest;
* it does not perturb the guest at all -- relevant to method-rule 19, an
  instrument that perturbs the thing it measures.  Reading VM 140's screen is
  therefore *not* an experiment and does not take VM 140's exclusivity lock;
* the campaign board already records that ``qm monitor`` screendump is more
  reliable than the in-guest click helper's ``shot``.

spicy runs pve-manager 9.2 / QEMU 11, whose ``screendump`` takes ``-f png``
natively, so no PPM conversion and no netpbm/ImageMagick is needed on the host
(none is installed there, and this tool installs nothing).  The PNGs are written
to a per-run temp dir on the host, streamed back inside a base64 tar in the same
SSH round trip, and the temp dir is removed before the SSH call returns -- so
nothing is left behind on spicy even if this script is killed.

Relationship to vmwatch
-----------------------
``tools/vmwatch.py`` is the live version of this: an always-on service on spicy
serving an auto-refreshing wall at http://192.168.3.201:8140/.  When that
service is reachable this script pulls its already-captured frames over HTTP
instead of opening its own SSH session -- same frames, no extra load on the
guests, and it works from anywhere on the LAN without SSH.  If the service is
down it falls back to the SSH + ``qm monitor`` path described below, so the CLI
never depends on the service being up.

Usage
-----
    tools/vmshot.py                     # contact sheet of every sots-* guest
    tools/vmshot.py --open              # ... and open it in the default viewer
    tools/vmshot.py 140 141             # only these ids
    tools/vmshot.py --one 140           # full-size single guest, no sheet
    tools/vmshot.py --cols 3 --width 900

**Run it as ``tools/vmshot.py``, not ``python3 tools/vmshot.py``.** The shebang is
``uv run --with pillow``, so invoking the file directly supplies Pillow even on a host
that has none; putting an interpreter in front of it throws that away. Lane BR hit
exactly this and concluded the tool was unusable from the WSL host. ``--one`` now works
either way -- it writes the guest's framebuffer bytes verbatim and needs no Pillow --
and only the contact sheet requires it, with the install line in the error.

For a single frame during a measurement run, going straight to the host is better still
and is what lane BR ended up doing::

    ssh spicy "echo 'screendump /tmp/x.png -f png' | qm monitor 146"

That is the *live* framebuffer. The wall at http://192.168.3.201:8140/ serves a cached
frame up to a poll cycle (~5 s) old, which reads exactly like a swallowed click and has
cost two lanes wasted clicks.

Output lands in ``~/sots-re/dumps/vmshot/`` (``dumps/`` is gitignored).

A guest that is stopped, paused, or whose screendump fails gets a labelled
placeholder tile with the host-side error rather than aborting the sheet.
"""

from __future__ import annotations

import argparse
import base64
import io
import json
import os
import re
import subprocess
import urllib.error
import urllib.request
import sys
import tarfile
import time
from datetime import datetime
from pathlib import Path

# Pillow is needed only to composite the contact sheet. `--one` writes the guest's
# framebuffer bytes straight through, and that is the mode a measurement lane actually
# wants -- lane BR could not use this tool at all from the WSL host because the import
# was unconditional and Pillow is not installed there. Fail at the point of use, with
# the fix in the message, rather than at import.
try:
    from PIL import Image, ImageDraw, ImageFont

    _PIL_ERROR = None
except ImportError as e:  # pragma: no cover - depends on the host
    Image = ImageDraw = ImageFont = None  # type: ignore[assignment]
    _PIL_ERROR = e


def _require_pil() -> None:
    if _PIL_ERROR is not None:
        sys.exit(
            f"vmshot: the contact sheet needs Pillow, which is not installed here ({_PIL_ERROR}).\n"
            "        `--one <id>` works without it and writes the guest's framebuffer verbatim.\n"
            "        To build sheets on this host: uv pip install pillow"
        )

HOST = "spicy"
SERVICE = os.environ.get("VMWATCH_URL", "http://192.168.3.201:8140")
OUT_DIR = Path.home() / "sots-re" / "dumps" / "vmshot"
FLEET_PATTERN = re.compile(r"sots", re.I)
FONT_PATH = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
FONT_BOLD = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"

# Tile chrome
LABEL_H = 34
PAD = 10
BG = (24, 26, 30)
LABEL_BG = (38, 42, 50)
LABEL_BG_BAD = (74, 34, 34)
FG = (232, 234, 238)
FG_DIM = (150, 156, 166)
BORDER = (60, 66, 78)
PLACEHOLDER_BG = (44, 30, 30)


def ssh(cmd: str, timeout: int = 90) -> subprocess.CompletedProcess:
    return subprocess.run(
        ["ssh", "-o", "BatchMode=yes", HOST, cmd],
        capture_output=True,
        timeout=timeout,
    )


def service_state(timeout: float = 3.0) -> list[dict] | None:
    """Fleet state from a reachable vmwatch service, or None if it is not up."""
    try:
        with urllib.request.urlopen(f"{SERVICE}/api/state", timeout=timeout) as r:
            return json.load(r)
    except (urllib.error.URLError, OSError, ValueError, json.JSONDecodeError):
        return None


def service_grab(ids: list[int], state: list[dict]) -> dict[int, dict]:
    """Pull already-captured frames from the vmwatch service. No guest load."""
    by_id = {g["id"]: g for g in state}
    out: dict[int, dict] = {}
    for vid in ids:
        g = by_id.get(vid)
        if g is None:
            out[vid] = {"png": None, "status": "absent", "log": "not known to vmwatch"}
            continue
        png = None
        if g.get("ok"):
            try:
                with urllib.request.urlopen(
                    f"{SERVICE}/shot/{vid}.png?t={g.get('ts', '')}", timeout=10
                ) as r:
                    png = r.read()
            except (urllib.error.URLError, OSError):
                png = None
        out[vid] = {
            "png": png,
            "status": g.get("status", "?"),
            "log": g.get("error", "") if not png else "",
        }
    return out


def discover_fleet() -> list[tuple[int, str, str]]:
    """Return [(vmid, name, status)] for every VM whose name matches sots-*."""
    r = ssh("qm list", timeout=30)
    if r.returncode != 0:
        sys.exit(f"vmshot: `qm list` on {HOST} failed: {r.stderr.decode().strip()}")
    fleet = []
    for line in r.stdout.decode().splitlines()[1:]:
        parts = line.split()
        if len(parts) < 3 or not parts[0].isdigit():
            continue
        vmid, name, status = int(parts[0]), parts[1], parts[2]
        if FLEET_PATTERN.search(name):
            fleet.append((vmid, name, status))
    return sorted(fleet)


# Remote script: dump every requested guest in parallel, tar the results back,
# always clean up.  Per-guest stderr is kept so a failure explains itself.
REMOTE = r"""
set -u
D=$(mktemp -d /tmp/vmshot.XXXXXX)
trap 'rm -rf "$D"' EXIT INT TERM
for id in %(ids)s; do
  (
    st=$(qm status "$id" 2>&1 | awk '{print $2}')
    echo "$st" > "$D/$id.status"
    timeout %(t)s qm monitor "$id" <<< "screendump $D/$id.png -f png" \
      > "$D/$id.log" 2>&1
    # qm monitor exits 0 even when the monitor command errors; the monitor
    # echoes the failure into the log, so treat "no file" as the real test.
    [ -s "$D/$id.png" ] || echo "no framebuffer written" >> "$D/$id.log"
  ) &
done
wait
tar -C "$D" -cf - . | base64 -w0
"""


def grab(ids: list[int], per_vm_timeout: int = 20) -> dict[int, dict]:
    """Screendump each id on the host; return {id: {'png': bytes|None, 'status', 'log'}}."""
    remote = REMOTE % {"ids": " ".join(str(i) for i in ids), "t": per_vm_timeout}
    r = ssh(f"bash -s <<'VMSHOT_EOF'\n{remote}\nVMSHOT_EOF", timeout=per_vm_timeout + 60)
    if r.returncode != 0 and not r.stdout.strip():
        sys.exit(f"vmshot: remote grab failed: {r.stderr.decode().strip()[:500]}")

    out: dict[int, dict] = {i: {"png": None, "status": "?", "log": ""} for i in ids}
    try:
        blob = base64.b64decode(r.stdout.strip())
        with tarfile.open(fileobj=io.BytesIO(blob), mode="r:") as tf:
            for m in tf.getmembers():
                name = Path(m.name).name
                if not m.isfile():
                    continue
                stem, _, ext = name.rpartition(".")
                if not stem.isdigit():
                    continue
                vid = int(stem)
                if vid not in out:
                    continue
                data = tf.extractfile(m).read()
                if ext == "png":
                    out[vid]["png"] = data
                elif ext == "status":
                    out[vid]["status"] = data.decode(errors="replace").strip() or "?"
                elif ext == "log":
                    out[vid]["log"] = data.decode(errors="replace").strip()
    except Exception as e:  # noqa: BLE001 - a mangled tar must not lose the whole sheet
        for v in out.values():
            v["log"] = v["log"] or f"could not unpack remote payload: {e}"
    return out


def _font(size: int, bold: bool = False):
    try:
        return ImageFont.truetype(FONT_BOLD if bold else FONT_PATH, size)
    except OSError:
        return ImageFont.load_default()


def _clean_monitor_log(log: str) -> str:
    """Strip the monitor banner/prompt noise, keep the actual complaint."""
    keep = []
    for ln in log.splitlines():
        ln = ln.replace("\x1b", "").strip()
        ln = re.sub(r"^(QEMU \d[\w.]* monitor.*|qm> ?)", "", ln).strip()
        if not ln or ln.startswith("Entering QEMU Monitor") or ln.startswith("Type 'help'"):
            continue
        keep.append(ln)
    return "; ".join(keep)[:180]


def make_tile(vmid: int, name: str, info: dict, width: int, stamp: str):
    png = info.get("png")
    status = info.get("status", "?")
    err = _clean_monitor_log(info.get("log", ""))

    if png:
        shot = Image.open(io.BytesIO(png)).convert("RGB")
        native = f"{shot.width}x{shot.height}"
        h = max(1, round(shot.height * width / shot.width))
        shot = shot.resize((width, h), Image.LANCZOS)
        detail = f"{status}  {native}  {stamp}"
        bad = False
    else:
        h = round(width * 3 / 4)
        shot = Image.new("RGB", (width, h), PLACEHOLDER_BG)
        d = ImageDraw.Draw(shot)
        msg = err or f"no screendump ({status})"
        d.text(
            (width // 2, h // 2 - 12),
            "NO SIGNAL",
            font=_font(max(18, width // 22), bold=True),
            fill=(210, 120, 120),
            anchor="mm",
        )
        f = _font(max(10, width // 60))
        # crude wrap
        line, lines = "", []
        for word in msg.split():
            if len(line) + len(word) + 1 > 52:
                lines.append(line)
                line = word
            else:
                line = f"{line} {word}".strip()
        lines.append(line)
        for i, ln in enumerate(lines[:4]):
            d.text(
                (width // 2, h // 2 + 18 + i * 15),
                ln,
                font=f,
                fill=(190, 150, 150),
                anchor="mm",
            )
        detail = f"{status}  {stamp}"
        bad = True

    tile = Image.new("RGB", (width, h + LABEL_H), LABEL_BG_BAD if bad else LABEL_BG)
    tile.paste(shot, (0, LABEL_H))
    d = ImageDraw.Draw(tile)
    f_id, f_name, f_detail = _font(16, bold=True), _font(14), _font(11)
    y = LABEL_H // 2

    d.text((8, y), str(vmid), font=f_id, fill=FG, anchor="lm")
    name_x = 8 + round(d.textlength(str(vmid), font=f_id)) + 10

    # A narrow tile must not let the two labels overlap: drop the timestamp
    # first, then ellipsise the guest name to whatever room is left.
    for candidate in (detail, detail.rsplit("  ", 1)[0], status):
        detail_w = d.textlength(candidate, font=f_detail)
        if name_x + d.textlength(name, font=f_name) + 12 + detail_w + 8 <= width:
            detail = candidate
            break
    else:
        detail = status
        detail_w = d.textlength(detail, font=f_detail)

    room = width - 8 - detail_w - 12 - name_x
    shown = name
    while shown and d.textlength(shown + "…", font=f_name) > room:
        shown = shown[:-1]
    if shown != name:
        shown = (shown + "…") if shown else ""
    if room > 0 and shown:
        d.text((name_x, y), shown, font=f_name, fill=FG, anchor="lm")
    d.text((width - 8, y), detail, font=f_detail, fill=FG_DIM, anchor="rm")
    d.rectangle([0, 0, width - 1, h + LABEL_H - 1], outline=BORDER)
    return tile


def contact_sheet(fleet, shots, tile_w: int, cols: int, stamp: str):
    _require_pil()
    tiles = [make_tile(vid, name, shots.get(vid, {}), tile_w, stamp) for vid, name, _ in fleet]
    rows = (len(tiles) + cols - 1) // cols
    row_h = [
        max((t.height for t in tiles[r * cols : (r + 1) * cols]), default=0) for r in range(rows)
    ]
    header = 40
    W = PAD + cols * (tile_w + PAD)
    H = header + PAD + sum(h + PAD for h in row_h)
    sheet = Image.new("RGB", (W, H), BG)
    d = ImageDraw.Draw(sheet)
    live = sum(1 for v in shots.values() if v.get("png"))
    d.text((PAD, header // 2), "SOTS lab guests", font=_font(18, bold=True), fill=FG, anchor="lm")
    d.text(
        (W - PAD, header // 2),
        f"{live}/{len(fleet)} framebuffers captured   {stamp}   host={HOST}",
        font=_font(12),
        fill=FG_DIM,
        anchor="rm",
    )
    y = header + PAD
    for r in range(rows):
        x = PAD
        for t in tiles[r * cols : (r + 1) * cols]:
            sheet.paste(t, (x, y))
            x += tile_w + PAD
        y += row_h[r] + PAD
    return sheet


def main() -> int:
    ap = argparse.ArgumentParser(description="Contact sheet of the SOTS lab VM screens.")
    ap.add_argument("ids", nargs="*", type=int, help="VM ids (default: every sots-* guest)")
    ap.add_argument("--one", type=int, metavar="ID", help="save one guest full-size, no sheet")
    ap.add_argument("--cols", type=int, default=3)
    ap.add_argument("--width", type=int, default=760, help="tile width in px")
    ap.add_argument("--out", type=Path, help="output png path")
    ap.add_argument("--timeout", type=int, default=20, help="per-guest screendump timeout (s)")
    ap.add_argument("--keep-raw", action="store_true", help="also save each guest's raw png")
    ap.add_argument("--json", action="store_true", help="print a machine-readable result line")
    ap.add_argument("--open", action="store_true", help="xdg-open the result")
    ap.add_argument(
        "--ssh",
        action="store_true",
        help="always capture over SSH, even if the vmwatch service is reachable",
    )
    args = ap.parse_args()

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    t0 = time.time()

    state = None if args.ssh else service_state()
    if state is not None:
        names = {g["id"]: g["name"] for g in state}
        all_ids = [g["id"] for g in state]
        source = SERVICE
    else:
        discovered = discover_fleet()
        names = {v: n for v, n, _ in discovered}
        all_ids = [v for v, _, _ in discovered]
        source = f"ssh {HOST}"

    if args.one is not None:
        wanted = [args.one]
    elif args.ids:
        wanted = args.ids
    else:
        wanted = all_ids
    if not wanted:
        sys.exit(f"vmshot: no sots-* guests found via {source}")
    fleet = [(v, names.get(v, f"vm{v}"), "") for v in wanted]

    shots = service_grab(wanted, state) if state is not None else grab(wanted, args.timeout)
    now = datetime.now()
    stamp = now.strftime("%Y-%m-%d %H:%M:%S")
    tag = now.strftime("%Y%m%d-%H%M%S")

    if args.keep_raw or args.one is not None:
        for vid, info in shots.items():
            if info.get("png"):
                p = OUT_DIR / f"vm{vid}-{tag}.png"
                p.write_bytes(info["png"])

    if args.one is not None:
        info = shots.get(args.one, {})
        if not info.get("png"):
            print(
                f"vmshot: no framebuffer from {args.one}: "
                f"{_clean_monitor_log(info.get('log', '')) or info.get('status', '?')}",
                file=sys.stderr,
            )
            return 1
        out = args.out or (OUT_DIR / f"vm{args.one}-{tag}.png")
        out.write_bytes(info["png"])
    else:
        sheet = contact_sheet(fleet, shots, args.width, args.cols, stamp)
        out = args.out or (OUT_DIR / f"sheet-{tag}.png")
        sheet.save(out)
        if args.out is None:  # only the default location keeps a `latest` alias
            (OUT_DIR / "latest.png").write_bytes(out.read_bytes())

    live = sum(1 for v in shots.values() if v.get("png"))
    if args.json:
        print(
            json.dumps(
                {
                    "out": str(out),
                    "captured": live,
                    "requested": len(wanted),
                    "source": source,
                    "seconds": round(time.time() - t0, 1),
                    "guests": {
                        str(v): {
                            "name": names.get(v, ""),
                            "status": i.get("status"),
                            "ok": bool(i.get("png")),
                            "error": _clean_monitor_log(i.get("log", "")) or None,
                        }
                        for v, i in sorted(shots.items())
                    },
                },
                indent=2,
            )
        )
    else:
        for v, i in sorted(shots.items()):
            mark = "ok " if i.get("png") else "FAIL"
            note = "" if i.get("png") else "  " + (_clean_monitor_log(i.get("log", "")) or "?")
            print(f"  {mark} {v:<5} {names.get(v, ''):<20} {i.get('status', '?')}{note}")
        print(
            f"vmshot: {live}/{len(wanted)} captured in {time.time() - t0:.1f}s "
            f"via {source} -> {out}"
        )

    if args.open:
        subprocess.run(["xdg-open", str(out)], check=False)
    return 0 if live else 2


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