#!/usr/bin/env python3
"""Zero-dependency request-load monitor for the single-slot LLM endpoints.

Why this exists (fox 2026-05-21): hermes (3090, vLLM) and qwen (4090,
llama.cpp) are single-slot endpoints fed real internet traffic through a
Caddy/nginx proxy (ai.foxhop.net). vLLM's per-request counters reset on
every restart, so a *time-series* is the only way to answer "are we being
swamped?". The swamp signal is the **queue depth** (`num_requests_waiting`
on vLLM, `requests_deferred` on llama.cpp): a single GPU slot that can't
keep up backs requests up and e2e latency balloons.

Three subcommands, all stdlib (urllib + sqlite3 + hand-rolled SVG) so it
runs in arborist's python+sqlite3 core with nothing to install:

  * ``poll``   — scrape each endpoint's Prometheus ``/metrics`` on an
                 interval, store one snapshot row per target in SQLite.
                 Prunes rows older than ``--retention-days`` every cycle
                 (bounded store — no cancer growth).
  * ``graph``  — render the stored snapshots to a multi-panel SVG:
                 queue depth, requests running, req/s, tok/s, e2e latency.
                 Deltas are computed at graph time from the raw counters.
  * ``access`` — parse a Caddy (JSON) or nginx/Apache (combined) access
                 log for the **real client IPs** the backend can't see
                 (everything collapses to the proxy IP at the backend),
                 tally req/min per IP, print the top talkers and render a
                 per-IP request-rate SVG. This is the "who is hammering
                 us" view; the backend ``/metrics`` is the "how much".

The backend answers HOW MUCH; the proxy log answers WHO. Use both.

Examples::

    # poll both GPU boxes every 10s into the default DB, forever
    python3 bench/load_monitor.py poll

    # one-shot sample (smoke test)
    python3 bench/load_monitor.py poll --once

    # render the last 6h to an SVG
    python3 bench/load_monitor.py graph --since-hours 6 --out load.svg

    # who hit the proxy (copy the Caddy log over first, or run on .27)
    python3 bench/load_monitor.py access /var/log/caddy/access.log --top 20
"""
from __future__ import annotations

import argparse
import json
import os
import sqlite3
import subprocess
import sys
import time
import urllib.request
from datetime import datetime, timezone

# --- default targets -------------------------------------------------------
# name -> Prometheus /metrics URL. Override with repeated --target name=url.
DEFAULT_TARGETS = {
    "hermes-3090": "http://3090-ai.foxhop.net:18888/metrics",
    "qwen-4090": "http://ai.foxhop.net:18888/metrics",
}

DEFAULT_DB = os.path.expanduser("~/.arborist/load_monitor.db")

# Map our common schema field -> ordered candidate Prometheus metric names.
# First name found in the scrape wins (vLLM first, llama.cpp fallback).
FIELD_CANDIDATES = {
    "running": ["vllm:num_requests_running", "llamacpp:requests_processing"],
    "waiting": ["vllm:num_requests_waiting", "llamacpp:requests_deferred"],
    "success": ["vllm:request_success_total"],
    "prompt_tok": ["vllm:prompt_tokens_total", "llamacpp:prompt_tokens_total"],
    "gen_tok": ["vllm:generation_tokens_total", "llamacpp:tokens_predicted_total"],
    "lat_sum": ["vllm:e2e_request_latency_seconds_sum"],
    "lat_count": ["vllm:e2e_request_latency_seconds_count"],
}
FIELDS = list(FIELD_CANDIDATES)

# series rendered by `graph`: (db_field_or_derived, panel title, unit)
PANELS = [
    ("waiting", "queue depth (waiting) — SWAMP SIGNAL", "reqs"),
    ("gpu_power", "GPU power draw — SATURATION SIGNAL (vs cap)", "W"),
    ("gpu_util", "GPU utilization", "%"),
    ("req_per_s", "request rate", "req/s"),
    ("tok_per_s", "generation throughput", "tok/s"),
    ("latency", "mean e2e latency", "s"),
]

