#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# semfix_uvmt.py - single-QMP tracking-semaphore/resolution watcher.
#
# One QMP client does everything semfix_watch.py and uvmt_poll.py did, because
# QEMU's QMP socket accepts a single client: running the two tools side by side
# silently broke semfix's pinning.
#
# Jobs:
#  1. keep the client's tracking-semaphore pool pages and push-buffer pages
#     pinned with fresh guest truth (semfix_watch.py, imported);
#  2. poll every tracking semaphore seen in the guest trace and print the exact
#     word the client's wait polls: pool GPU VA (gpu_base + idx*4), its guest PA
#     (gva2gpa of the CPU mapping), the payload value and the +0x14/+0x28
#     fields. Every change is appended to a timeline file, so the last wait
#     before a failure and the release that should clear it are both visible;
#  3. mirror the host's decoded releases (dmesg `sema-push ... rel=... pay=...`)
#     into the same timeline, so a wait can be matched to the channel/push that
#     releases its word.
#
# Usage: tools/guest-diag/semfix_uvmt.py [interval-seconds]
# Env:   TIMELINE=/tmp/opencode/uvmt_timeline.log
#        UVMT_HEARTBEAT=30 (seconds between unchanged-state heartbeats)
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 (  # noqa: E402
    Qmp, QMP_SOCK, is_kern, is_gpu, diag_str, guest_trace,
    resolve_sem_pool, pin, pin_bulk, resolve_from_channel,
    resolve_from_pushbuffer, resolve_from_sem, resolve_pushbuffer_wide,
    pin_pushbuffer, PUSHBUF_MAX_PAGES,
)

TIMELINE = os.environ.get("TIMELINE", "/tmp/opencode/uvmt_timeline.log")
HEARTBEAT = float(os.environ.get("UVMT_HEARTBEAT", "30"))
DLOG = os.environ.get("UVMT_DMESG", "1") == "1"
VMEXEC = os.path.normpath(os.path.join(HERE, "..", "..", "scripts", "vgpu_vm.py"))
COUNTER_LO = int(os.environ.get("COUNTER_LO", "0x20400f000"), 16)
COUNTER_HI = int(os.environ.get("COUNTER_HI", "0x204010000"), 16)
SEMSTREAM = os.environ.get("SEM_STREAM_FILE", "/tmp/opencode/sem_stream.txt")
TRACE_EVERY = int(os.environ.get("UVMT_TRACE_EVERY", "5"))
# The push-buffer resolution walks the channel/pool/manager chain with many
# QMP reads and can block the poller for tens of seconds. When the GR window
# stream (pin_gr_stream.py) already supplies the push-buffer pages, skip it so
# the tracking-sem poller stays live during the client's deadline window.
SKIP_PB = os.environ.get("SEMFIX_SKIP_PB", "0") == "1"
# Seconds between counter-page pin attempts (the guest-agent exec blocks for
# its full timeout when the client saturates the guest).
COUNTER_RETRY_S = float(os.environ.get("SEMFIX_COUNTER_RETRY", "10"))


class Stream:
    """Tail the guest scanner's control lines mirrored by pin_gr_stream.py.

    The guest emits "#pid <n>" and "#sem 0x<k>" over vsock; the listener
    appends them here. This replaces per-loop guest-agent trace reads, which
    were too slow to catch short-lived tracking semaphores.
    """

    def __init__(self, path):
        self.path = path
        self.off = 0
        self.sems = []
        self.pids = []
        self.pbs = []
        self.new_pbs = []

    def reset(self):
        """Forget everything learned about the previous client. A new client
        allocates its own channel pools; keeping the old sems/seen-pin marks
        makes the watcher chase stale kernel pointers and pin the previous
        client's (freed) pool pages - the stale-translation trap."""
        self.sems = []
        self.pbs = []
        self.pids = []
        self.new_pbs = []

    def poll(self):
        try:
            with open(self.path) as f:
                f.seek(0, os.SEEK_END)
                size = f.tell()
                if size < self.off:
                    self.off = 0
                f.seek(self.off)
                data = f.read()
                self.off = f.tell()
        except OSError:
            return [], []
        new_sems, new_pids = [], []
        self.new_pbs = []
        for line in data.splitlines():
            if line.startswith("#sem"):
                m = re.search(r"0x([0-9a-fA-F]+)", line)
                if m:
                    v = int(m.group(1), 16)
                    if v not in self.sems:
                        self.sems.append(v)
                        new_sems.append(v)
            elif line.startswith("#pb"):
                m = re.search(r"0x([0-9a-fA-F]+)", line)
                if m:
                    v = int(m.group(1), 16)
                    if v not in self.pbs:
                        self.pbs.append(v)
                        self.new_pbs.append(v)
            elif line.startswith("#pid"):
                m = re.search(r"(\d+)", line)
                if m:
                    v = int(m.group(1))
                    if v not in self.pids:
                        self.pids.append(v)
                        new_pids.append(v)
        return new_sems, new_pids


