#!/usr/bin/env python3
"""vmwatch -- a fishtank for the SOTS lab Windows guests.

A small always-on HTTP service that serves an auto-refreshing page of live
guest screens.  Leave the tab open and watch the lab.

    http://192.168.3.201:8140/

Where it runs
-------------
On the Proxmox host **spicy** itself, because it needs `qm`.  Installed as
``/opt/vmwatch/vmwatch.py`` with a systemd unit ``vmwatch.service``.  The
canonical copy lives here in ``~/sots-re/tools/`` -- edit here, then reinstall
with ``tools/vmwatch-install.sh``.

Dependencies: **Python 3 standard library only.**  Nothing is installed on
spicy beyond this file and its unit -- no pip, no ImageMagick, no netpbm (none
of which are present there).

Mechanism
---------
A QMP ``screendump`` sent to each guest's ``/var/run/qemu-server/<id>.qmp``
unix socket -- the same socket, and the same command, that
``qm monitor <id> <<< "screendump <file> -f png"`` uses, minus the Perl.  That
detail is not cosmetic: forking ``qm`` once per guest per poll interval cost
about 90%% of a core on the host, which is not acceptable on a box that is also
running the guests.  ``qm monitor`` remains the fallback when the socket is
absent.  Either way this reads the guest framebuffer straight out of QEMU:

* no guest agent, no guest network, nothing running inside the guest;
* it does **not** perturb the guest -- relevant to method-rule 19, an
  instrument that perturbs the thing it measures has already burned this
  campaign once.  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`` accepts ``-f png``
natively.  On an older QEMU without ``-f``, this falls back to a PPM dump and a
hand-rolled PPM->PNG encoder (zlib + a stored-filter IDAT), so no external
image tooling is needed either way.  Temp files are read into memory and
unlinked immediately; nothing accumulates in /tmp.

Politeness
----------
A single background poller grabs every guest on a fixed interval and serves
every browser from that one in-memory copy, so N open tabs still cost one
screendump per guest per interval.  The poller goes idle when no client has
asked for anything in IDLE_AFTER seconds, and wakes on the next request.

Routes
------
    GET /                 the wall
    GET /one/<id>         one guest, full size
    GET /api/state        JSON: per-guest id, name, status, ok, error, age
    GET /shot/<id>.png    the latest cached framebuffer for that guest
    GET /healthz          plaintext ok

Read-only: this service never starts, stops, resets or reconfigures anything.
"""

from __future__ import annotations

import argparse
import html
import json
import os
import re
import socket
import struct
import subprocess
import sys
import tempfile
import threading
import time
import zlib
from concurrent.futures import ThreadPoolExecutor
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

FLEET_PATTERN = re.compile(r"sots-re", re.I)
DEFAULT_PORT = 8140
DEFAULT_INTERVAL = 5.0  # seconds between framebuffer grabs
FLEET_TTL = 20.0  # seconds between `qm list` refreshes
IDLE_AFTER = 120.0  # stop polling if nobody has looked for this long
DUMP_TIMEOUT = 15  # per-guest screendump timeout, seconds


# --------------------------------------------------------------------------
# host-side capture
# --------------------------------------------------------------------------


CONF_DIR = "/etc/pve/qemu-server"
QMP_DIR = "/var/run/qemu-server"


def _run(cmd: list[str], timeout: int, stdin: str | None = None) -> subprocess.CompletedProcess:
    return subprocess.run(
        cmd, input=stdin, capture_output=True, text=True, timeout=timeout, check=False
    )


class QmpError(Exception):
    pass


