"""Observe only the launched Ori process; trace Wwise thread setup failures.

Uses Windows debugging APIs and Python's standard library. No DLLs are
installed, and no game files are changed. INT3 breakpoints are confined to
the child process and restored before resuming each instruction.
"""
import argparse
import ctypes as C
from ctypes import wintypes as W
import hashlib
import json
import os
from pathlib import Path
import struct
import sys
import time

sys.stdout.reconfigure(encoding="utf-8", errors="replace", line_buffering=True)
K = C.WinDLL("kernel32", use_last_error=True)
HANDLE = C.c_void_p
DWORD = C.c_uint32
U64 = C.c_uint64

class STARTUPINFO(C.Structure):
    _fields_ = [("cb", DWORD), ("reserved", W.LPWSTR), ("desktop", W.LPWSTR),
                ("title", W.LPWSTR), ("x", DWORD), ("y", DWORD),
                ("xs", DWORD), ("ys", DWORD), ("xc", DWORD), ("yc", DWORD),
                ("fill", DWORD), ("flags", DWORD), ("show", W.WORD),
                ("cb2", W.WORD), ("reserved2", HANDLE),
                ("stdin", HANDLE), ("stdout", HANDLE), ("stderr", HANDLE)]

class PROCESSINFO(C.Structure):
    _fields_ = [("process", HANDLE), ("thread", HANDLE),
                ("pid", DWORD), ("tid", DWORD)]

class DEBUGEVENT(C.Structure):
    _fields_ = [("code", DWORD), ("pid", DWORD), ("tid", DWORD),
                ("payload", C.c_ubyte * 164)]

def api(name, restype, *argtypes):
    fn = getattr(K, name)
    fn.restype, fn.argtypes = restype, argtypes
    return fn

CreateProcess = api("CreateProcessW", W.BOOL, W.LPCWSTR, W.LPWSTR, HANDLE,
                    HANDLE, W.BOOL, DWORD, HANDLE, W.LPCWSTR,
                    C.POINTER(STARTUPINFO), C.POINTER(PROCESSINFO))
Read = api("ReadProcessMemory", W.BOOL, HANDLE, HANDLE, HANDLE, C.c_size_t, HANDLE)
Write = api("WriteProcessMemory", W.BOOL, HANDLE, HANDLE, HANDLE, C.c_size_t, HANDLE)
Protect = api("VirtualProtectEx", W.BOOL, HANDLE, HANDLE, C.c_size_t, DWORD, C.POINTER(DWORD))
Flush = api("FlushInstructionCache", W.BOOL, HANDLE, HANDLE, C.c_size_t)
Wait = api("WaitForDebugEvent", W.BOOL, C.POINTER(DEBUGEVENT), DWORD)
Continue = api("ContinueDebugEvent", W.BOOL, DWORD, DWORD, DWORD)
OpenThread = api("OpenThread", HANDLE, DWORD, W.BOOL, DWORD)
GetContext = api("GetThreadContext", W.BOOL, HANDLE, HANDLE)
SetContext = api("SetThreadContext", W.BOOL, HANDLE, HANDLE)
Close = api("CloseHandle", W.BOOL, HANDLE)
Terminate = api("TerminateProcess", W.BOOL, HANDLE, DWORD)
Detach = api("DebugActiveProcessStop", W.BOOL, DWORD)
SetKill = api("DebugSetProcessKillOnExit", W.BOOL, W.BOOL)
GetAffinity = api("GetProcessAffinityMask", W.BOOL, HANDLE, C.POINTER(U64), C.POINTER(U64))

REGS = {"rax":120,"rcx":128,"rdx":136,"rbx":144,"rsp":152,
        "rbp":160,"rsi":168,"rdi":176,"r8":184,"r9":192,
        "r10":200,"r11":208,"r12":216,"r13":224,"r14":232,
        "r15":240,"rip":248}
BREAKS = {0x5B2B9:"before_priority", 0x5B2BF:"after_priority",
          0x5B2D8:"before_affinity", 0x5B2DE:"after_affinity",
          0xAD473:"before_affinity_audio", 0xAD479:"after_affinity_audio",
          0xBB212:"before_affinity_bank", 0xBB218:"after_affinity_bank",
          0xC3364:"before_affinity_monitor", 0xC336A:"after_affinity_monitor"}