def log(msg):
    line = "[%s] %s" % (time.strftime("%H:%M:%S"), msg)
    print(line, flush=True)
    try:
        with open(TIMELINE, "a") as f:
            f.write(line + "\n")
    except OSError:
        pass


def uvmt_state(q, sem):
    """Read one tracking semaphore; returns a dict or None.

    Layout (uvm 570.144, verified by probe): sem+0 = tracking semaphore page
    object, sem+4 low u16 = pool index; page+0 = uvm_rm_mem; mem+0x10 = vas[];
    the payload the client's update_completed_value polls is the u32 at
    cpu_va + idx*4 (its GPU alias is gpu_va + idx*4, the VA the releasing push
    writes).
    """
    vals = q.read64(sem, 12, cache=False)
    if len(vals) < 2 or not is_kern(vals[0]):
        return None
    pv = q.read64(vals[0], cache=False)
    if not pv or not is_kern(pv[0]):
        return None
    mem = pv[0]
    vas = q.read64(mem + 0x10, 16, cache=False)
    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
    word_pa = None
    if cpu is not None:
        w = q.read64(cpu + idx * 4, cache=False)
        payload = w[0] if w else None
        word_pa = q.gva2gpa(cpu + idx * 4)
    word_va = (gpu[0] + idx * 4) if gpu else None
    return {
        "page": vals[0],
        "mem": mem,
        "idx": idx,
        "gpu": gpu,
        "cpu": cpu,
        "payload": payload,
        "word_va": word_va,
        "word_pa": word_pa,
        "w14": vals[5] if len(vals) > 5 else 0,
        "w28": vals[10] if len(vals) > 10 else 0,
    }


def state_key(st):
    if not st:
        return None
    return (st["idx"], st["payload"], st["w14"], st["w28"], st["word_pa"],
            tuple(st["gpu"]))


def state_str(sem, st):
    if not st:
        return "sem=0x%x unreadable" % sem
    return ("sem=0x%x idx=%d word_va=%s word_pa=%s payload=%s "
            "+14=0x%x +28=0x%x gpu=%s" %
            (sem, st["idx"],
             "0x%x" % st["word_va"] if st["word_va"] is not None else "?",
             "0x%x" % st["word_pa"] if st["word_pa"] is not None else "?",
             "0x%x" % st["payload"] if st["payload"] is not None else "?",
             st["w14"], st["w28"],
             ",".join("0x%x" % v for v in st["gpu"])))


def dmesg_releases(seen):
    """Print/append new host-decoded releases (sema-push lines).

    Only lines after the last "sema-svc: thread started" marker count: dmesg
    keeps the previous module load's lines and the old client's releases.
    """
    try:
        r = subprocess.run(
            ["sudo", "-n", "sh", "-c",
             "dmesg -T | awk '/sema-svc: thread started/{n=NR} "
             "{l[NR]=$0} END{for(i=n+1;i<=NR;i++) print l[i]}' "
             "| grep 'sema-push' | tail -300"],
            capture_output=True, text=True, timeout=10)
    except Exception:
        return
    for line in r.stdout.splitlines():
        line = line.strip()
        if not line or line in seen:
            continue
        seen.add(line)
        log("push: %s" % line.replace("vgsp: ", ""))


def cdt_pid():
    try:
        r = subprocess.run([VMEXEC, "exec", "pgrep -x cdt | head -1"],
                           capture_output=True, text=True, timeout=30)
    except Exception:
        return None
    pids = re.findall(r"(?m)^(\d+)\s*$", r.stdout)
    return int(pids[-1]) if pids else None


