#!/usr/bin/env python3
"""
Patch script for Interstate '76 (GOG Repack, 2009)
Fixes: EXCEPTION_PRIV_INSTRUCTION crash at RVA 0x99997 (cli instruction)

Usage:
    python3 patch_i76.py          # apply patch
    python3 patch_i76.py --revert # restore original
    python3 patch_i76.py --check  # verify current state

Patch:
    File offset 0x98D8B: 73 0A (jae +10) → 90 90 (nop; nop)
    Effect: Always take the "is Pentium" code path, skip the dead cli routine.

Original MD5:  9a232dcc2c164648cff20c414c1f9698
Patched MD5:  712ead811bea234a5f5774d13da17ede
"""

import hashlib
import shutil
import sys
import os

EXE = "i76.exe"
BACKUP = "i76.exe.original"
PATCH_OFFSET = 0x98D8B
PATCH_BYTES = b"\x90\x90"      # nop; nop
ORIG_BYTES = b"\x73\x0a"       # jae +0x0a
ORIG_MD5 = "9a232dcc2c164648cff20c414c1f9698"
PATCHED_MD5 = "712ead811bea234a5f5774d13da17ede"


def md5(path):
    with open(path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()


def current_bytes():
    with open(EXE, "rb") as f:
        f.seek(PATCH_OFFSET)
        return f.read(2)


def apply():
    if not os.path.exists(EXE):
        print(f"ERROR: {EXE} not found.")
        sys.exit(1)

    cur = current_bytes()

    if cur == PATCH_BYTES:
        print(f"Already patched (bytes: {cur.hex()}). Nothing to do.")
        return

    if cur != ORIG_BYTES:
        print(f"WARNING: Expected bytes {ORIG_BYTES.hex()} at offset 0x{PATCH_OFFSET:X}, "
              f"found {cur.hex()}. Patch may not apply correctly.")
        resp = input("Continue anyway? [y/N] ")
        if resp.lower() != "y":
            sys.exit(0)

    # Backup if not already backed up
    if not os.path.exists(BACKUP):
        shutil.copyfile(EXE, BACKUP)
        print(f"Backup saved: {BACKUP}")
        print(f"  MD5: {md5(BACKUP)}")
    else:
        print(f"Backup already exists: {BACKUP}")

    # Patch
    with open(EXE, "r+b") as f:
        f.seek(PATCH_OFFSET)
        f.write(PATCH_BYTES)

    print(f"Patch applied: offset 0x{PATCH_OFFSET:X} → {PATCH_BYTES.hex()}")
    print(f"  New MD5: {md5(EXE)}")


def revert():
    if not os.path.exists(BACKUP):
        print(f"ERROR: No backup found ({BACKUP}). Cannot revert.")
        sys.exit(1)

    shutil.copyfile(BACKUP, EXE)
    print(f"Reverted {EXE} from {BACKUP}")
    print(f"  MD5: {md5(EXE)}")


def check():
    if not os.path.exists(EXE):
        print(f"ERROR: {EXE} not found.")
        sys.exit(1)

    cur = current_bytes()
    exe_md5 = md5(EXE)

    if cur == PATCH_BYTES:
        status = "PATCHED ✅"
    elif cur == ORIG_BYTES:
        status = "ORIGINAL"
    else:
        status = f"UNKNOWN (bytes: {cur.hex()})"

    print(f"{EXE}: {status}")
    print(f"  Offset 0x{PATCH_OFFSET:X}: {cur.hex()} (expected: {ORIG_BYTES.hex()} or {PATCH_BYTES.hex()})")
    print(f"  MD5: {exe_md5}")

    if exe_md5 == PATCHED_MD5:
        print("  State: ✅ Correctly patched")
    elif exe_md5 == ORIG_MD5:
        print("  State: ⬜ Original (not patched)")
    else:
        print("  State: ⚠️  Unknown state")


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] in ("--revert", "-r"):
        revert()
    elif len(sys.argv) > 1 and sys.argv[1] in ("--check", "-c"):
        check()
    else:
        apply()