"""Windows only. Unattended game run under cdb (WinDbg console debugger) with logging breakpoints.

Launch method (tested, see NOTES.md): with Steam running, BatmanAC.exe can be started directly
with SteamAppId=200260 in the environment (no relaunch through Steam). SteamStub decrypts .text
at runtime, so we wait until a known code signature is readable, then attach cdb non-intrusively
except for int3 breakpoints placed on already-decrypted code.
"""
import ctypes
import ctypes.wintypes as wt
import os
import re
import subprocess
import time

APPID = 200260


def find_cdb():
    """x86 cdb.exe from the WinDbg Store package (its path changes on every auto-update)."""
    if os.environ.get("ACKIT_CDB"):
        return os.environ["ACKIT_CDB"]
    out = subprocess.run(["powershell", "-NoProfile", "-Command",
                          "(Get-AppxPackage Microsoft.WinDbg).InstallLocation"],
                         capture_output=True, text=True).stdout.strip()
    for base in out.splitlines()[::-1]:
        p = os.path.join(base.strip(), "x86", "cdb.exe")
        if os.path.exists(p):
            return p
    for p in (r"C:\Program Files (x86)\Windows Kits\10\Debuggers\x86\cdb.exe",):
        if os.path.exists(p):
            return p
    raise FileNotFoundError("x86 cdb.exe not found (install WinDbg or set ACKIT_CDB)")

# Runtime addresses (image base 0x400000; ASLR off for this exe). See docs/reverse-engineering/decoder.md.
ADDR = dict(
    get_or_create_setup=0x010B3400,   # __thiscall(cache, AkSrcInfo*, ...) ; info+0x9c hash
    parse_setup=0x010B4D10,           # WwiseVorbisParseSetupHeader(info, channels, reader, pool) -> 0 ok
    parse_setup_ret=0x010B351A,       # in get_or_create_setup, right after the call (eax = result)
    sig_site=0x010B347B,              # mov ebx,[ebp+0x94]  (decryption probe)
)
SIG = bytes.fromhex("8B9D94000000")

k32 = ctypes.WinDLL("kernel32", use_last_error=True)
PROCESS_VM_READ = 0x0010
PROCESS_QUERY_INFORMATION = 0x0400
k32.OpenProcess.restype = wt.HANDLE
k32.ReadProcessMemory.argtypes = [wt.HANDLE, wt.LPCVOID, wt.LPVOID, ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t)]


def read_mem(pid, addr, n):
    h = k32.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
    if not h:
        return None
    try:
        buf = ctypes.create_string_buffer(n)
        got = ctypes.c_size_t()
        if not k32.ReadProcessMemory(h, ctypes.c_void_p(addr), buf, n, ctypes.byref(got)):
            return None
        return buf.raw[:got.value]
    finally:
        k32.CloseHandle(h)


class _MBI(ctypes.Structure):
    _fields_ = [("BaseAddress", ctypes.c_void_p), ("AllocationBase", ctypes.c_void_p),
                ("AllocationProtect", wt.DWORD), ("RegionSize", ctypes.c_size_t), ("State", wt.DWORD),
                ("Protect", wt.DWORD), ("Type", wt.DWORD)]


k32.VirtualQueryEx.argtypes = [wt.HANDLE, ctypes.c_void_p, ctypes.POINTER(_MBI), ctypes.c_size_t]
k32.VirtualQueryEx.restype = ctypes.c_size_t
MEM_COMMIT, PAGE_NOACCESS, PAGE_GUARD = 0x1000, 0x01, 0x100


def search_memory(pid, needles, limit=1 << 32):
    """Which of `needles` (bytes) occur anywhere in the committed memory of the process -> set."""
    h = k32.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
    if not h:
        return set()
    found = set()
    try:
        addr, mbi = 0, _MBI()
        while addr < limit and len(found) < len(needles):
            if not k32.VirtualQueryEx(h, ctypes.c_void_p(addr), ctypes.byref(mbi), ctypes.sizeof(mbi)):
                break
            size = mbi.RegionSize
            if mbi.State == MEM_COMMIT and not mbi.Protect & (PAGE_NOACCESS | PAGE_GUARD):
                buf = ctypes.create_string_buffer(size)
                got = ctypes.c_size_t()
                if k32.ReadProcessMemory(h, ctypes.c_void_p(addr), buf, size, ctypes.byref(got)):
                    data = buf.raw[:got.value]
                    for n in needles:
                        if n not in found and n in data:
                            found.add(n)
            addr = (mbi.BaseAddress or 0) + size
    finally:
        k32.CloseHandle(h)
    return found


