"""Temporarily bump python.exe's main thread stack size from 1.94MB to 4MB via PE header patch."""
import struct
import shutil
import os
import sys
from pathlib import Path

PY = Path(os.environ.get("PYTHON_EXE", r"C:\projects\manual_slop_tier2\.venv\Scripts\python.exe"))
BACKUP = PY.with_suffix(".exe.stackbackup")

# PE header structure (simplified for stack size fields)
# DOS header -> e_lfanew at offset 0x3C -> NT headers
# NT headers: signature (4), FileHeader (20), OptionalHeader
# OptionalHeader: Magic (2), MajorLinkerVersion (1), MinorLinkerVersion (1),
#   SizeOfCode (4), SizeOfInitializedData (4), SizeOfUninitializedData (4),
#   AddressOfEntryPoint (4), BaseOfCode (4), BaseOfData (4),
#   ImageBase (4 for 32-bit PE, 8 for 64-bit), SectionAlignment (4),
#   FileAlignment (4), ... then at offset 0x48 (for 64-bit):
#   SizeOfStackReserve (8), SizeOfStackCommit (8)

def get_pe_stack_reserve(python_path: Path) -> int:
    with open(python_path, "rb") as f:
        data = f.read()
    e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
    # Check PE signature
    pe_sig = data[e_lfanew:e_lfanew+4]
    if pe_sig != b"PE\0\0":
        raise ValueError(f"Not a valid PE file at {python_path}")
    # Optional header magic at e_lfanew + 24
    opt_magic = struct.unpack_from("<H", data, e_lfanew + 24)[0]
    if opt_magic == 0x10b:
        # PE32 (32-bit)
        stack_offset = e_lfanew + 24 + 28  # SizeOfStackReserve at offset 28 from OptionalHeader start
        fmt = "<I"
    elif opt_magic == 0x20b:
        # PE32+ (64-bit)
        stack_offset = e_lfanew + 24 + 56  # SizeOfStackReserve at offset 56 from OptionalHeader start
        fmt = "<Q"
    else:
        raise ValueError(f"Unknown PE optional header magic: 0x{opt_magic:x}")
    return struct.unpack_from(fmt, data, stack_offset)[0]

def set_pe_stack_reserve(python_path: Path, new_size: int) -> None:
    with open(python_path, "rb") as f:
        data = bytearray(f.read())
    e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
    opt_magic = struct.unpack_from("<H", data, e_lfanew + 24)[0]
    if opt_magic == 0x20b:
        # PE32+
        stack_offset = e_lfanew + 24 + 56
        fmt = "<Q"
    elif opt_magic == 0x10b:
        stack_offset = e_lfanew + 24 + 28
        fmt = "<I"
    else:
        raise ValueError(f"Unknown PE optional header magic: 0x{opt_magic:x}")
    struct.pack_into(fmt, data, stack_offset, new_size)
    with open(python_path, "wb") as f:
        f.write(data)

if not BACKUP.exists():
    shutil.copy2(PY, BACKUP)
    print(f"Backed up to {BACKUP}")
else:
    print(f"Backup already exists at {BACKUP}")

orig_size = get_pe_stack_reserve(PY)
print(f"Original SizeOfStackReserve: {orig_size} bytes ({orig_size / 1024 / 1024:.2f} MB)")

# Set to 4MB
new_size = 4 * 1024 * 1024
set_pe_stack_reserve(PY, new_size)
print(f"Patched SizeOfStackReserve to: {new_size} bytes ({new_size / 1024 / 1024:.2f} MB)")

# Verify
new_actual = get_pe_stack_reserve(PY)
print(f"Verified SizeOfStackReserve: {new_actual} bytes ({new_actual / 1024 / 1024:.2f} MB)")
