#!/usr/bin/env python3
# pin_gr_stream.py - pin the CUDA client's GR push-buffer pages from a vsock
# stream produced by guest_pin_stream.py (guest truth via /proc/<pid>/pagemap).
#
# The guest scanner runs in the guest and streams only deltas, so the host
# learns a page the moment it is mapped - no guest-agent round trips on the
# critical path (the agent is only used once to install/start the scanner).
#
# Usage: tools/guest-diag/pin_gr_stream.py
# Env: PINWIN_LO/PINWIN_HI/PINWIN_PA_MAX as in pin_gr_window.py,
#      PINWIN_PORT vsock port (default 45678)
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/vgsp/parameters"
GUEST_SCRIPT = os.path.join(HERE, "guest_pin_stream.py")

PORT = int(os.environ.get("PINWIN_PORT", "45678"))
# 0x200000000 covers the GR GPFIFO rings/UVM channel mappings below the
# old 0x200400000 start; without them the host learns bogus vidmem pairs
# for the rings and never decodes GR pushes.
LO = os.environ.get("PINWIN_LO", "0x200000000")
HI = os.environ.get("PINWIN_HI", "0x206000000")
PAMAX = os.environ.get("PINWIN_PA_MAX", "0x180000000")
SEMSTREAM = os.environ.get("SEM_STREAM_FILE", "/tmp/opencode/sem_stream.txt")


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


def _mirror(line):
    """Append guest control lines (#pid/#sem) to the shared host file so the
    single-QMP watcher can consume them without guest-agent round trips."""
    try:
        with open(SEMSTREAM, "a") as f:
            f.write(line + "\n")
    except OSError:
        pass


def vexec(cmd, timeout=40):
    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 _start_cmd():
    with open(GUEST_SCRIPT) as f:
        b64 = subprocess.run(["base64", "-w0"], stdin=f,
                             capture_output=True, text=True).stdout
    prefixes = os.environ.get("CLIENT_PREFIXES", "cdt,va,victory")
    vma_every = os.environ.get("STREAM_VMA_EVERY", "1")
    return ("echo '%s' | base64 -d > /tmp/guest_pin_stream.py; "
            "if [ -f /tmp/guest_pin_stream.pid ]; then "
            "p=$(cat /tmp/guest_pin_stream.pid 2>/dev/null); "
            "[ -n \"$p\" ] && grep -q guest_pin_stream "
            "/proc/$p/cmdline 2>/dev/null && sudo kill $p 2>/dev/null; "
            "fi; "
            "sudo rm -f /tmp/guest_pin_stream.pid; "
            "sudo setsid nohup env CLIENT_PREFIXES='%s' "
            "STREAM_VMA_EVERY='%s' "
            "python3 /tmp/guest_pin_stream.py %d %d %s %s %s "
            ">/tmp/guest_pin_stream.log 2>&1 & sleep 2; "
            "sudo cat /tmp/guest_pin_stream.pid 2>/dev/null"
            % (b64, prefixes, vma_every, 2, PORT, LO, HI, PAMAX))


def install_and_start():
    # The guest-agent exec can fail transiently (agent busy just after boot,
    # or a stale qga client); retry a few times instead of giving up, because
    # a boot without the scanner streams no pins and the client then hangs.
    for attempt in range(5):
        out = vexec(_start_cmd(), 60)
        if any(ch.isdigit() for ch in out):
            return True
        log("WARN: scanner start attempt %d failed; retrying" % (attempt + 1))
        time.sleep(5)
    return False


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:
            return True
        time.sleep(0.05)
    return False


def main():
    if not install_and_start():
        log("WARN: could not start guest scanner (will retry via vsock wait)")
    srv = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((socket.VMADDR_CID_ANY, PORT))
    srv.listen(4)
    log("listening on vsock port %d (window %s..%s)" % (PORT, LO, HI))

    seen = {}
    pin_total = 0
    buf = b""
    while True:
        try:
            conn, _ = srv.accept()
        except OSError as e:
            log("accept failed: %s" % e)
            time.sleep(1)
            continue
        conn.settimeout(5)
        log("guest connected")
        buf = b""
        runs = []
        while True:
            try:
                chunk = conn.recv(1 << 20)
            except socket.timeout:
                if runs:
                    _flush(runs, seen)
                    runs = []
                continue
            except OSError:
                break
            if not chunk:
                break
            buf += chunk
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                text = line.decode(errors="replace").strip()
                if not text:
                    continue
                if text.startswith("#pid"):
                    m = re.search(r"(\d+)", text)
                    seen = {}
                    log("guest reports cdt pid %s" % (m.group(1) if m else "?"))
                    _mirror(text)
                    continue
                if text.startswith("#sem"):
                    _mirror(text)
                    continue
                if text.startswith("#pb"):
                    _mirror(text)
                    continue
                m = re.match(r"^(0x[0-9a-f]+),(0x[0-9a-f]+)(?:,(\d+))?$", text)
                if m:
                    runs.append((int(m.group(1), 16), int(m.group(2), 16),
                                 int(m.group(3)) if m.group(3) else 1))
            if runs:
                pin_total += _flush(runs, seen)
                runs = []
        log("guest disconnected (pins this session: %d)" % pin_total)
        try:
            conn.close()
        except Exception:
            pass


def _flush(runs, seen):
    """Inject run lines (one batched sudo pass), return number of pages pinned.

    The old implementation spawned one sudo process per run; with ~400 runs
    that cost ~4 s, which landed the GR pins after the client's context-init
    deadline. The batch writes all runs in a single shell, retrying only the
    ones the kernel queue rejected (ENOSPC) individually.
    """
    todo = []
    for va, pa, n in runs:
        if seen.get(va) == pa and n == 1:
            continue
        todo.append((va, pa, n))
    if not todo:
        return 0

    def mark(va, pa, n):
        for i in range(n):
            seen[va + i * 0x1000] = pa + i * 0x1000

    script = []
    for va, pa, n in todo:
        script.append(
            "printf '0x%x,0x%x,0x%x,1' > %s/sema_map_add 2>/dev/null || "
            "echo 'FAIL 0x%x 0x%x 0x%x'"
            % (va, pa, n * 0x1000, PARAM, va, pa, n * 0x1000))
    failed = set()
    CHUNK = 200
    for start in range(0, len(script), CHUNK):
        batch = script[start:start + CHUNK]
        try:
            r = subprocess.run(["sudo", "-n", "sh", "-c", "\n".join(batch)],
                               capture_output=True, text=True, timeout=60)
            for line in r.stdout.splitlines():
                m = re.match(r"FAIL (0x[0-9a-f]+) (0x[0-9a-f]+) (0x[0-9a-f]+)",
                             line.strip())
                if m:
                    failed.add((int(m.group(1), 16), int(m.group(2), 16),
                                int(m.group(3), 16) // 0x1000))
        except Exception as e:
            log("batch inject failed: %s" % e)
            for va, pa, n in todo[start:start + CHUNK]:
                failed.add((va, pa, n))

    pages = 0
    for va, pa, n in todo:
        if (va, pa, n) in failed:
            continue
        mark(va, pa, n)
        pages += n
    for va, pa, n in failed:
        for i in range(n):
            if inject(va + i * 0x1000, pa + i * 0x1000, 0x1000):
                mark(va + i * 0x1000, pa + i * 0x1000, 1)
                pages += 1
    if pages:
        log("pinned %d page(s) in %d run(s)" % (pages, len(todo)))
        for va, pa, n in todo[:3]:
            log("  va=0x%x -> pa=0x%x (%d)" % (va, pa, n))
    return pages


if __name__ == "__main__":
    main()