class Qmp:
    """Minimal QMP client over the per-VM unix socket PVE already exposes.

    This is the same socket `qm monitor` uses, opened the same way (connect,
    negotiate, one command, close).  Going straight to it instead of shelling
    out to `qm` matters: `qm` is a Perl program and forking one per guest per
    poll interval cost ~90%% of a core on the host.  Falls back to `qm monitor`
    if the socket is missing.
    """

    def __init__(self, vmid: int, timeout: float = 10.0):
        self.path = os.path.join(QMP_DIR, f"{vmid}.qmp")
        self.timeout = timeout
        self.sock: socket.socket | None = None
        self.buf = b""

    def __enter__(self) -> Qmp:
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        s.settimeout(self.timeout)
        s.connect(self.path)
        self.sock = s
        self._read_json()  # greeting
        self.execute("qmp_capabilities")
        return self

    def __exit__(self, *exc) -> None:
        if self.sock:
            try:
                self.sock.close()
            except OSError:
                pass

    def _read_json(self) -> dict:
        while True:
            nl = self.buf.find(b"\n")
            if nl >= 0:
                line, self.buf = self.buf[:nl], self.buf[nl + 1 :]
                if line.strip():
                    msg = json.loads(line)
                    if "event" in msg:  # async event: not our reply
                        continue
                    return msg
                continue
            chunk = self.sock.recv(65536)  # type: ignore[union-attr]
            if not chunk:
                raise QmpError("monitor closed the connection")
            self.buf += chunk

    def execute(self, cmd: str, **args) -> dict:
        payload = {"execute": cmd}
        if args:
            payload["arguments"] = args
        self.sock.sendall(json.dumps(payload).encode() + b"\n")  # type: ignore[union-attr]
        reply = self._read_json()
        if "error" in reply:
            raise QmpError(reply["error"].get("desc", json.dumps(reply["error"]))[:200])
        return reply.get("return", {})


def _conf_name(vmid: int) -> str | None:
    """Guest name from its PVE config, ignoring the [snapshot] sections."""
    try:
        with open(os.path.join(CONF_DIR, f"{vmid}.conf"), encoding="utf-8") as fh:
            for line in fh:
                if line.startswith("["):  # snapshot section begins
                    break
                if line.startswith("name:"):
                    return line.split(":", 1)[1].strip()
    except OSError:
        return None
    return None


def list_fleet() -> list[dict]:
    """[{id, name, status}] for every VM whose name matches sots-re.

    Read from /etc/pve/qemu-server + the QMP sockets rather than `qm list`,
    which is another Perl fork.  Falls back to `qm list` if that is not
    readable (e.g. running somewhere that is not a PVE node).
    """
    fleet = []
    try:
        entries = os.listdir(CONF_DIR)
    except OSError:
        entries = []
    for entry in entries:
        stem, _, ext = entry.rpartition(".")
        if ext != "conf" or not stem.isdigit():
            continue
        vmid = int(stem)
        name = _conf_name(vmid)
        if not name or not FLEET_PATTERN.search(name):
            continue
        status = "stopped"
        if os.path.exists(os.path.join(QMP_DIR, f"{vmid}.qmp")):
            try:
                with Qmp(vmid, timeout=5) as q:
                    status = q.execute("query-status").get("status", "running")
            except (OSError, QmpError, ValueError):
                status = "unreachable"
        fleet.append({"id": vmid, "name": name, "status": status})
    if fleet:
        return sorted(fleet, key=lambda g: g["id"])

    try:
        r = _run(["qm", "list"], timeout=20)
    except (subprocess.TimeoutExpired, OSError):
        return []
    for line in r.stdout.splitlines()[1:]:
        parts = line.split()
        if len(parts) < 3 or not parts[0].isdigit():
            continue
        if FLEET_PATTERN.search(parts[1]):
            fleet.append({"id": int(parts[0]), "name": parts[1], "status": parts[2]})
    return sorted(fleet, key=lambda g: g["id"])


