#!/usr/bin/env python3
# fence_assist.py - content-locate the client's ring pages in guest RAM from the
# service's own GPFIFO entry decodes (dmesg) and inject the TRUE VA->PA pairs
# via sema_map_add, so the pump's push reads (and its pushbuffer pairing) have
# a correct anchor.
import mmap, re, struct, subprocess, sys, threading, time

F = "/dev/shm/vgpu-guest-ram"
PARAM = "/sys/module/vgsp/parameters/sema_map_add"

def gpa_of(fo):
    return fo if fo < 0x80000000 else fo - 0x80000000 + 0x100000000

def inject(va, pa, ln=0x1000, flag=1):
    cmd = "echo 0x%x,0x%x,0x%x,%d > %s" % (va & ~0xfff, pa & ~0xfff, ln, flag, PARAM)
    r = subprocess.run(["sudo", "sh", "-c", cmd], capture_output=True)
    print("INJECT va=0x%x pa=0x%x : %s" % (va & ~0xfff, pa & ~0xfff,
          "ok" if r.returncode == 0 else "FAIL"), flush=True)

def search(mm, pats, results):
    CH = 8 << 20
    size = len(mm)
    off = 0
    while off < size:
        d = mm[off:off+CH]
        for (name, pat) in pats.items():
            if name in results:
                continue
            i = d.find(pat)
            if i >= 0:
                results[name] = gpa_of(off + i)
        off += CH - 32

def main():
    dur = float(sys.argv[1]) if len(sys.argv) > 1 else 120
    f = open(F, "rb")
    mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    end = time.time() + dur
    pats = {}       # name -> pattern bytes
    ring_va = {}    # name -> ring VA
    seen = set()    # raw entry lines processed
    done = set()    # injected names
    print("# fence_assist start (%.0fs)" % dur, flush=True)
    while time.time() < end:
        try:
            with open("/tmp/opencode/dmsg_w.txt", "r", errors="ignore") as df:
                tail = df.read().splitlines()[-6000:]
            out = "\n".join(tail)
        except OSError:
            out = ""
            tail = []
        tail = [ln for ln in tail if "sema-chan:" in ln or "sema-push" in ln]
        cur = {}
        for ln in tail:
            m = re.search(r"sema-chan: entry\[(\d+)\] h=0x([0-9a-f]+) eng=\d+ raw=([0-9a-f]+) va=0x([0-9a-f]+)", ln)
            if m:
                idx, h, raw, va = int(m.group(1)), m.group(2), int(m.group(3), 16), int(m.group(4), 16)
                cur.setdefault(h, {})[idx] = raw
                if idx == 0:
                    ring_va[h] = va  # entry[0] va is the push VA (ring_va - 0x20000)
        for h, ents in cur.items():
            key = (h, tuple(sorted(ents.items())))
            if key in seen or len(ents) < 2:
                continue
            seen.add(key)
            name = h
            pat = b"".join(struct.pack("<Q", ents[i]) for i in sorted(ents)[:2])
            pats[name] = pat
            print("# channel h=0x%s entries=%s pattern ready" % (h, sorted(ents)), flush=True)
        if pats:
            results = {}
            threads = []
            # split the map among 4 threads
            LO = 0x100000000; HI = 0x180000000
            size = len(mm)
            step = size // 4
            class Sl(shim := object):
                def __init__(self, a, b):
                    self.a, self.b = a, b
            def worker(a, b):
                CH = 8 << 20
                off = a
                while off < b:
                    if off < 0x80000000 and b <= 0x80000000: pass
                    d = mm[off:off+CH]
                    for (nm, p) in list(pats.items()):
                        if nm in results: continue
                        i = d.find(p)
                        if i >= 0:
                            results[nm] = gpa_of(off + i)
                    off += CH - 32
            for k in range(4):
                a = 0x80000000 + k * (0x80000000 // 4)
                b = 0x80000000 + (k + 1) * (0x80000000 // 4)
                t = threading.Thread(target=worker, args=(a, b))
                t.start(); threads.append(t)
            for t in threads: t.join()
            for nm, gpa in results.items():
                if nm in done: continue
                done.add(nm)
                # ring page content starts with entry[0]; the ring VA is the
                # channel's gpfifo VA = we reconstruct: entry0 va is the push,
                # ring = first "gpfifo=0x..." dmesg line for this h
                m = re.search(r"sema-chan: h=0x%s eng=\d+ cid=\d+ gpfifo=0x([0-9a-f]+)" % nm, out)
                rva = int(m.group(1), 16) if m else None
                if rva:
                    print("RING FOUND h=0x%s gpfifo_va=0x%x -> gpa=0x%x" % (nm, rva, gpa), flush=True)
                    inject(rva, gpa, 0x1000)
        time.sleep(0.6)
    print("# fence_assist done, injected %d" % len(done), flush=True)

main()