#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# uvmt_poll.py - poll the live state of the tracking semaphores the CUDA
# client spins on, from the guest trace kptrs, via QMP (no guest agent
# besides the trace read). Prints the payload word the client actually
# reads (pool page + index*4) on every change, so the last stall before a
# failure is visible.
#
# Usage: tools/guest-diag/uvmt_poll.py [interval-seconds]
import os
import re
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from semfix_watch import Qmp, QMP_SOCK, is_kern, is_gpu  # noqa: E402

VMEXEC = os.path.normpath(os.path.join(HERE, "..", "..", "scripts", "vgpu_vm.py"))


def log(msg):
    print("[%s] %s" % (time.strftime("%H:%M:%S"), msg), flush=True)


def trace_sems():
    inner = ("tail -c 262144 /sys/kernel/tracing/trace "
             "| grep -E \"uvmt\" | tail -40")
    try:
        r = subprocess.run([VMEXEC, "exec", "sudo sh -c '%s'" % inner],
                           capture_output=True, text=True, timeout=60)
        out = r.stdout
    except Exception:
        return set()
    return {int(m, 16) for m in re.findall(r"sem=0x([0-9a-fA-F]+)", out)}


def state(q, sem):
    vals = q.read64(sem, 12)
    if len(vals) < 2 or not is_kern(vals[0]):
        return None
    page = vals[0]
    pv = q.read64(page)
    if not pv or not is_kern(pv[0]):
        return None
    mem = pv[0]
    vas = q.read64(mem + 0x10, 16)
    cpu = next((v for v in vas if is_kern(v)), None)
    gpu = [v for v in vas if is_gpu(v)]
    idx = vals[1] & 0xffff
    payload = None
    if cpu:
        w = q.read64(cpu + idx * 4)
        payload = w[0] if w else None
    words = " ".join("+%x=0x%x" % (4 * i, vals[i])
                     for i in range(2, min(len(vals), 12)))
    return (page, mem, idx, gpu, payload, words)


def main():
    interval = float(sys.argv[1]) if len(sys.argv) > 1 else 2.0
    q = Qmp(QMP_SOCK)
    last = {}
    sems_seen = set()
    log("polling uvmt sems (interval %.1fs)" % interval)
    while True:
        sems = trace_sems()
        if not sems and not sems_seen:
            time.sleep(interval)
            continue
        for sem in sems:
            st = None
            try:
                st = state(q, sem)
            except Exception:
                st = None
            key = st[:3] + (st[4],) if st else None
            if key != last.get(sem):
                last[sem] = key
                if st:
                    print("uvmt sem=0x%x page=0x%x mem=0x%x idx=%d gpu=%s "
                          "payload=%s %s"
                          % (sem, st[0], st[1], st[2],
                             ["0x%x" % v for v in st[3]],
                             "0x%x" % st[4] if st[4] is not None else "?",
                             st[5]), flush=True)
                else:
                    print("uvmt sem=0x%x state unreadable" % sem, flush=True)
        sems_seen |= sems
        time.sleep(interval)


if __name__ == "__main__":
    main()