def ppm_to_png(data: bytes) -> bytes:
    """Minimal binary-P6 PPM -> PNG. Only used if QEMU lacks screendump -f png."""
    if not data.startswith(b"P6"):
        raise ValueError("not a binary P6 PPM")
    # header: P6, then width height maxval, whitespace separated, # comments allowed
    fields, pos = [], 2
    while len(fields) < 3:
        while pos < len(data) and data[pos : pos + 1].isspace():
            pos += 1
        if data[pos : pos + 1] == b"#":
            while pos < len(data) and data[pos] != 0x0A:
                pos += 1
            continue
        start = pos
        while pos < len(data) and not data[pos : pos + 1].isspace():
            pos += 1
        fields.append(int(data[start:pos]))
    pos += 1  # single whitespace byte after maxval
    w, h, maxval = fields
    if maxval != 255:
        raise ValueError(f"unsupported PPM maxval {maxval}")
    px = data[pos : pos + w * h * 3]
    stride = w * 3
    raw = bytearray()
    for y in range(h):
        raw.append(0)  # filter type 0 (None)
        raw += px[y * stride : (y + 1) * stride]

    def chunk(tag: bytes, payload: bytes) -> bytes:
        return (
            struct.pack(">I", len(payload))
            + tag
            + payload
            + struct.pack(">I", zlib.crc32(tag + payload) & 0xFFFFFFFF)
        )

    return (
        b"\x89PNG\r\n\x1a\n"
        + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
        + chunk(b"IDAT", zlib.compress(bytes(raw), 6))
        + chunk(b"IEND", b"")
    )


def screendump(vmid: int, png_native: bool = True) -> tuple[bytes | None, str]:
    """Grab one guest framebuffer. Returns (png_bytes, error_message); never raises."""
    suffix = ".png" if png_native else ".ppm"
    fd, path = tempfile.mkstemp(prefix=f"vmwatch-{vmid}-", suffix=suffix, dir="/tmp")
    os.close(fd)
    os.unlink(path)  # QEMU creates it itself; we only wanted a unique name
    err = ""
    try:
        if os.path.exists(os.path.join(QMP_DIR, f"{vmid}.qmp")):
            args = {"filename": path}
            if png_native:
                args["format"] = "png"
            with Qmp(vmid, timeout=DUMP_TIMEOUT) as q:
                q.execute("screendump", **args)
        else:
            # not running, or not a PVE node: fall back to the CLI monitor
            cmd = f"screendump {path}" + (" -f png" if png_native else "")
            r = _run(["qm", "monitor", str(vmid)], timeout=DUMP_TIMEOUT, stdin=cmd + "\n")
            # `qm monitor` exits 0 even when the monitor command itself errors;
            # it echoes the failure to stdout, so "no file" is the real test.
            err = _clean_monitor_text(r.stdout + r.stderr)
    except FileNotFoundError:
        _unlink(path)
        return None, "guest is not running (no monitor socket)"
    except (socket.timeout, subprocess.TimeoutExpired):
        _unlink(path)
        return None, "screendump timed out"
    except (QmpError, OSError, ValueError) as e:
        _unlink(path)
        return None, str(e) or f"{type(e).__name__}"

    try:
        with open(path, "rb") as fh:
            blob = fh.read()
    except OSError:
        blob = b""
    finally:
        _unlink(path)

    if not blob:
        return None, err or "no framebuffer written"
    if png_native and blob.startswith(b"\x89PNG"):
        return blob, ""
    try:
        return ppm_to_png(blob), ""
    except Exception as e:  # noqa: BLE001
        return None, f"could not decode framebuffer: {e}"


def _unlink(path: str) -> None:
    try:
        os.unlink(path)
    except OSError:
        pass


def _clean_monitor_text(text: str) -> str:
    keep = []
    for ln in text.replace("\x1b", "").splitlines():
        ln = re.sub(r"^(QEMU \d[\w.]* monitor.*|qm> ?)", "", ln.strip()).strip()
        if not ln or ln.startswith(("Entering QEMU Monitor", "Type 'help'", "screendump ")):
            continue
        keep.append(ln)
    return "; ".join(keep)[:200]


# --------------------------------------------------------------------------
# poller
# --------------------------------------------------------------------------


