# waitscan: dump all threads' registers + wait-handle object types of a hung
# bridge. Goal: find what Nexus attached() and its worker threads wait on.
import ctypes, ctypes.wintypes as w, struct, subprocess, sys, os, time

sys.stdout.reconfigure(encoding='utf-8', errors='replace')

SHM_NAME = "SonicForge_DAW_IPC_VERIFY"
SIZE = 32768
BRIDGE = r"C:\Users\locpham\SonicForgeStudio\native_bridge\build\Release\daw_vst_bridge.exe"
LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "waitscan.log")

k32 = ctypes.windll.kernel32
k32.CreateFileMappingA.argtypes = [ctypes.c_void_p, ctypes.c_void_p, w.DWORD, w.DWORD, w.DWORD, ctypes.c_char_p]
k32.CreateFileMappingA.restype = ctypes.c_void_p
k32.MapViewOfFile.argtypes = [ctypes.c_void_p, w.DWORD, w.DWORD, w.DWORD, ctypes.c_size_t]
k32.MapViewOfFile.restype = ctypes.c_void_p
hMap = k32.CreateFileMappingA(ctypes.c_void_p(-1).value, None, 0x04, 0, SIZE, SHM_NAME.encode())
k32.MapViewOfFile(hMap, 0xF001F, 0, 0, 0)

def wait_log(pattern, timeout):
    end = time.time() + timeout
    while time.time() < end:
        try:
            with open(LOG, 'r', encoding='utf-8', errors='replace') as f:
                if pattern in f.read():
                    return True
        except FileNotFoundError:
            pass
        time.sleep(0.2)
    return False

extra = dict(os.environ)
for k in list(extra):
    if k.startswith("SF_"):
        del extra[k]
extra.update({"SF_SHM_NAME": SHM_NAME, "SF_SAMPLE_RATE": "48000",
              "SF_BLOCK_SIZE": "256", "SF_ONCE": "1", "SF_AUTOGUI": "0",
              "SF_ONCE_PROBEWIN": "1"})
proc = subprocess.Popen([BRIDGE, "--shm", SHM_NAME], env=extra,
                        stdout=open(LOG, 'wb'), stderr=subprocess.STDOUT)
pid = proc.pid
print("pid", pid)
if not wait_log("isPlatformTypeSupported=0", 120):
    print("TIMEOUT"); sys.exit(1)
time.sleep(3)

PROCESS_QUERY_INFORMATION = 0x0400
PROCESS_VM_READ = 0x0010
PROCESS_DUP_HANDLE = 0x0040
TH32CS_SNAPTHREAD = 0x4
THREAD_GET_CONTEXT = 0x0008
THREAD_QUERY_INFORMATION = 0x0040
THREAD_SUSPEND_RESUME = 0x0002

k32.OpenProcess.argtypes = [w.DWORD, w.BOOL, w.DWORD]
k32.OpenProcess.restype = ctypes.c_void_p
hProc = k32.OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_DUP_HANDLE, False, pid)
if not hProc:
    print("OpenProcess failed", ctypes.get_last_error()); sys.exit(1)

# module list for name resolution
psapi = ctypes.windll.psapi
class MODULEINFO(ctypes.Structure):
    _fields_ = [("lpBaseOfDll", ctypes.c_void_p), ("SizeOfImage", w.DWORD),
                ("pad", w.DWORD), ("EntryPoint", ctypes.c_void_p)]
psapi.EnumProcessModulesEx.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), w.DWORD, ctypes.POINTER(w.DWORD), w.DWORD]
psapi.GetModuleInformation.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.POINTER(MODULEINFO), w.DWORD]
psapi.GetModuleBaseNameA.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_char_p, w.DWORD]
nmods = w.DWORD(0)
psapi.EnumProcessModulesEx(hProc, None, 0, ctypes.byref(nmods), 3)
arr = (ctypes.c_void_p * (nmods.value // 8))()
psapi.EnumProcessModulesEx(hProc, arr, nmods.value, ctypes.byref(nmods), 3)
mods = []
for m in arr:
    if not m:
        continue
    mi = MODULEINFO()
    psapi.GetModuleInformation(hProc, m, ctypes.byref(mi), ctypes.sizeof(MODULEINFO))
    name = ctypes.create_string_buffer(260)
    psapi.GetModuleBaseNameA(hProc, m, name, 260)
    mods.append((mi.lpBaseOfDll, mi.SizeOfImage, name.value.decode('latin1')))
def resolve(addr):
    for base, size, name in mods:
        if base <= addr < base + size:
            return f"{name}+0x{addr-base:x}"
    return f"0x{addr:x}"

class THREADENTRY32(ctypes.Structure):
    _fields_ = [("dwSize", w.DWORD), ("cntUsage", w.DWORD), ("th32ThreadID", w.DWORD),
                ("th32OwnerProcessID", w.DWORD), ("tpBasePri", w.LONG),
                ("tpDeltaPri", w.LONG), ("dwFlags", w.DWORD)]
te = THREADENTRY32()
te.dwSize = ctypes.sizeof(THREADENTRY32)
snap = k32.CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)
ok = k32.Thread32First(snap, ctypes.byref(te))
threads = []
while ok:
    if te.th32OwnerProcessID == pid:
        threads.append(te.th32ThreadID)
    ok = k32.Thread32Next(snap, ctypes.byref(te))
k32.CloseHandle(snap)
print(f"threads: {threads}")

class CONTEXT(ctypes.Structure):
    # full x64 CONTEXT is 0x4d0 bytes; we only parse the GPR area (0x00-0xf8).
    _fields_ = [("raw", ctypes.c_ubyte * 0x4d0)]

CONTEXT_AMD64 = 0x00100000
CONTEXT_CONTROL = CONTEXT_AMD64 | 0x1
CONTEXT_INTEGER = CONTEXT_AMD64 | 0x2
CONTEXT_FULL = CONTEXT_CONTROL | CONTEXT_INTEGER
def ctx_regs(ctx):
    b = bytes(ctx.raw)
    u = lambda o: struct.unpack_from('<Q', b, o)[0]
    return dict(rip=u(0xF8), rsp=u(0x98), rcx=u(0x80), rdx=u(0x88), r8=u(0xB8), r9=u(0xC0),
                rbx=u(0x90), rax=u(0x78))
k32.GetThreadContext.argtypes = [ctypes.c_void_p, ctypes.POINTER(CONTEXT)]
k32.GetThreadContext.restype = w.BOOL
k32.SuspendThread.argtypes = [ctypes.c_void_p]
k32.SuspendThread.restype = w.DWORD
k32.ResumeThread.argtypes = [ctypes.c_void_p]
k32.ResumeThread.restype = w.DWORD
k32.OpenThread.argtypes = [w.DWORD, w.BOOL, w.DWORD]
k32.OpenThread.restype = ctypes.c_void_p
k32.CloseHandle.argtypes = [ctypes.c_void_p]
k32.CloseHandle.restype = w.BOOL
k32.DuplicateHandle.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), w.DWORD, w.BOOL, w.DWORD]
k32.DuplicateHandle.restype = w.BOOL
k32.GetCurrentProcess.restype = ctypes.c_void_p

