#!/usr/bin/env python3
"""
patch_hotline.py - Binary patcher for Cymatics Hotline VST3 under Wine.

Addresses two core compatibility issues in JUCE 8 under Wine:
1. Direct2D / DirectComposition initialization failure / freeze:
   Forces HWNDComponentPeer to instantiate the CPU-based Software Renderer
   instead of Direct2D 1.3 / DirectComposition (dcomp.dll).

2. VBlankThread infinite hang / frozen GUI repainting:
   Replaces the failing IDXGIOutput::WaitForVBlank() call with a 16ms Sleep
   timer (~60 FPS), enabling regular UI repaints and graceful thread exit
   when closing/switching windows in REAPER.
"""

import sys
import os
import shutil
import pefile

DEFAULT_PATH = "Cymatics Hotline.vst3/Contents/x86_64-win/Cymatics Hotline.vst3"

# --- Patch Definitions ---
# 1. HWNDComponentPeer Constructor (0x1800b1756, file offset 0xb0b56)
#    Change: call 0x180357d10 (createD2DContext) -> call 0x180357e90 (createSoftwareContext)
OFFSET_CALL_FACTORY = 0xb0b56
ORIG_CALL_FACTORY   = bytes([0xe8, 0xb5, 0x65, 0x2a, 0x00]) # call 0x180357d10
PATCH_CALL_FACTORY  = bytes([0xe8, 0x35, 0x67, 0x2a, 0x00]) # call 0x180357e90

# 2. HWNDComponentPeer Constructor check (0x1800b1727, file offset 0xb0b27)
#    Change: cmp r15d, 1 (is Direct2D active) -> cmp r15d, 0 (is Software active)
OFFSET_CMP_ENGINE   = 0xb0b27
ORIG_CMP_ENGINE     = bytes([0x41, 0x83, 0xff, 0x01]) # cmp r15d, 1
PATCH_CMP_ENGINE    = bytes([0x41, 0x83, 0xff, 0x00]) # cmp r15d, 0

# 3. RenderContext Factory Table Entry 1 (0x1809ede48, file offset 0x9ecc48)
#    Change: Pointer to createD2DContext (0x180357d10) -> createSoftwareContext (0x180357e90)
OFFSET_TABLE_ENTRY1 = 0x9ecc48
ORIG_TABLE_ENTRY1   = (0x180357d10).to_bytes(8, 'little')
PATCH_TABLE_ENTRY1  = (0x180357e90).to_bytes(8, 'little')

# 4. VBlankThread::run loop (0x1800af510, file offset 0xae910, 21 bytes)
#    Original:
#      mov rcx, [rbx + 0x228]
#      mov rax, [rcx]
#      call [rax + 0x50]      ; IDXGIOutput::WaitForVBlank()
#      test eax, eax
#      jne 0x1800af629        ; jumps to Sleep(1) loop without checking exit!
#    Patch:
#      mov ecx, 16            ; Sleep 16 ms (~60 Hz refresh)
#      call [rip + 0x3a1c55]  ; KERNEL32.dll Sleep
#      xor eax, eax           ; eax = 0 (S_OK)
#      nop (8 bytes)
OFFSET_VBLANK_LOOP  = 0xae910
ORIG_VBLANK_LOOP    = bytes.fromhex("488b8b28020000488b01ff505085c00f8504010000")
PATCH_VBLANK_LOOP   = bytes([
    0xb9, 0x10, 0x00, 0x00, 0x00,              # mov ecx, 16
    0xff, 0x15, 0x55, 0x1c, 0x3a, 0x00,        # call qword ptr [rip + 0x3a1c55] (Sleep)
    0x31, 0xc0,                                # xor eax, eax
    0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90, 0x90 # 8x nop
])


def verify(filepath: str):
    """Check current patch status of the binary."""
    with open(filepath, "rb") as f:
        data = f.read()

    status = {}
    status["factory_call"] = "patched" if data[OFFSET_CALL_FACTORY:OFFSET_CALL_FACTORY+5] == PATCH_CALL_FACTORY else (
        "original" if data[OFFSET_CALL_FACTORY:OFFSET_CALL_FACTORY+5] == ORIG_CALL_FACTORY else "unknown"
    )
    status["cmp_engine"] = "patched" if data[OFFSET_CMP_ENGINE:OFFSET_CMP_ENGINE+4] == PATCH_CMP_ENGINE else (
        "original" if data[OFFSET_CMP_ENGINE:OFFSET_CMP_ENGINE+4] == ORIG_CMP_ENGINE else "unknown"
    )
    status["factory_table"] = "patched" if data[OFFSET_TABLE_ENTRY1:OFFSET_TABLE_ENTRY1+8] == PATCH_TABLE_ENTRY1 else (
        "original" if data[OFFSET_TABLE_ENTRY1:OFFSET_TABLE_ENTRY1+8] == ORIG_TABLE_ENTRY1 else "unknown"
    )
    status["vblank_loop"] = "patched" if data[OFFSET_VBLANK_LOOP:OFFSET_VBLANK_LOOP+21] == PATCH_VBLANK_LOOP else (
        "original" if data[OFFSET_VBLANK_LOOP:OFFSET_VBLANK_LOOP+21] == ORIG_VBLANK_LOOP else "unknown"
    )

    print(f"Status for '{filepath}':")
    for k, v in status.items():
        print(f"  {k:15s}: {v}")
    return all(v == "patched" for v in status.values())


def patch(filepath: str, make_backup: bool = True):
    """Apply all 4 patches and update PE checksum."""
    if make_backup:
        bak = filepath + ".bak"
        if not os.path.exists(bak):
            shutil.copyfile(filepath, bak)
            print(f"[*] Backup created: {bak}")

    with open(filepath, "r+b") as f:
        # Patch 1
        f.seek(OFFSET_CALL_FACTORY)
        f.write(PATCH_CALL_FACTORY)
        # Patch 2
        f.seek(OFFSET_CMP_ENGINE)
        f.write(PATCH_CMP_ENGINE)
        # Patch 3
        f.seek(OFFSET_TABLE_ENTRY1)
        f.write(PATCH_TABLE_ENTRY1)
        # Patch 4
        f.seek(OFFSET_VBLANK_LOOP)
        f.write(PATCH_VBLANK_LOOP)

    # Recompute PE checksum
    pe = pefile.PE(filepath)
    pe.OPTIONAL_HEADER.CheckSum = pe.generate_checksum()
    pe.write(filepath)
    pe.close()
    print(f"[+] Successfully patched '{filepath}' and recomputed PE checksum.")


def restore(filepath: str):
    """Restore from backup file."""
    bak = filepath + ".bak"
    if not os.path.exists(bak):
        print(f"[-] No backup file found at '{bak}'")
        return
    shutil.copyfile(bak, filepath)
    print(f"[+] Restored '{filepath}' from '{bak}'.")


if __name__ == "__main__":
    target = sys.argv[2] if len(sys.argv) > 2 else DEFAULT_PATH
    mode = sys.argv[1] if len(sys.argv) > 1 else "--status"

    if mode in ("--verify", "--status"):
        verify(target)
    elif mode == "--patch":
        patch(target)
        verify(target)
    elif mode == "--restore":
        restore(target)
        verify(target)
    else:
        print(f"Usage: {sys.argv[0]} [--status|--patch|--restore] [path/to/binary]")