#!/usr/bin/env python3
# semfix_watch.py - keep UVM semaphore-pool VA->PA translations fresh from
# guest truth (host-side dev tool, no guest changes).
#
# Why: the emulation's content-scanned VA->PA map can hold a *stale*
# translation for a semaphore-pool page (same VA, previous client's page).
# The GPU-tracking-semaphore waits in nvidia-uvm then never clear because the
# release writes land in the old page. The guest kernel knows the real
# backing: the tracking semaphore object's `page->memory->vas[]` array holds
# both the kernel CPU VA and the GPU VA of the pool page, and QEMU can turn
# the CPU VA into the GPA (`gva2gpa`). Pinning that pair with sema_map_add
# makes the service's release path write the page the client actually polls.
#
# Sources:
#   - guest: kprobe `uvmt` (uvm_gpu_tracking_semaphore_update_completed_value)
#     trace lines carry the tracking-semaphore kernel pointer as `sem=0x...`.
#   - host: QEMU QMP `human-monitor-command` for `x/..gx` memory reads and
#     `gva2gpa`.
#
# Usage: tools/guest-diag/semfix_watch.py [interval-seconds]
# Runs until killed; prints one line per newly resolved/pinned page.

import json
import os
import re
import socket
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
VMEXEC = os.path.normpath(os.path.join(HERE, "..", "..", "scripts", "vgpu_vm.py"))
PARAM = "/sys/module/open_iov_vfio/parameters"
QMP_SOCK = os.environ.get("QMP_SOCK", "/tmp/qemu-qmp.sock")

GPU_VA_LO = 0x100000000
GPU_VA_HI = 0x2000000000


class Qmp:
    def __init__(self, path):
        self.s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.s.connect(path)
        self.s.settimeout(15)
        self._rd()
        self.cmd("qmp_capabilities")

    def _rd(self):
        time.sleep(0.05)
        try:
            return self.s.recv(1 << 20).decode()
        except Exception:
            return ""

    def cmd(self, name, **args):
        try:
            self.s.sendall(json.dumps({"execute": name, "arguments": args}).encode() + b"\n")
            return json.loads(self._rd() or "{}")
        except Exception:
            return {}

    def hmp(self, line):
        return self.cmd("human-monitor-command", **{"command-line": line}).get("return", "")

    def read64(self, addr, n=1):
        out = self.hmp("x/%dgx 0x%x" % (n, addr))
        if ":" not in out:
            return []
        body = out.split(":", 1)[1]
        return [int(m, 16) for m in re.findall(r"0x([0-9a-fA-F]+)", body)]

    def gva2gpa(self, va):
        out = self.hmp("gva2gpa 0x%x" % va)
        m = re.search(r"gpa: 0x([0-9a-fA-F]+)", out)
        return int(m.group(1), 16) if m else None


def guest_sems():
    """Kernel pointers of tracking semaphores the client is spinning on."""
    cmd = ("sudo sh -c 'echo > /sys/kernel/tracing/trace; sleep 2; "
           "grep uvmt /sys/kernel/tracing/trace | tail -4'")
    try:
        out = subprocess.run([VMEXEC, "exec", cmd], capture_output=True,
                             text=True, timeout=60).stdout
    except Exception:
        return set()
    return {int(m, 16) for m in re.findall(r"sem=0x([0-9a-fA-F]+)", out)}


def pin(va, pa):
    cmd = "printf '0x%x,0x%x' > %s/sema_map_add" % (va, pa, PARAM)
    return subprocess.run(["sudo", "-n", "sh", "-c", cmd],
                          capture_output=True, text=True).returncode == 0


def resolve(q, sem):
    """(gpu_va, gpa) pairs for the pool page behind this tracking semaphore."""
    page = q.read64(sem)
    if not page or not page[0]:
        return []
    mem = q.read64(page[0])
    if not mem or not mem[0]:
        return []
    vas = q.read64(mem[0] + 0x10, 16)
    if not vas:
        return []
    cpu_kva = next((v for v in vas if v and (v >> 56) == 0xFF), None)
    gpu_vas = [v for v in vas if v and GPU_VA_LO <= v < GPU_VA_HI]
    if not cpu_kva or not gpu_vas:
        return []
    gpa = q.gva2gpa(cpu_kva)
    if not gpa:
        return []
    return [(gv, gpa) for gv in gpu_vas]


def main():
    interval = float(sys.argv[1]) if len(sys.argv) > 1 else 3.0
    q = Qmp(QMP_SOCK)
    seen = set()
    print("semfix: watching (interval %.1fs)" % interval, flush=True)
    while True:
        for sem in guest_sems() - seen:
            seen.add(sem)
            try:
                pairs = resolve(q, sem)
            except Exception as e:
                print("semfix: sem=0x%x resolve failed: %s" % (sem, e), flush=True)
                continue
            if not pairs:
                print("semfix: sem=0x%x unresolved (retry)" % sem, flush=True)
                seen.discard(sem)
                continue
            for gv, gpa in pairs:
                ok = pin(gv, gpa)
                print("semfix: sem=0x%x va=0x%x -> gpa=0x%x pin=%s"
                      % (sem, gv, gpa, "ok" if ok else "FAILED"), flush=True)
        time.sleep(interval)


if __name__ == "__main__":
    main()