#!/usr/bin/env python3
# semfix_watch.py - keep UVM VA->PA translations fresh from guest truth
# (host-side dev tool, no guest changes).
#
# Two jobs:
#  1. Semaphore pool: the emulation's content-scanned map can hold a *stale*
#     translation for a tracking-semaphore pool page (same VA, previous
#     client's page), so GPU-tracking-semaphore waits never clear. The
#     tracking semaphore object (kprobe `uvmt` gives its kernel pointer) has
#     `semaphore.page->memory->vas[]` with the pool page's CPU and GPU VAs;
#     QEMU turns the CPU VA into a GPA (`gva2gpa`) which we pin.
#  2. CE pushbuffer: the pushes live in a UVM rm_mem with GPU VA
#     `0x120000000` whose CPU VA is vmalloc-mapped guest RAM. Its pages map
#     linearly *descending* in GPA (page i -> base_gpa - i*0x1000). Pinning
#     the range lets the host pump read the pushes directly and execute the
#     PTE memsets/copies instead of only acknowledging their releases.
#
# Sources: guest kprobe on uvm_gpu_tracking_semaphore_update_completed_value
# (`sem=0x...`), QEMU QMP `x/..gx` 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
PUSHBUF_VA_LO = 0x120000000
PUSHBUF_VA_HI = 0x120400000
PUSHBUF_MAX_PAGES = 1024


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.03)
        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))
        m = re.search(r":(.*)$", out, re.S)
        if not m:
            return []
        return [int(x, 16) for x in re.findall(r"0x([0-9a-fA-F]+)", m.group(1))][:n]

    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 is_kern(v):
    return v and (v >> 48) == 0xFFFF


def is_gpu(v):
    return v and GPU_VA_LO <= v < GPU_VA_HI


def guest_trace():
    """(sem kptrs, pushbuffer kptrs) from the trace ring, then clear it.

    Read *before* clearing: uvm_pushbuffer_create fires once per client and
    would be wiped by a clear that races it.
    """
    cmd = ("sudo sh -c 'grep -E \"uvmt|uvpb\" /sys/kernel/tracing/trace | tail -200; "
           "echo > /sys/kernel/tracing/trace'")
    try:
        out = subprocess.run([VMEXEC, "exec", cmd], capture_output=True,
                             text=True, timeout=90).stdout
    except Exception:
        return set(), set()
    sems = {int(m, 16) for m in re.findall(r"sem=0x([0-9a-fA-F]+)", out)}
    pbs = {int(m, 16) for m in re.findall(r"pb=0x([0-9a-fA-F]+)", out)}
    return sems, pbs


def resolve_pb_ptr(q, pb):
    """pushbuffer pointer -> rm_mem -> vas[] -> (gpu_vas, cpu, gpa)."""
    mem = q.read64(pb + 8)
    if not mem or not is_kern(mem[0]):
        return [], None, None
    vas = q.read64(mem[0] + 0x10, 16)
    gpu_vas = [v for v in vas if PUSHBUF_VA_LO <= v < PUSHBUF_VA_HI]
    cpu = next((v for v in vas if is_kern(v)), None)
    if not gpu_vas or not cpu:
        return [], None, None
    gpa = q.gva2gpa(cpu)
    if not gpa:
        return [], None, None
    return gpu_vas, cpu, gpa


def pin(va, pa):
    cmd = "printf '0x%x,0x%x' > %s/sema_map_add" % (va, pa, PARAM)
    for _ in range(40):
        r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
                           capture_output=True, text=True)
        if r.returncode == 0 and "timed out" not in (r.stderr or ""):
            return True
        time.sleep(0.05)  # async queue full: drain and retry
    return False


def resolve_sem_pool(q, sem):
    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)
    cpu = next((v for v in vas if is_kern(v)), None)
    gpu_vas = [v for v in vas if is_gpu(v)]
    if not cpu or not gpu_vas:
        return []
    gpa = q.gva2gpa(cpu)
    if not gpa:
        return []
    return [(gv, gpa) for gv in gpu_vas]