class Tank:
    def __init__(self, interval: float):
        self.interval = interval
        self.lock = threading.Lock()
        self.shots: dict[int, dict] = {}  # id -> {png, error, ts}
        self.fleet: list[dict] = []
        self.fleet_ts = 0.0
        self.last_client = time.time()
        self.png_native = True
        self.cycles = 0
        self.started = time.time()
        self.wake = threading.Event()

    def touch(self) -> None:
        idle = time.time() - self.last_client > IDLE_AFTER
        self.last_client = time.time()
        if idle:
            self.wake.set()  # first look after a nap: refresh right away

    def get_fleet(self) -> list[dict]:
        with self.lock:
            fresh = time.time() - self.fleet_ts < FLEET_TTL and self.fleet
            if fresh:
                return list(self.fleet)
        fleet = list_fleet()
        with self.lock:
            if fleet or not self.fleet:
                self.fleet, self.fleet_ts = fleet, time.time()
            return list(self.fleet)

    def cycle(self) -> None:
        fleet = self.get_fleet()
        if not fleet:
            return
        with ThreadPoolExecutor(max_workers=max(1, len(fleet))) as pool:
            results = list(pool.map(lambda g: (g["id"], screendump(g["id"], self.png_native)), fleet))
        now = time.time()
        with self.lock:
            for vmid, (png, err) in results:
                prev = self.shots.get(vmid, {})
                self.shots[vmid] = {
                    "png": png if png else prev.get("png"),
                    "fresh": bool(png),
                    "error": err,
                    "ts": now if png else prev.get("ts", 0.0),
                }
            live = {g["id"] for g in fleet}
            for gone in set(self.shots) - live:
                del self.shots[gone]
            self.cycles += 1
        # one-time downgrade if this QEMU has no `-f png`
        if self.png_native and all(not p for _, (p, _) in results):
            errs = " ".join(e for _, (_, e) in results)
            if "format" in errs.lower() or "invalid" in errs.lower() or "-f" in errs:
                self.png_native = False
                print("vmwatch: QEMU lacks screendump -f png; falling back to PPM", flush=True)

    def loop(self) -> None:
        while True:
            if time.time() - self.last_client <= IDLE_AFTER:
                try:
                    self.cycle()
                except Exception as e:  # noqa: BLE001 - the poller must never die
                    print(f"vmwatch: cycle error: {e}", file=sys.stderr, flush=True)
                self.wake.wait(self.interval)
            else:
                self.wake.wait(2.0)
            self.wake.clear()

    def state(self) -> list[dict]:
        now = time.time()
        out = []
        with self.lock:
            shots = dict(self.shots)
        for g in self.get_fleet():
            s = shots.get(g["id"], {})
            out.append(
                {
                    **g,
                    "ok": bool(s.get("png")),
                    "fresh": bool(s.get("fresh")),
                    "error": s.get("error") or "",
                    "age": round(now - s["ts"], 1) if s.get("ts") else None,
                    "ts": time.strftime("%H:%M:%S", time.localtime(s["ts"])) if s.get("ts") else "",
                }
            )
        return out

    def png(self, vmid: int) -> bytes | None:
        with self.lock:
            return (self.shots.get(vmid) or {}).get("png")


# --------------------------------------------------------------------------
# page
# --------------------------------------------------------------------------