EXPECTED_HASH = "8f90a4a3856540cb97d03ca07b4bb10b7a08f7b4ab3c524be2bec1afdd29feca"

def run(args):
    game = Path(args.game).resolve()
    dll = game / "oriwotw_Data/Plugins/AkSoundEngine.dll"
    digest = hashlib.sha256(dll.read_bytes()).hexdigest()
    if digest != EXPECTED_HASH:
        raise RuntimeError("Unexpected AkSoundEngine.dll hash; refusing breakpoints")
    assert C.sizeof(STARTUPINFO)==104 and C.sizeof(PROCESSINFO)==24
    assert C.sizeof(DEBUGEVENT)==176
    print(json.dumps({"game":str(game),"dll_sha256":digest,"check_only":args.check},ensure_ascii=False))
    if args.check:
        return
    out = Path(args.output).resolve()
    out.mkdir(parents=True, exist_ok=True)
    log = (out / "trace.jsonl").open("w",encoding="utf-8")
    def emit(**kw):
        s = json.dumps({"time":time.time(),**kw},ensure_ascii=False)
        print(s)
        log.write(s+"\n")
        log.flush()
    pi, si = PROCESSINFO(), STARTUPINFO()
    si.cb = C.sizeof(si)
    # The command runner uses a private desktop. An explicitly authorized
    # game test must use the user's interactive desktop, or Unity may pause
    # before audio initialization while its window is invisible.
    si.desktop = "winsta0\\default"
    si.flags = 1
    si.show = 1
    cmd = C.create_unicode_buffer('"'+str(game / "oriwotw.exe")+'" -screen-fullscreen 0 -screen-width 1280 -screen-height 720 -logFile "'+str(out / "game.log")+'"')
    if not CreateProcess(str(game / "oriwotw.exe"),cmd,None,None,False,2,None,str(game),C.byref(si),C.byref(pi)):
        raise C.WinError(C.get_last_error())
    emit(event="launched",pid=pi.pid)
    process_mask, system_mask = U64(), U64()
    if not GetAffinity(pi.process,C.byref(process_mask),C.byref(system_mask)):
        raise C.WinError(C.get_last_error())
    emit(event="affinity_limits",process_mask=hex(process_mask.value),system_mask=hex(system_mask.value))
    threads, bps = {}, {}
    pending = None
    alive = True
    start = time.monotonic()
    def read(addr,n):
        buf = C.create_string_buffer(n)
        if not Read(pi.process,addr,buf,n,None):
            return b""
        return buf.raw
    def patch(addr,data):
        old = DWORD()
        if not Protect(pi.process,addr,len(data),0x40,C.byref(old)):
            raise C.WinError(C.get_last_error())
        try:
            if not Write(pi.process,addr,data,len(data),None):
                raise C.WinError(C.get_last_error())
            Flush(pi.process,addr,len(data))
        finally:
            unused = DWORD()
            Protect(pi.process,addr,len(data),old.value,C.byref(unused))
    def context(tid):
        h = OpenThread(0x1A,False,tid)
        if not h:
            raise C.WinError(C.get_last_error())
        raw = C.create_string_buffer(1248)
        addr = (C.addressof(raw)+15)&~15
        C.memset(addr,0,1232)
        C.c_uint32.from_address(addr+48).value=0x10000B
        if not GetContext(h,addr):
            Close(h)
            raise C.WinError(C.get_last_error())
        return h,raw,addr
    try:
        while time.monotonic()-start < args.seconds:
            ev = DEBUGEVENT()
            if not Wait(C.byref(ev),500):
                err=C.get_last_error()
                if err==121:
                    continue
                raise C.WinError(err)
            raw=bytes(ev)
            status=0x10002
            if ev.code==3:
                fh,_,th,_,_,_,teb=struct.unpack_from("<QQQQIIQ",raw,16)
                threads[ev.tid]=teb
                if fh: Close(fh)
            elif ev.code==2:
                th,teb,_=struct.unpack_from("<QQQ",raw,16)
                threads[ev.tid]=teb
                if th: Close(th)
            elif ev.code==6:
                fh,base=struct.unpack_from("<QQ",raw,16)
                # Read the PE export directory name to identify the module.
                hdr=read(base,4096)
                name=""
                if hdr[:2]==b"MZ":
                    pe=struct.unpack_from("<I",hdr,60)[0]
                    export=struct.unpack_from("<I",hdr,pe+24+112)[0]
                    ed=read(base+export,40) if export else b""
                    if ed:
                        nr=struct.unpack_from("<I",ed,12)[0]
                        name=read(base+nr,128).split(b"\0")[0].decode("ascii","replace")
                if fh: Close(fh)
                if name.lower()=="aksoundengine.dll":
                    emit(event="wwise_loaded",base=hex(base))
                    if args.automatic_affinity:
                        addr=base+0x121F0
                        original=bytes.fromhex("48 85 c9 74 03 89 51 04 c3 cc cc cc cc cc cc cc")
                        if read(addr,16)!=original:
                            raise RuntimeError("Affinity setter bytes differ; refusing memory patch")
                        # Preserve the null check; assign the documented zero
                        # mask, meaning inherit Windows scheduling, then return.
                        patch(addr,bytes.fromhex("48 85 c9 74 07 c7 41 04 00 00 00 00 c3 cc cc cc"))
                        emit(event="automatic_affinity_memory_patch",rva="0x121f0")
                    for rva,label in BREAKS.items():
                        addr=base+rva
                        original=read(addr,1)
                        bps[addr]=(original,label)
                        patch(addr,b"\xcc")
            elif ev.code==1:
                code=struct.unpack_from("<I",raw,16)[0]
                addr=struct.unpack_from("<Q",raw,32)[0]
                first=struct.unpack_from("<I",raw,168)[0]
                if code==0x80000003 and addr in bps:
                    original,label=bps.pop(addr)
                    h,buf,ctx=context(ev.tid)
                    try:
                        regs={k:C.c_uint64.from_address(ctx+v).value for k,v in REGS.items()}
                        teb=threads.get(ev.tid,0)
                        er=read(teb+0x68,4) if teb else b""
                        props=read(regs["rdi"],12)
                        emit(event=label,tid=ev.tid,registers={k:hex(v) for k,v in regs.items()},last_error=struct.unpack("<I",er)[0] if er else None,thread_properties=list(struct.unpack("<iII",props)) if props else None)
                        if label.startswith("before_affinity") and args.clamp_affinity:
                            fixed=(regs["rdx"] & process_mask.value) or process_mask.value
                            C.c_uint64.from_address(ctx+REGS["rdx"]).value=fixed
                            emit(event="clamped_affinity",old=hex(regs["rdx"]),new=hex(fixed))
                        patch(addr,original)
                        C.c_uint64.from_address(ctx+248).value=addr
                        if not SetContext(h,ctx): raise C.WinError(C.get_last_error())
                    finally:
                        Close(h)
                elif code==0x80000003:
                    pass
                else:
                    status=0x80010001
                    if code not in (0x406D1388,0xE06D7363):
                        emit(event="exception",code=hex(code),address=hex(addr),first_chance=bool(first),tid=ev.tid)
            elif ev.code==5:
                emit(event="exit",exit_code=struct.unpack_from("<I",raw,16)[0])
                alive=False
            Continue(ev.pid,ev.tid,status)
            if not alive: break
        if alive:
            emit(event="time_limit",seconds=args.seconds)
    finally:
        if alive:
            # Stop only the exact process created above; never attach to or stop
            # any unrelated game session.
            Terminate(pi.process,0)
            for _ in range(100):
                ev=DEBUGEVENT()
                if not Wait(C.byref(ev),100): break
                Continue(ev.pid,ev.tid,0x10002)
                if ev.code==5: break
        Close(pi.thread)
        Close(pi.process)
        log.close()

if __name__=="__main__":
    p=argparse.ArgumentParser()
    p.add_argument("--game",required=True,help="game root directory containing oriwotw.exe")
    p.add_argument("--output",default="ori-evidence/runtime-original")
    p.add_argument("--seconds",type=int,default=40)
    p.add_argument("--check",action="store_true")
    p.add_argument("--clamp-affinity",action="store_true")
    p.add_argument("--automatic-affinity",action="store_true")
    run(p.parse_args())