def resolve_pushbuffer(q, sem):
    """Find the pushbuffer rm_mem near the channel and return (gpu_vas, cpu, gpa)."""
    base = sem & ~0xFFF
    words = []
    for i in range(64):
        chunk = q.read64(base - 0x40000 + i * 0x1000, 512)
        if not chunk:
            continue
        words += chunk
    # The manager/pushbuffer live in the same slab region as the channel.
    prefix = sem >> 32
    cands = list(dict.fromkeys(w for w in words
                               if is_kern(w) and (w >> 32) == prefix))
    # Chain 1: value is the channel manager; +8 pushbuffer, +8 rm_mem.
    deadline = time.time() + 45
    for M in cands[:1200]:
        if time.time() > deadline:
            break
        pb = q.read64(M + 8)
        if not pb or not is_kern(pb[0]):
            continue
        mem = q.read64(pb[0] + 8)
        if not mem or not is_kern(mem[0]):
            continue
        vas = q.read64(mem[0] + 0x10, 16)
        gpu_vas = [v for v in vas if PUSHBUF_VA_LO <= v < PUSHBUF_VA_HI]
        if not gpu_vas:
            continue
        cpu = next((v for v in vas if is_kern(v)), None)
        if not cpu:
            continue
        gpa = q.gva2gpa(cpu)
        if gpa:
            return gpu_vas, cpu, gpa
    # Chain 2: value is the pushbuffer itself; +8 = rm_mem.
    # Two QMP reads per candidate, so cap the scan (bounded in time).
    deadline = time.time() + 30
    for V in cands[:400]:
        if time.time() > deadline:
            break
        mem = q.read64(V + 8)
        if not mem or not is_kern(mem[0]):
            continue
        g5 = q.read64(mem[0] + 0x38)
        if not g5 or not (PUSHBUF_VA_LO <= g5[0] < PUSHBUF_VA_HI):
            continue
        g4 = q.read64(mem[0] + 0x30)
        cpu = g4[0] if g4 and is_kern(g4[0]) else None
        if not cpu:
            continue
        gpa = q.gva2gpa(cpu)
        if gpa:
            return [g5[0]], cpu, gpa
    return [], None, None


def pin_pushbuffer(q, gpu_vas, cpu, gpa):
    """Pin the whole pushbuffer with *per-page* translations (the layout is
    not linear across the allocation: page 32 jumped on a real client)."""
    pinned = 0
    for gv in gpu_vas:
        miss = 0
        for i in range(PUSHBUF_MAX_PAGES):
            g = q.gva2gpa(cpu + i * 0x1000)
            if not g:
                miss += 1
                if miss >= 32:
                    print("semfix: pushbuf va=0x%x: 32 unmapped pages at %d, stop (%d pinned)"
                          % (gv, i, pinned), flush=True)
                    return pinned
                continue
            miss = 0
            if not pin(gv + i * 0x1000, g):
                print("semfix: pushbuf va=0x%x: pin failed at page %d" % (gv, i), flush=True)
                return pinned
            pinned += 1
        print("semfix: pushbuf va=0x%x cpu=0x%x gpa=0x%x: pinned %d page(s)"
              % (gv, cpu, gpa, pinned), flush=True)
    return pinned


def main():
    interval = float(sys.argv[1]) if len(sys.argv) > 1 else 3.0
    q = Qmp(QMP_SOCK)
    seen = set()
    seen_pb = set()
    pinned_ranges = set()

    def pin_range_once(gpu_vas, cpu, gpa):
        key = tuple(gpu_vas), cpu
        if key in pinned_ranges:
            return
        pinned_ranges.add(key)
        pin_pushbuffer(q, gpu_vas, cpu, gpa)

    print("semfix: watching (interval %.1fs)" % interval, flush=True)
    while True:
        sems, pbs = guest_trace()
        for pb in pbs - seen_pb:
            seen_pb.add(pb)
            try:
                gpu_vas, cpu, gpa = resolve_pb_ptr(q, pb)
            except Exception as e:
                print("semfix: pb=0x%x resolve failed: %s" % (pb, e), flush=True)
                continue
            if gpu_vas:
                pin_range_once(gpu_vas, cpu, gpa)
            else:
                print("semfix: pb=0x%x no pushbuffer rm_mem (retry)" % pb, flush=True)
        for sem in sems - seen:
            try:
                pairs = resolve_sem_pool(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)
                continue
            seen.add(sem)
            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)
            try:
                gpu_vas, cpu, gpa = resolve_pushbuffer(q, sem)
            except Exception as e:
                print("semfix: pushbuf resolve failed: %s" % e, flush=True)
                gpu_vas = []
            if gpu_vas:
                pin_range_once(gpu_vas, cpu, gpa)
            else:
                print("semfix: pushbuf not found near sem=0x%x (retry)" % sem, flush=True)
        time.sleep(interval)


if __name__ == "__main__":
    main()