PAGE = """<!doctype html>
<html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>%(title)s</title>
<link rel="icon" href="data:image/svg+xml,%%3Csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 16 16'%%3E%%3Crect width='16' height='16' rx='2' fill='%%232a2f39'/%%3E%%3Crect x='2' y='3' width='12' height='8' rx='1' fill='%%2368c07a'/%%3E%%3Crect x='6' y='12' width='4' height='1.5' fill='%%2398a0ad'/%%3E%%3C/svg%%3E">
<style>
 :root{--bg:#16181c;--tile:#22262e;--tile2:#2a2f39;--line:#3a4150;--fg:#e8eaee;
       --dim:#98a0ad;--bad:#4a2222;--badfg:#e08a8a;--ok:#68c07a}
 *{box-sizing:border-box}
 body{margin:0;background:var(--bg);color:var(--fg);
      font:13px/1.4 "Segoe UI",system-ui,-apple-system,sans-serif}
 header{display:flex;align-items:center;gap:12px;padding:8px 14px;
        background:#1c1f25;border-bottom:1px solid var(--line);
        position:sticky;top:0;z-index:5}
 header h1{font-size:14px;font-weight:600;margin:0;letter-spacing:.2px}
 header .sp{flex:1}
 header .meta{color:var(--dim);font-size:11px;font-variant-numeric:tabular-nums}
 .dot{width:7px;height:7px;border-radius:50%%;background:var(--ok);display:inline-block;
      margin-right:5px;transition:opacity .3s}
 .dot.blink{opacity:.25}
 main{display:grid;gap:10px;padding:10px;
      grid-template-columns:repeat(auto-fill,minmax(%(minw)spx,1fr))}
 main.one{grid-template-columns:1fr;padding:10px}
 .tile{background:var(--tile);border:1px solid var(--line);border-radius:5px;
       overflow:hidden;text-decoration:none;color:inherit;display:block}
 .tile.bad{background:var(--bad);border-color:#6a3030}
 .bar{display:flex;align-items:center;gap:8px;padding:5px 9px;background:var(--tile2);
      border-bottom:1px solid var(--line)}
 .tile.bad .bar{background:#3a1e1e;border-bottom-color:#6a3030}
 .id{font-weight:700;font-size:13px;font-variant-numeric:tabular-nums}
 .nm{color:var(--fg);font-size:12px}
 .st{margin-left:auto;color:var(--dim);font-size:11px;font-variant-numeric:tabular-nums;
     white-space:nowrap}
 .shot{display:block;width:100%%;aspect-ratio:4/3;object-fit:contain;background:#000}
 .ph{display:flex;flex-direction:column;align-items:center;justify-content:center;gap:6px;
     aspect-ratio:4/3;background:#2c1e1e;color:var(--badfg);text-align:center;padding:14px}
 .ph b{font-size:15px;letter-spacing:1.5px}
 .ph span{font-size:11px;color:#b89696;max-width:34em;word-break:break-word}
 a.back{color:var(--dim);text-decoration:none;font-size:12px}
 a.back:hover{color:var(--fg)}
</style></head><body>
<header>
  <span class="dot" id="dot"></span>
  <h1>%(title)s</h1>
  %(nav)s
  <span class="sp"></span>
  <span class="meta" id="meta">connecting...</span>
</header>
<main id="wall" class="%(cls)s"></main>
<script>
const ONLY = %(only)s, INTERVAL = %(interval)s;
const wall = document.getElementById('wall'), meta = document.getElementById('meta'),
      dot = document.getElementById('dot');
const esc = s => String(s).replace(/[&<>"]/g, c =>
  ({'&':'&amp;','<':'&lt;','>':'&gt;','"':'&quot;'}[c]));

function render(guests){
  const want = guests.map(g => String(g.id)).join(',');
  if (wall.dataset.keys !== want){ wall.innerHTML = ''; wall.dataset.keys = want;
    for (const g of guests){
      const el = document.createElement(ONLY ? 'div' : 'a');
      el.className = 'tile'; el.id = 'g' + g.id;
      if (!ONLY) el.href = '/one/' + g.id;
      el.innerHTML = '<div class="bar"><span class="id"></span><span class="nm"></span>'
                   + '<span class="st"></span></div><div class="body"></div>';
      wall.appendChild(el);
    }
  }
  for (const g of guests){
    const el = document.getElementById('g' + g.id); if (!el) continue;
    el.classList.toggle('bad', !g.ok);
    el.querySelector('.id').textContent = g.id;
    el.querySelector('.nm').textContent = g.name;
    el.querySelector('.st').textContent =
      g.status + (g.ok ? '  ' + g.ts + (g.fresh ? '' : '  (stale)') : '  no signal');
    const body = el.querySelector('.body');
    if (g.ok){
      let img = body.querySelector('img');
      if (!img){ body.innerHTML = ''; img = document.createElement('img');
        img.className = 'shot'; img.alt = 'VM ' + g.id;
        // an undecodable frame must degrade to the placeholder, not a broken icon
        img.onerror = () => { el.classList.add('bad'); body.innerHTML =
          '<div class="ph"><b>NO SIGNAL</b><span>frame did not decode</span></div>'; };
        body.appendChild(img); }
      // cache-buster keyed to capture time: no refetch unless the frame changed
      const next = '/shot/' + g.id + '.png?t=' + encodeURIComponent(g.ts);
      if (img.getAttribute('src') !== next) img.src = next;
    } else {
      body.innerHTML = '<div class="ph"><b>NO SIGNAL</b><span>'
        + esc(g.error || g.status) + '</span></div>';
    }
  }
}

async function tick(){
  try{
    const r = await fetch('/api/state', {cache:'no-store'});
    let guests = await r.json();
    if (ONLY) guests = guests.filter(g => g.id == ONLY);
    render(guests);
    const live = guests.filter(g => g.ok).length;
    meta.textContent = live + '/' + guests.length + ' framebuffers  ·  '
      + new Date().toLocaleTimeString() + '  ·  every ' + INTERVAL + 's  ·  spicy';
    dot.style.background = 'var(--ok)';
  } catch(e){
    meta.textContent = 'lost contact with vmwatch — retrying';
    dot.style.background = '#c05a5a';
  }
  dot.classList.add('blink'); setTimeout(() => dot.classList.remove('blink'), 300);
}
tick(); setInterval(tick, INTERVAL * 1000);
</script></body></html>
"""