COLORS = ["#2563eb", "#dc2626", "#16a34a", "#9333ea", "#ea580c", "#0891b2"]


# --- prometheus parsing ----------------------------------------------------
def parse_prometheus(text: str) -> dict[str, float]:
    """Parse Prometheus text exposition into {metric_name: summed_value}.

    Labels are stripped and values summed across label sets, so a counter
    split by reason (vLLM's request_success_total has finished/aborted
    label sets) collapses to one total. Tolerant of label values that
    contain spaces (value is taken after the closing brace)."""
    out: dict[str, float] = {}
    for raw in text.splitlines():
        line = raw.strip()
        if not line or line[0] == "#":
            continue
        if "{" in line:
            name = line[: line.index("{")]
            rest = line[line.rindex("}") + 1 :].strip()
        else:
            sp = line.split(None, 1)
            name = sp[0]
            rest = sp[1] if len(sp) > 1 else ""
        if not rest:
            continue
        try:
            val = float(rest.split()[0])
        except (ValueError, IndexError):
            continue
        out[name] = out.get(name, 0.0) + val
    return out


def extract_fields(metrics: dict[str, float]) -> dict[str, float | None]:
    row: dict[str, float | None] = {}
    for field, candidates in FIELD_CANDIDATES.items():
        row[field] = next((metrics[c] for c in candidates if c in metrics), None)
    return row


# --- storage ---------------------------------------------------------------
def connect(db_path: str) -> sqlite3.Connection:
    os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True)
    con = sqlite3.connect(db_path)
    con.execute(
        """CREATE TABLE IF NOT EXISTS samples (
               ts REAL NOT NULL,
               target TEXT NOT NULL,
               ok INTEGER NOT NULL,
               running REAL, waiting REAL, success REAL,
               prompt_tok REAL, gen_tok REAL,
               lat_sum REAL, lat_count REAL)"""
    )
    con.execute("CREATE INDEX IF NOT EXISTS idx_samples_ts ON samples(ts)")
    # GPU columns added later — migrate an older DB in place (idempotent).
    have = {r[1] for r in con.execute("PRAGMA table_info(samples)")}
    for col in ("gpu_util", "gpu_power", "gpu_power_limit", "gpu_mem", "gpu_temp"):
        if col not in have:
            con.execute(f"ALTER TABLE samples ADD COLUMN {col} REAL")
    con.commit()
    return con


def scrape(url: str, timeout: float) -> dict[str, float] | None:
    try:
        with urllib.request.urlopen(url, timeout=timeout) as resp:
            return parse_prometheus(resp.read().decode("utf-8", "replace"))
    except Exception:
        return None


def sample_gpu(host: str, timeout: float) -> dict[str, float] | None:
    """SSH to a GPU box and read nvidia-smi (first GPU). Stdlib subprocess —
    nvidia-smi must run on the box; the poller drives it over SSH."""
    query = "utilization.gpu,power.draw,power.limit,memory.used,temperature.gpu"
    cmd = ["ssh", "-o", "BatchMode=yes", "-o", f"ConnectTimeout={int(timeout)}",
           host, f"nvidia-smi --query-gpu={query} --format=csv,noheader,nounits"]
    try:
        out = subprocess.run(cmd, capture_output=True, text=True,
                             timeout=timeout + 5).stdout.strip()
        first = out.splitlines()[0]
        u, p, plim, mem, temp = (float(x) for x in first.split(","))
        return {"gpu_util": u, "gpu_power": p, "gpu_power_limit": plim,
                "gpu_mem": mem, "gpu_temp": temp}
    except Exception:
        return None