def running_game_pids():
    out = subprocess.run(["tasklist", "/FI", "IMAGENAME eq BatmanAC.exe", "/FO", "CSV", "/NH"],
                         capture_output=True, text=True).stdout
    return [int(m.group(1)) for m in re.finditer(r'"BatmanAC\.exe","(\d+)"', out)]


def kill_game():
    for pid in running_game_pids():
        subprocess.run(["taskkill", "/F", "/PID", str(pid)], capture_output=True)


def launch_game(exe, extra_args=()):
    """Start BatmanAC.exe directly (Steam must be running; SteamAppId makes steam_api accept it)."""
    env = dict(os.environ, SteamAppId=str(APPID), SteamGameId=str(APPID))
    return subprocess.Popen([exe] + list(extra_args), cwd=os.path.dirname(exe), env=env)


def wait_decrypted(pid, timeout=60):
    t0 = time.time()
    while time.time() - t0 < timeout:
        if read_mem(pid, ADDR["sig_site"], len(SIG)) == SIG:
            return time.time() - t0
        time.sleep(0.05)
    raise TimeoutError("code never decrypted (signature not found)")


def cdb_script(extra_bps="", dump_info=False):
    """Debugger command file: log codec setup creation + parse result; log 2nd-chance crashes."""
    a = ADDR
    dump = "dd poi(@esp+4) L0x30; " if dump_info else ""
    return "\n".join([
        ".echotime",
        "sxd av",
        'sxd -c2 ".echo ACKIT_CRASH; .echotime; r; kv 30; .ecxr; q" av',
        "sxn ch", "sxn sse", "sxn eh", "sxn clr", "sxn ld", "sxn ud",
        "sxi ibp",
        # stream codec creation: log the AkVorbisInfo fields of the stream being opened
        'bp %#x ".printf \\"ACKIT_SETUP hash=%%08x frames=%%u alloc=%%u ch=%%u rate=%%u\\\\n\\", '
        'poi(poi(@esp+4)+0x9c), poi(poi(@esp+4)+0x78), poi(poi(@esp+4)+0x94), '
        'wo(poi(@esp+4)+0x5a), poi(poi(@esp+4)+0x5c); %sgc"' % (a["get_or_create_setup"], dump),
        # setup parse result (0 = OK)
        'bp %#x ".printf \\"ACKIT_PARSE ret=%%x\\\\n\\", @eax; gc"' % a["parse_setup_ret"],
        extra_bps,
        ".echo ACKIT_ATTACHED",
        "g",
    ])


class Run:
    def __init__(self, exe, log_path, script_extra="", dump_info=False):
        self.exe = exe
        self.log_path = log_path
        self.script_extra = script_extra
        self.dump_info = dump_info
        self.game = None
        self.dbg = None

    def start(self):
        kill_game()
        self.t0 = time.time()
        self.game = launch_game(self.exe)
        pid = self.game.pid
        dt = wait_decrypted(pid)
        script = self.log_path + ".cdbcmds"
        open(script, "w").write(cdb_script(self.script_extra, self.dump_info))
        self.dbg = subprocess.Popen([find_cdb(), "-p", str(pid), "-pd", "-logo", self.log_path,
                                     "-y", "", "-cf", script],
                                    stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                                    stderr=subprocess.DEVNULL)
        return dict(pid=pid, decrypt_s=dt)

    def alive(self):
        return self.game.poll() is None

    def stop(self):
        kill_game()
        if self.dbg:
            try:
                self.dbg.wait(15)
            except subprocess.TimeoutExpired:
                self.dbg.kill()

    def log(self):
        try:
            return open(self.log_path, encoding="utf-8", errors="replace").read()
        except FileNotFoundError:
            return ""