ntdll = ctypes.windll.ntdll
class UNICODE_STRING(ctypes.Structure):
    _fields_ = [("Length", w.USHORT), ("MaximumLength", w.USHORT), ("Buffer", ctypes.c_void_p)]
class OBJECT_NAME_INFORMATION(ctypes.Structure):
    _fields_ = [("Name", UNICODE_STRING)]
class OBJECT_TYPE_INFORMATION(ctypes.Structure):
    _fields_ = [("Name", UNICODE_STRING), ("TotalNumberOfObjects", w.ULONG), ("TotalNumberOfHandles", w.ULONG),
                ("TotalPagedPoolUsage", w.ULONG), ("TotalNonPagedPoolUsage", w.ULONG),
                ("TotalNamePoolUsage", w.ULONG), ("TotalHandleTableUsage", w.ULONG),
                ("HighWaterNumberOfObjects", w.ULONG), ("HighWaterNumberOfHandles", w.ULONG),
                ("HighWaterPagedPoolUsage", w.ULONG), ("HighWaterNonPagedPoolUsage", w.ULONG),
                ("HighWaterNamePoolUsage", w.ULONG), ("HighWaterHandleTableUsage", w.ULONG),
                ("InvalidAttributes", w.ULONG), ("GenericMapping", w.BYTE * 16), ("ValidAccess", w.ULONG),
                ("SecurityRequired", w.BYTE), ("MaintainHandleCount", w.BYTE),
                ("MaintainTypeList", w.BYTE), ("Reserved", w.BYTE * 9)]

def obj_info(h):
    # duplicate to get a handle valid in OUR process for querying
    dup = ctypes.c_void_p()
    if not k32.DuplicateHandle(hProc, ctypes.c_void_p(h), k32.GetCurrentProcess(),
                               ctypes.byref(dup), 0, False, 0x2):  # DUPLICATE_SAME_ACCESS
        return "dup-fail"
    # type
    buf = ctypes.create_string_buffer(512)
    sz = w.ULONG(0)
    r = ntdll.NtQueryObject(dup, 2, buf, 512, ctypes.byref(sz))  # ObjectTypeInformation=2
    typ = "?"
    if r == 0:
        ti = OBJECT_TYPE_INFORMATION.from_buffer(buf)
        if ti.Name.Buffer:
            typ = ctypes.wstring_at(ti.Name.Buffer)
    # name
    buf2 = ctypes.create_string_buffer(1024)
    sz2 = w.ULONG(0)
    name = ""
    r = ntdll.NtQueryObject(dup, 1, buf2, 1024, ctypes.byref(sz2))  # ObjectNameInformation=1
    if r == 0:
        oi = OBJECT_NAME_INFORMATION.from_buffer(buf2)
        if oi.Name.Buffer:
            try:
                name = ctypes.wstring_at(oi.Name.Buffer)
            except Exception:
                name = "(name-err)"
    k32.CloseHandle(dup)
    return f"{typ}|{name}"

for tid in threads:
    ht = k32.OpenThread(THREAD_GET_CONTEXT | THREAD_QUERY_INFORMATION | THREAD_SUSPEND_RESUME, False, tid)
    if not ht:
        print(f"tid {tid}: open fail {ctypes.get_last_error()}"); continue
    k32.SuspendThread(ht)
    ctx = CONTEXT()
    struct.pack_into('<I', ctx.raw, 0x30, CONTEXT_FULL)
    if k32.GetThreadContext(ht, ctypes.byref(ctx)):
        r = ctx_regs(ctx)
        rip = r['rip']; rcx = r['rcx']
        line = f"tid {tid}: rip={resolve(rip)} rcx=0x{rcx:x} rdx=0x{r['rdx']:x} r8=0x{r['r8']:x} r9=0x{r['r9']:x} rsp=0x{r['rsp']:x}"
        if rip and rcx:
            try:
                info = obj_info(rcx)
                line += f" | wait_obj: {info}"
            except Exception as e:
                line += f" | obj_err {e}"
        print(line)
    else:
        print(f"tid {tid}: GetThreadContext fail {ctypes.get_last_error()}")
    k32.ResumeThread(ht)
    k32.CloseHandle(ht)