#!/usr/bin/env python3
# pin_gr_window.py - pin the CUDA client's GR push-buffer pages from guest
# truth, per boot (host-side dev tool, no guest changes beyond the helper).
#
# Why: the compute (GR) channels' GPFIFO entries point at pushes that live in
# the client's GPU VA window (observed one allocation per channel starting at
# 0x200400000 + stride). The pushes are sysmem pages mapped into the client
# process, so the guest's /proc/<pid>/pagemap gives their GPA. Injecting those
# translations (sema_map_add) is what lets the pump read the GR pushes and
# apply the tracking-semaphore releases they carry - the wall documented in
# PART 61 (pushes_nomap / deferred releases).
#
# The tool re-scans on an interval so pages mapped after the first pass are
# caught, and re-injects when a new client pid appears (pins are per-process
# pages; the same VA usually has a different PA on the next client/boot).
#
# Usage: tools/guest-diag/pin_gr_window.py [interval-seconds]
# Env:
#   PINWIN_LO/PINWIN_HI  VA window to scan (default 0x200400000..0x206000000)
#   PINWIN_MAX           max present pages per pass (default 20000)
#   PINWIN_PA_MAX        ignore PAs at/above this (skip device BARs)
#                        (default 0x180000000)
#   PINWIN_VERIFY        1 = probe the first pins on the first batch (default)

import os
import re
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/vgsp/parameters"
HELPER = os.path.join(HERE, "guest_pmap_win.py")

LO = int(os.environ.get("PINWIN_LO", "0x200400000"), 16)
HI = int(os.environ.get("PINWIN_HI", "0x206000000"), 16)
MAXP = int(os.environ.get("PINWIN_MAX", "20000"))
PAMAX = int(os.environ.get("PINWIN_PA_MAX", "0x180000000"), 16)
VERIFY = os.environ.get("PINWIN_VERIFY", "1") == "1"


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


def vexec(cmd, timeout=60):
    try:
        r = subprocess.run([VMEXEC, "exec", "--timeout", str(timeout - 5), cmd],
                           capture_output=True, text=True, timeout=timeout)
        return r.stdout
    except Exception:
        return ""


def push_helper():
    with open(HELPER) as f:
        b64 = subprocess.run(["base64", "-w0"], stdin=f,
                             capture_output=True, text=True).stdout
    for _ in range(5):
        vexec("echo '%s' | base64 -d > /tmp/guest_pmap_win.py" % b64, 40)
        out = vexec("test -s /tmp/guest_pmap_win.py && echo PUSHED", 20)
        if "PUSHED" in out:
            return True
        time.sleep(1)
    return False


def cdt_pid():
    out = vexec("pgrep -x cdt", 25)
    m = re.search(r"^\s*(\d+)\s*$", out, re.M)
    return int(m.group(1)) if m else None


def scan(pid):
    out = vexec("sudo python3 /tmp/guest_pmap_win.py %d 0x%x 0x%x %d 0x%x"
                % (pid, LO, HI, MAXP, PAMAX), 90)
    pairs = []
    for line in out.splitlines():
        m = re.match(r"^(0x[0-9a-f]+),(0x[0-9a-f]+)(?:,(\d+))?$", line.strip())
        if not m:
            continue
        va = int(m.group(1), 16)
        pa = int(m.group(2), 16)
        n = int(m.group(3)) if m.group(3) else 1
        for i in range(n):
            pairs.append((va + i * 0x1000, pa + i * 0x1000))
    return pairs


def inject(va, pa, ln):
    cmd = "printf '0x%x,0x%x,0x%x,1' > %s/sema_map_add" % (va, pa, ln, PARAM)
    for _ in range(60):
        r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
                           capture_output=True, text=True)
        if r.returncode == 0 and "denied" not in (r.stderr or ""):
            return True
        time.sleep(0.05)
    return False


def runs_of(pairs):
    """Group (va,pa) pairs into contiguous VA+PA runs."""
    pairs = sorted(set(pairs))
    runs = []
    cur_va, cur_pa, cur_n = pairs[0][0], pairs[0][1], 1
    for va, pa in pairs[1:]:
        if va == cur_va + cur_n * 0x1000 and pa == cur_pa + cur_n * 0x1000:
            cur_n += 1
            continue
        runs.append((cur_va, cur_pa, cur_n * 0x1000))
        cur_va, cur_pa, cur_n = va, pa, 1
    runs.append((cur_va, cur_pa, cur_n * 0x1000))
    return runs


def probe(va):
    subprocess.run(["sudo", "-n", "sh", "-c",
                    "printf '0x%x,2' > %s/sema_probe" % (va, PARAM)],
                   capture_output=True, text=True)
    time.sleep(0.4)
    r = subprocess.run(["sudo", "-n", "dmesg"], capture_output=True, text=True)
    tail = r.stdout.splitlines()[-14:]
    vals = []
    for line in tail:
        m = re.search(r"sema-probe:.*\[(\d+)\] 0x([0-9a-f]+)", line)
        if m:
            vals.append((int(m.group(1)), int(m.group(2), 16)))
    return vals


def main():
    interval = float(sys.argv[1]) if len(sys.argv) > 1 else 1.0
    if not push_helper():
        log("FATAL: could not push guest helper")
        return 1
    log("watching window 0x%x..0x%x (pamax 0x%x, max %d, interval %.1fs)"
        % (LO, HI, PAMAX, MAXP, interval))
    seen = {}
    pid = None
    verified = False
    while True:
        p = cdt_pid()
        if not p:
            pid = None
            time.sleep(interval)
            continue
        if p != pid:
            pid = p
            seen = {}
            log("new cdt pid %d" % pid)
        try:
            pairs = scan(pid)
        except Exception as e:
            log("scan failed: %s" % e)
            time.sleep(interval)
            continue
        new = [(va, pa) for va, pa in pairs if seen.get(va) != pa]
        if new:
            runs = runs_of(new)
            npages = 0
            nfail = 0
            for va, pa, ln in runs:
                if inject(va, pa, ln):
                    npages += ln // 0x1000
                else:
                    nfail += 1
                    for i in range(ln // 0x1000):
                        if inject(va + i * 0x1000, pa + i * 0x1000, 0x1000):
                            npages += 1
                        else:
                            nfail += 1
            for va, pa in new:
                seen[va] = pa
            log("pid=%d present=%d new=%d pinned=%dp/%dr fail=%d"
                % (pid, len(pairs), len(new), npages, len(runs), nfail))
            for va, pa in new[:4]:
                log("  va=0x%x -> pa=0x%x" % (va, pa))
            if VERIFY and not verified and new:
                verified = True
                for va, _pa in new[:2]:
                    vals = probe(va)
                    log("  probe va=0x%x: %s" % (va, " ".join(
                        "0x%x" % v for _, v in vals) or "(no reply)"))
        time.sleep(interval)


if __name__ == "__main__":
    sys.exit(main())