#!/bin/bash
# Lightweight GPU fleet monitor — run on pve cron (every 5 min).
# Logs JSONL to /var/log/blair-fleet-monitor.jsonl; updates blair-share latest snapshot.
# Alerts in log when IoPsiFull or flush latency exceed thresholds.
set -euo pipefail

export BLAIR_SHARE="${BLAIR_SHARE:-/tank/blair-share}"
LOG="${BLAIR_FLEET_MONITOR_LOG:-/var/log/blair-fleet-monitor.jsonl}"
LATEST="${BLAIR_SHARE}/fleet/monitor/latest.json"
NODE="${BLAIR_PVE_NODE:-pve}"
GPU_VMIDS=(230 231 232 233 234 235)
IO_WARN=15
IO_CRIT=30
FLUSH_WARN_MS=20

if ! command -v python3 >/dev/null 2>&1; then
  echo 'python3 required' >&2
  exit 1
fi

exec python3 - <<'PY'
import json
import os
import subprocess
from datetime import datetime, timezone

LOG = os.environ.get('BLAIR_FLEET_MONITOR_LOG', '/var/log/blair-fleet-monitor.jsonl')
LATEST = os.path.join(os.environ.get('BLAIR_SHARE', '/tank/blair-share'), 'fleet/monitor/latest.json')
NODE = os.environ.get('BLAIR_PVE_NODE', 'pve')
GPU_VMIDS = [230, 231, 232, 233, 234, 235]
IO_WARN = int(os.environ.get('BLAIR_IO_WARN', '15'))
IO_CRIT = int(os.environ.get('BLAIR_IO_CRIT', '30'))
FLUSH_WARN_MS = float(os.environ.get('BLAIR_FLUSH_WARN_MS', '20'))


def pvesh_get(path: str) -> dict:
    proc = subprocess.run(
        ['pvesh', 'get', path, '--output-format', 'json'],
        capture_output=True, text=True, check=False,
    )
    if proc.returncode != 0:
        return {}
    return json.loads(proc.stdout)


def blockstat_latency(blockstat: dict) -> list[dict]:
    rows = []
    if not blockstat:
        return rows
    for dev, d in blockstat.items():
        if not isinstance(d, dict):
            continue
        rd_ops = float(d.get('rd_operations') or 0)
        wr_ops = float(d.get('wr_operations') or 0)
        fl_ops = float(d.get('flush_operations') or 0)
        rows.append({
            'device': dev,
            'rd_avg_ms': round(float(d.get('rd_total_time_ns') or 0) / rd_ops / 1e6, 2) if rd_ops else 0,
            'wr_avg_ms': round(float(d.get('wr_total_time_ns') or 0) / wr_ops / 1e6, 2) if wr_ops else 0,
            'flush_avg_ms': round(float(d.get('flush_total_time_ns') or 0) / fl_ops / 1e6, 2) if fl_ops else 0,
            'rd_ops': int(rd_ops),
            'wr_ops': int(wr_ops),
            'flush_ops': int(fl_ops),
        })
    return rows


def main() -> int:
    now = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
    alerts: list[dict] = []
    host = pvesh_get(f'/nodes/{NODE}/status')
    stor = pvesh_get(f'/nodes/{NODE}/storage')
    pool = next((s for s in stor if s.get('storage') == 'local-2tb'), {})
    pool_pct = round(100 * pool.get('used', 0) / pool['total'], 1) if pool.get('total') else None

    host_summary = {
        'cpu_pct': round(float(host.get('cpu', 0)) * 100, 1),
        'mem_used_gb': round(float(host.get('memory', {}).get('used', 0)) / (1024**3), 1),
        'mem_total_gb': round(float(host.get('memory', {}).get('total', 0)) / (1024**3), 1),
        'pool_used_pct': pool_pct,
    }

    gpu_rows = []
    for vmid in GPU_VMIDS:
        cfg = pvesh_get(f'/nodes/{NODE}/qemu/{vmid}/config')
        cur = pvesh_get(f'/nodes/{NODE}/qemu/{vmid}/status/current')
        if not cur or cur.get('status') != 'running':
            gpu_rows.append({'vmid': vmid, 'name': cfg.get('name'), 'status': 'stopped'})
            continue
        io_full = round(float(cur.get('pressureiofull') or 0), 1)
        io_some = round(float(cur.get('pressureiosome') or 0), 1)
        lat = blockstat_latency(cur.get('blockstat') or {})
        disk = cfg.get('sata0') or cfg.get('scsi0') or cfg.get('virtio0')
        hostpci = next((v for k, v in cfg.items() if str(k).startswith('hostpci')), '')
        row = {
            'vmid': vmid,
            'name': cur.get('name') or cfg.get('name'),
            'status': 'running',
            'cpu_pct': round(float(cur.get('cpu', 0)) * 100, 1),
            'mem_gb': round(float(cur.get('mem', 0)) / (1024**3), 2),
            'io_psi_full': io_full,
            'io_psi_some': io_some,
            'uptime_h': round(float(cur.get('uptime', 0)) / 3600, 1),
            'vga': cfg.get('vga'),
            'disk': disk,
            'gpu_mapping': hostpci,
            'latency': lat,
        }
        gpu_rows.append(row)
        if io_full >= IO_CRIT:
            alerts.append({'level': 'CRIT', 'vmid': vmid, 'msg': f'IO stall {io_full}%'})
        elif io_full >= IO_WARN:
            alerts.append({'level': 'WARN', 'vmid': vmid, 'msg': f'IO stall {io_full}%'})
        for l in lat:
            if l['flush_avg_ms'] >= FLUSH_WARN_MS and l['flush_ops'] > 10:
                alerts.append({
                    'level': 'WARN', 'vmid': vmid,
                    'msg': f"{l['device']} flush avg {l['flush_avg_ms']}ms",
                })

    record = {
        'timestamp': now,
        'host': host_summary,
        'gpu_fleet': gpu_rows,
        'alerts': alerts,
    }

    os.makedirs(os.path.dirname(LOG), exist_ok=True)
    with open(LOG, 'a', encoding='utf-8') as fh:
        fh.write(json.dumps(record, separators=(',', ':')) + '\n')

    os.makedirs(os.path.dirname(LATEST), exist_ok=True)
    tmp = LATEST + '.tmp'
    with open(tmp, 'w', encoding='utf-8') as fh:
        json.dump(record, fh, indent=2)
        fh.write('\n')
    os.replace(tmp, LATEST)

    if alerts:
        for a in alerts:
            print(f"[{a['level']}] VM {a.get('vmid', '-')} {a['msg']}")
    return 0


if __name__ == '__main__':
    raise SystemExit(main())
PY