class Handler(BaseHTTPRequestHandler):
    server_version = "vmwatch"
    tank: Tank = None  # type: ignore[assignment]

    def log_message(self, fmt, *a):  # quieter than the default one-line-per-image
        pass

    def _send(self, code, ctype, body: bytes, cache: str = "no-store"):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", cache)
        self.end_headers()
        try:
            self.wfile.write(body)
        except (BrokenPipeError, ConnectionResetError):
            pass

    def _page(self, only: int | None):
        title = f"SOTS guest {only}" if only else "SOTS lab guests"
        nav = '<a class="back" href="/">&larr; all guests</a>' if only else ""
        body = PAGE % {
            "title": html.escape(title),
            "nav": nav,
            "cls": "one" if only else "",
            "only": only if only else "null",
            "interval": self.tank.interval,
            "minw": 900 if only else 460,
        }
        self._send(HTTPStatus.OK, "text/html; charset=utf-8", body.encode())

    def do_GET(self):  # noqa: N802
        path = self.path.split("?", 1)[0]
        self.tank.touch()

        if path == "/":
            return self._page(None)
        if m := re.fullmatch(r"/one/(\d+)/?", path):
            return self._page(int(m.group(1)))
        if path == "/api/state":
            return self._send(
                HTTPStatus.OK, "application/json", json.dumps(self.tank.state()).encode()
            )
        if m := re.fullmatch(r"/shot/(\d+)\.png", path):
            png = self.tank.png(int(m.group(1)))
            if png:
                # immutable: the URL carries the capture timestamp
                return self._send(HTTPStatus.OK, "image/png", png, "max-age=30")
            return self._send(HTTPStatus.NOT_FOUND, "text/plain", b"no framebuffer\n")
        if path == "/healthz":
            t = self.tank
            up = int(time.time() - t.started)
            return self._send(
                HTTPStatus.OK,
                "text/plain",
                f"ok guests={len(t.get_fleet())} cycles={t.cycles} uptime={up}s\n".encode(),
            )
        return self._send(HTTPStatus.NOT_FOUND, "text/plain", b"not found\n")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--port", type=int, default=DEFAULT_PORT)
    ap.add_argument("--bind", default="0.0.0.0")
    ap.add_argument("--interval", type=float, default=DEFAULT_INTERVAL)
    args = ap.parse_args()

    tank = Tank(args.interval)
    Handler.tank = tank
    threading.Thread(target=tank.loop, daemon=True, name="poller").start()
    srv = ThreadingHTTPServer((args.bind, args.port), Handler)
    srv.daemon_threads = True
    print(
        f"vmwatch: http://{args.bind}:{args.port}/  interval={args.interval}s  "
        f"guests={[g['id'] for g in tank.get_fleet()]}",
        flush=True,
    )
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass
    return 0


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