def insert_sample(con, ts, target, fields_or_none, gpu=None):
    g = gpu or {}
    if fields_or_none is None and not g:
        con.execute(
            "INSERT INTO samples (ts, target, ok) VALUES (?, ?, 0)", (ts, target)
        )
        return
    f = fields_or_none or {k: None for k in FIELDS}
    con.execute(
        """INSERT INTO samples
           (ts, target, ok, running, waiting, success,
            prompt_tok, gen_tok, lat_sum, lat_count,
            gpu_util, gpu_power, gpu_power_limit, gpu_mem, gpu_temp)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
        (ts, target, 1 if fields_or_none is not None else 0,
         f["running"], f["waiting"], f["success"],
         f["prompt_tok"], f["gen_tok"], f["lat_sum"], f["lat_count"],
         g.get("gpu_util"), g.get("gpu_power"), g.get("gpu_power_limit"),
         g.get("gpu_mem"), g.get("gpu_temp")),
    )


def cmd_poll(args) -> int:
    targets = parse_targets(args.target) or DEFAULT_TARGETS
    gpus = parse_targets(args.gpu) or {}
    con = connect(args.db)
    retain_s = args.retention_days * 86400
    print(f"polling {len(targets)} target(s) every {args.interval}s -> {args.db}",
          flush=True)
    for name, url in targets.items():
        gtag = f"  +gpu via {gpus[name]}" if name in gpus else ""
        print(f"  {name}: {url}{gtag}", flush=True)
    start = time.time()
    try:
        while True:
            ts = time.time()
            line = [datetime.now(timezone.utc).strftime("%H:%M:%S")]
            for name, url in targets.items():
                metrics = scrape(url, args.timeout)
                gpu = sample_gpu(gpus[name], args.timeout) if name in gpus else None
                f = extract_fields(metrics) if metrics is not None else None
                insert_sample(con, ts, name, f, gpu)
                if metrics is None and gpu is None:
                    line.append(f"{name}=DOWN")
                else:
                    parts = [name]
                    if f is not None:
                        parts.append(f"wait={fmt(f['waiting'])} run={fmt(f['running'])}")
                    if gpu is not None:
                        parts.append(f"gpu={fmt(gpu['gpu_power'])}W/{fmt(gpu['gpu_util'])}%")
                    line.append(" ".join(parts))
            if retain_s > 0:
                con.execute("DELETE FROM samples WHERE ts < ?", (ts - retain_s,))
            con.commit()
            print("  ".join(line), flush=True)
            if args.once:
                break
            if args.duration and (time.time() - start) >= args.duration:
                break
            time.sleep(args.interval)
    except KeyboardInterrupt:
        print("\nstopped.", flush=True)
    finally:
        con.close()
    return 0


def fmt(v) -> str:
    return "-" if v is None else (f"{v:.0f}" if v == int(v) else f"{v:.1f}")


# --- graphing --------------------------------------------------------------
def load_series(con, since_hours: float):
    cutoff = time.time() - since_hours * 3600
    rows = con.execute(
        """SELECT ts, target, ok, running, waiting, success, gen_tok,
                  lat_sum, lat_count, gpu_power, gpu_util
           FROM samples WHERE ts >= ? ORDER BY target, ts""",
        (cutoff,),
    ).fetchall()
    by_target: dict[str, list] = {}
    for row in rows:
        by_target.setdefault(row[1], []).append(row)
    return by_target


def derive_points(rows, panel_field):
    """Return [(ts, value), ...] for a panel from one target's raw rows.

    Gauges (waiting/running) read directly; rate panels differentiate the
    cumulative counters between consecutive ok samples."""
    pts = []
    prev = None
    gauges = {"waiting": None, "running": None, "gpu_power": None, "gpu_util": None}
    for r in rows:
        (ts, _t, ok, running, waiting, success, gen_tok,
         lat_sum, lat_count, gpu_power, gpu_util) = r
        # GPU rows may carry ok=0 (metrics down) but still hold gpu data.
        if not ok and panel_field not in ("gpu_power", "gpu_util"):
            prev = None  # gap the line across downtime
            continue
        gauges.update(waiting=waiting, running=running,
                      gpu_power=gpu_power, gpu_util=gpu_util)
        if panel_field in gauges:
            v = gauges[panel_field]
        elif panel_field in ("req_per_s", "tok_per_s", "latency"):
            v = None
            if prev is not None:
                dt = ts - prev[0]
                if dt > 0:
                    if panel_field == "req_per_s" and success is not None and prev[1] is not None:
                        d = success - prev[1]
                        v = d / dt if d >= 0 else None
                    elif panel_field == "tok_per_s" and gen_tok is not None and prev[2] is not None:
                        d = gen_tok - prev[2]
                        v = d / dt if d >= 0 else None
                    elif panel_field == "latency" and lat_count is not None and prev[3] is not None:
                        dc = lat_count - prev[3]
                        ds = lat_sum - prev[4] if (lat_sum is not None and prev[4] is not None) else None
                        v = (ds / dc) if (dc and dc > 0 and ds is not None) else None
            prev = (ts, success, gen_tok, lat_count, lat_sum)
        else:
            v = None
        if panel_field in gauges:
            prev = (ts,)
        if v is not None:
            pts.append((ts, v))
    return pts


def svg_panel(target_pts, title, unit, x0, y0, w, h, t_min, t_max):
    """Render one panel (all targets overlaid) as SVG fragment string."""
    parts = []
    # frame + title
    parts.append(f'<rect x="{x0}" y="{y0}" width="{w}" height="{h}" '
                 f'fill="#fafafa" stroke="#ccc"/>')
    parts.append(f'<text x="{x0+6}" y="{y0+16}" font-size="13" '
                 f'font-family="sans-serif" font-weight="bold">{esc(title)}</text>')
    all_vals = [v for pts in target_pts.values() for _, v in pts]
    v_max = max(all_vals) if all_vals else 1.0
    v_max = v_max if v_max > 0 else 1.0
    pad_top = 24
    plot_h = h - pad_top - 8
    plot_y0 = y0 + pad_top

    def sx(t):
        return x0 + 4 + (w - 8) * ((t - t_min) / (t_max - t_min) if t_max > t_min else 0)

    def sy(v):
        return plot_y0 + plot_h * (1 - v / v_max)

    # y-axis max label
    parts.append(f'<text x="{x0+w-4}" y="{plot_y0+10}" font-size="10" '
                 f'fill="#888" text-anchor="end" font-family="sans-serif">'
                 f'{v_max:.2f} {esc(unit)}</text>')
    for i, (name, pts) in enumerate(sorted(target_pts.items())):
        color = COLORS[i % len(COLORS)]
        if pts:
            d = " ".join(f"{sx(t):.1f},{sy(v):.1f}" for t, v in pts)
            parts.append(f'<polyline points="{d}" fill="none" '
                         f'stroke="{color}" stroke-width="1.5"/>')
        # legend dot
        ly = plot_y0 + 12 + i * 14
        parts.append(f'<rect x="{x0+8}" y="{ly-8}" width="10" height="10" fill="{color}"/>')
        parts.append(f'<text x="{x0+22}" y="{ly}" font-size="11" '
                     f'font-family="sans-serif" fill="#333">{esc(name)}</text>')
    return "\n".join(parts)


def esc(s: str) -> str:
    return (str(s).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;"))


def cmd_graph(args) -> int:
    con = connect(args.db)
    by_target = load_series(con, args.since_hours)
    con.close()
    if not by_target:
        print("no samples in window — run `poll` first.", file=sys.stderr)
        return 1
    all_ts = [r[0] for rows in by_target.values() for r in rows]
    t_min, t_max = min(all_ts), max(all_ts)

    W = 900
    panel_h = 150
    margin = 20
    H = margin * 2 + 40 + panel_h * len(PANELS) + 10 * (len(PANELS) - 1)
    out = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{W}" height="{H}" '
           f'viewBox="0 0 {W} {H}">']
    out.append(f'<rect width="{W}" height="{H}" fill="white"/>')
    span_min = (t_max - t_min) / 60.0
    head = (f"LLM endpoint load — {datetime.fromtimestamp(t_min, timezone.utc):%Y-%m-%d %H:%M} "
            f"to {datetime.fromtimestamp(t_max, timezone.utc):%H:%M} UTC "
            f"({span_min:.0f} min, {len(all_ts)} samples)")
    out.append(f'<text x="{margin}" y="26" font-size="15" font-weight="bold" '
               f'font-family="sans-serif">{esc(head)}</text>')

    y = margin + 40
    pw = W - margin * 2
    for field, title, unit in PANELS:
        tp = {name: derive_points(rows, field) for name, rows in by_target.items()}
        out.append(svg_panel(tp, title, unit, margin, y, pw, panel_h, t_min, t_max))
        y += panel_h + 10
    out.append("</svg>")

    dest = args.out or os.path.join(
        "bench", "results",
        f"load_monitor_{datetime.now(timezone.utc):%Y-%m-%dT%H-%M-%SZ}.svg")
    os.makedirs(os.path.dirname(dest) or ".", exist_ok=True)
    with open(dest, "w") as fh:
        fh.write("\n".join(out))
    print(f"wrote {dest}  ({len(all_ts)} samples, {span_min:.0f} min, "
          f"targets: {', '.join(sorted(by_target))})")
    return 0


# --- access-log parsing (proxy real-client-IP view) ------------------------
def parse_access_line(line: str):
    """Return (epoch_ts, client_ip) from a Caddy-JSON or combined log line.

    Caddy logs JSON; prefers X-Forwarded-For (real client) over remote_ip
    (the upstream hop). nginx/Apache 'combined' format: IP is field 1."""
    line = line.strip()
    if not line:
        return None
    if line[0] == "{":
        try:
            o = json.loads(line)
        except ValueError:
            return None
        req = o.get("request", {})
        xff = (req.get("headers", {}) or {}).get("X-Forwarded-For")
        if xff:
            ip = (xff[0] if isinstance(xff, list) else xff).split(",")[0].strip()
        else:
            ip = req.get("remote_ip") or req.get("client_ip") or \
                 (req.get("remote_addr", "").rsplit(":", 1)[0])
        ts = o.get("ts")
        try:
            ts = float(ts) if ts is not None else None
        except (TypeError, ValueError):
            ts = None
        return (ts, ip) if ip else None
    # combined log format: IP - - [10/Oct/2000:13:55:36 -0700] "GET ..." ...
    sp = line.split()
    if len(sp) < 4:
        return None
    ip = sp[0]
    ts = None
    try:
        # [10/Oct/2000:13:55:36 +0000]
        raw = line.split("[", 1)[1].split("]", 1)[0]
        ts = datetime.strptime(raw, "%d/%b/%Y:%H:%M:%S %z").timestamp()
    except (IndexError, ValueError):
        ts = None
    return (ts, ip)


def is_lan(ip: str) -> bool:
    return ip.startswith(("192.168.", "10.", "127.")) or ip.startswith(
        tuple(f"172.{n}." for n in range(16, 32)))


def cmd_access(args) -> int:
    counts: dict[str, int] = {}
    per_min: dict[str, dict[int, int]] = {}
    total = 0
    with open(args.logfile, errors="replace") as fh:
        for line in fh:
            parsed = parse_access_line(line)
            if not parsed:
                continue
            ts, ip = parsed
            total += 1
            counts[ip] = counts.get(ip, 0) + 1
            if ts is not None:
                minute = int(ts // 60)
                per_min.setdefault(ip, {})[minute] = \
                    per_min.setdefault(ip, {}).get(minute, 0) + 1

    if not total:
        print("no parseable request lines found.", file=sys.stderr)
        return 1
    top = sorted(counts.items(), key=lambda kv: kv[1], reverse=True)[: args.top]
    print(f"{total} requests, {len(counts)} distinct client IPs\n")
    print(f"{'requests':>9}  {'%':>5}  scope  client IP")
    for ip, n in top:
        scope = "LAN " if is_lan(ip) else "WAN*"
        print(f"{n:>9}  {100*n/total:>4.1f}%  {scope}  {ip}")
    wan = sum(n for ip, n in counts.items() if not is_lan(ip))
    print(f"\nWAN (internet) requests: {wan} ({100*wan/total:.1f}%)  "
          f"| LAN: {total-wan} ({100*(total-wan)/total:.1f}%)")

    if args.out and per_min:
        render_access_svg(per_min, top, args.out)
        print(f"\nwrote {args.out}")
    elif args.out:
        print("\n(no timestamps parsed — cannot render per-minute SVG)",
              file=sys.stderr)
    return 0


def render_access_svg(per_min, top, dest):
    W, H, margin = 900, 360, 50
    mins = sorted({m for d in per_min.values() for m in d})
    if not mins:
        return
    m0, m1 = mins[0], mins[-1]
    peak = max((c for d in per_min.values() for c in d.values()), default=1)
    out = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{W}" height="{H}">',
           f'<rect width="{W}" height="{H}" fill="white"/>',
           f'<text x="{margin}" y="26" font-size="15" font-weight="bold" '
           f'font-family="sans-serif">proxy requests/min by client IP '
           f'(top {len(top)}) — peak {peak}/min</text>']
    px0, py0 = margin, 50
    pw, ph = W - margin * 2, H - 90

    def sx(m):
        return px0 + pw * ((m - m0) / (m1 - m0) if m1 > m0 else 0)

    def sy(c):
        return py0 + ph * (1 - c / peak)

    out.append(f'<rect x="{px0}" y="{py0}" width="{pw}" height="{ph}" '
               f'fill="#fafafa" stroke="#ccc"/>')
    for i, (ip, _n) in enumerate(top):
        color = COLORS[i % len(COLORS)]
        d = per_min.get(ip, {})
        pts = " ".join(f"{sx(m):.1f},{sy(d.get(m,0)):.1f}" for m in mins)
        out.append(f'<polyline points="{pts}" fill="none" stroke="{color}" '
                   f'stroke-width="1.5"/>')
        ly = py0 + 14 + i * 14
        if ly < py0 + ph - 4:
            out.append(f'<rect x="{px0+8}" y="{ly-8}" width="10" height="10" fill="{color}"/>')
            out.append(f'<text x="{px0+22}" y="{ly}" font-size="11" '
                       f'font-family="sans-serif">{esc(ip)}</text>')
    out.append("</svg>")
    os.makedirs(os.path.dirname(dest) or ".", exist_ok=True)
    with open(dest, "w") as fh:
        fh.write("\n".join(out))


# --- cli -------------------------------------------------------------------
def parse_targets(items):
    if not items:
        return None
    out = {}
    for it in items:
        if "=" not in it:
            raise SystemExit(f"--target must be name=url, got {it!r}")
        name, url = it.split("=", 1)
        out[name] = url
    return out


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--db", default=DEFAULT_DB, help=f"SQLite path (default {DEFAULT_DB})")
    sub = p.add_subparsers(dest="cmd", required=True)

    pp = sub.add_parser("poll", help="scrape /metrics on an interval")
    pp.add_argument("--target", action="append",
                    help="name=url (repeatable; default hermes+qwen)")
    pp.add_argument("--gpu", action="append",
                    help="name=sshhost (repeatable) — sample nvidia-smi on that "
                         "box for target 'name'; name must match a --target")
    pp.add_argument("--interval", type=float, default=10.0)
    pp.add_argument("--timeout", type=float, default=5.0)
    pp.add_argument("--retention-days", type=float, default=7.0,
                    help="prune samples older than this (0 = keep all)")
    pp.add_argument("--once", action="store_true", help="single sample then exit")
    pp.add_argument("--duration", type=float, default=0.0,
                    help="stop after N seconds (0 = run until Ctrl-C)")
    pp.set_defaults(func=cmd_poll)

    pg = sub.add_parser("graph", help="render stored samples to SVG")
    pg.add_argument("--since-hours", type=float, default=6.0)
    pg.add_argument("--out", help="SVG path (default bench/results/load_monitor_<ts>.svg)")
    pg.set_defaults(func=cmd_graph)

    pa = sub.add_parser("access", help="parse a proxy access log for real client IPs")
    pa.add_argument("logfile")
    pa.add_argument("--top", type=int, default=20)
    pa.add_argument("--out", help="optional per-minute-by-IP SVG path")
    pa.set_defaults(func=cmd_access)

    args = p.parse_args(argv)
    return args.func(args)


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