def pin_counter_page(pid, pinned_pids):
    """Pin the client's completion-counter page from its own pagemap.

    The GR channels release the counter page VA range (0x20400fec0..0x20400fff0
    observed); translating it from fresh guest truth before the first releases
    decode is what lets those releases land on the slots the client polls.
    """
    cmd = ("sudo python3 /tmp/guest_pmap_win.py %d 0x%x 0x%x 64"
           % (pid, COUNTER_LO, COUNTER_HI))
    try:
        r = subprocess.run([VMEXEC, "exec", cmd], capture_output=True,
                           text=True, timeout=25)
    except Exception:
        return False
    n = 0
    for m in re.finditer(r"^(0x[0-9a-f]+),(0x[0-9a-f]+),(\d+)$",
                         r.stdout, re.M):
        va, pa, pages = int(m.group(1), 16), int(m.group(2), 16), int(m.group(3))
        if pin_bulk(va, pa, pages * 0x1000):
            n += pages
    if n:
        pinned_pids.add(pid)
        log("counter: pinned %d page(s) at pid %d (0x%x..0x%x)"
            % (n, pid, COUNTER_LO, COUNTER_HI))
    return n > 0


def main():
    interval = float(sys.argv[1]) if len(sys.argv) > 1 else 2.0
    os.environ.setdefault("KEEP_TRACE", "1")
    q = Qmp(QMP_SOCK)
    os.makedirs(os.path.dirname(TIMELINE), exist_ok=True)
    log("semfix_uvmt: watching (interval %.1fs, qmp %s)" % (interval, QMP_SOCK))

    known_sems = set()
    seen_sem_pin = set()
    seen_pb = set()
    seen_ch = set()
    pinned_ranges = {}
    last_state = {}
    last_seen_t = {}
    dseen = set()
    counter_pids = set()
    counter_next_try = {}
    last_pid = None
    stream = Stream(SEMSTREAM)
    loop = 0

    def pin_range_once(gpu_vas, cpu, gpa, size=0):
        key = tuple(gpu_vas), cpu, size
        st = pinned_ranges.get(key)
        if st is None:
            st = pinned_ranges[key] = {"tries": 0, "done": False,
                                       "args": (gpu_vas, cpu, gpa, size)}
        if st["done"]:
            return
        st["tries"] += 1
        pinned = pin_pushbuffer(q, gpu_vas, cpu, gpa, size)
        expected = min(max(size >> 12, 1) if size else PUSHBUF_MAX_PAGES,
                       PUSHBUF_MAX_PAGES)
        if pinned >= expected:
            st["done"] = True
        elif st["tries"] >= 12:
            st["done"] = True
            log("semfix: pushbuf pin incomplete after %d tries: %d/%d pages"
                % (st["tries"], pinned, expected))

    while True:
        loop += 1
        sems = set()
        pbs = set()
        chans = set()

        new_sems, new_pids = stream.poll()
        for s in new_sems:
            log("stream: new uvmt sem 0x%x" % s)
        for p in new_pids:
            log("stream: client pid %d" % p)
            if last_pid is not None and p != last_pid:
                stream.reset()
                known_sems.clear()
                seen_sem_pin.clear()
                last_state.clear()
                last_seen_t.clear()
                pinned_ranges.clear()
                seen_ch.clear()
                seen_pb.clear()
                counter_pids.clear()
                log("semfix: new client pid %d - reset per-client pin state" % p)
            last_pid = p
        for pb in getattr(stream, "new_pbs", []):
            log("stream: new pushbuffer 0x%x" % pb)
            pbs.add(pb)
        sems.update(stream.sems)

        # The guest-agent trace read is slow under client load; use it only
        # as a fallback for channel/pushbuffer discovery and to recover sems
        # the scanner may have missed.
        if loop % TRACE_EVERY == 0:
            try:
                t_sems, t_pbs, t_chans = guest_trace()
            except Exception:
                t_sems, t_pbs, t_chans = set(), set(), set()
            sems.update(t_sems)
            pbs.update(t_pbs)
            chans.update(t_chans)

        # Poll every sem ever seen: the client's final wait may only appear in
        # the trace once, and the trace ring wraps while it spins.
        for sem in sorted(sems - known_sems):
            known_sems.add(sem)

        for sem in sorted(known_sems):
            try:
                st = uvmt_state(q, sem)
            except Exception:
                st = None
            key = state_key(st)
            now = time.time()
            if key != last_state.get(sem):
                last_state[sem] = key
                last_seen_t[sem] = now
                log("uvmt %s" % state_str(sem, st))
            elif now - last_seen_t.get(sem, 0) >= HEARTBEAT:
                last_seen_t[sem] = now
                log("uvmt(hb) %s" % state_str(sem, st))

            # Pin the pool page once per (sem, word) resolution. The payload
            # word is what the releasing push writes; pinning the page makes
            # the host able to apply that release.
            if st and sem not in seen_sem_pin:
                try:
                    pairs = resolve_sem_pool(q, sem)
                except Exception as e:
                    pairs = []
                    log("semfix: sem=0x%x resolve failed: %s" % (sem, e))
                if pairs:
                    seen_sem_pin.add(sem)
                    for gv, gpa in pairs:
                        ok = pin(gv, gpa)
                        log("semfix: sem=0x%x pool va=0x%x -> gpa=0x%x pin=%s"
                            % (sem, gv, gpa, "ok" if ok else "FAILED"))
                    if SKIP_PB:
                        continue
                    d = {}
                    try:
                        res = resolve_from_sem(q, sem, d)
                    except Exception as e:
                        res = None
                        d["exception"] = str(e)[:40]
                    if not res:
                        for k in ("chunks", "words", "cands", "stop_pool",
                                  "no_gpuva", "no_vmalloc", "empty_vas"):
                            d.pop(k, None)
                        try:
                            res = resolve_pushbuffer_wide(q, sem, d)
                        except Exception as e:
                            res = None
                            d["exception"] = str(e)[:40]
                    if res:
                        pin_range_once(*res)
                        log("semfix: sem=0x%x pushbuf -> %s"
                            % (sem, diag_str(d)))
                    else:
                        log("semfix: sem=0x%x pushbuf not found [%s]"
                            % (sem, diag_str(d)))
                else:
                    log("semfix: sem=0x%x pool unresolved (retry)" % sem)

        for ch in sorted(chans - seen_ch):
            d = {}
            try:
                res = resolve_from_channel(q, ch, d)
            except Exception as e:
                res = None
                d["exception"] = str(e)[:40]
            if res:
                seen_ch.add(ch)
                pin_range_once(*res)
                log("semfix: ch=0x%x -> %s" % (ch, diag_str(d)))
            else:
                log("semfix: ch=0x%x pushbuf miss [%s]" % (ch, diag_str(d)))
        for pb in sorted(pbs - seen_pb):
            d = {}
            try:
                res = resolve_from_pushbuffer(q, pb, d)
            except Exception as e:
                res = None
                d["exception"] = str(e)[:40]
            if res:
                seen_pb.add(pb)
                pin_range_once(*res)
                log("semfix: pb=0x%x -> %s" % (pb, diag_str(d)))
            else:
                log("semfix: pb=0x%x pushbuf miss [%s]" % (pb, diag_str(d)))

        if DLOG:
            dmesg_releases(dseen)

        # Retry incomplete push-buffer pins: gva2gpa of the kernel vmalloc
        # alias fails while the sampled vCPU runs in user mode (PTI), so a
        # pass can miss a run of pages that resolves a moment later. Only
        # every 8th loop - a full pass costs thousands of QMP reads.
        if loop % 8 == 0:
            for st in list(pinned_ranges.values()):
                if not st["done"] and st["tries"] > 0:
                    pin_range_once(*st["args"])

        # Pin the completion-counter page of a newly started client, before its
        # first GR pushes decode their releases. The pid comes from the guest
        # scanner; the guest-agent fallback is only a safety net. The exec can
        # block for its whole timeout under client load, so retry on a slow
        # cadence: hammering the agent here starves it and wedges the run.
        if last_pid is None and loop % TRACE_EVERY == 0:
            last_pid = cdt_pid()
            if last_pid:
                log("counter: client pid %d (agent)" % last_pid)
        if last_pid and last_pid not in counter_pids:
            now = time.time()
            if now >= counter_next_try.get(last_pid, 0):
                try:
                    pin_counter_page(last_pid, counter_pids)
                except Exception as e:
                    log("counter: pin failed: %s" % e)
                counter_next_try[last_pid] = now + COUNTER_RETRY_S

        time.sleep(interval)


if __name__ == "__main__":
    main()