import sys
import getpass
import os
import tempfile

# Check if user specified a path
if len(sys.argv) > 1:
    path = sys.argv[1]
else:
    path = '/home/' + getpass.getuser() + '/.local/share/Steam/steamapps/common/Trials Rising/datapack/trialsrising.exe'

# Some explanation: these bytes first serve as hex representations of the program instructions we want to change
# But if that was all they needed to be, they could each be 3 hex characters long:
# original_bytes = b'\x48\x89\x54'
# new_and_improved_bytes = b'\xB0\x01\xC3'
# However the issue is that this sequence of original_bytes shows up many many times in the executable,
# so we would end up patching many places that don't need to be patched,
# causing unintended functionality (likely) or else just causing the program to crash (extremely likely)
# instead, we specify *just* enough bytes such that the sequence only occurs in one place:
original_bytes = (
    b'\x48\x89\x54\x24\x10\x55\x56\x57\x48\x81\xEC\xF0\x00\x00\x00\x48\xC7\x44\x24\x78'
)
# translated (from Ghidra (used to disassemble this program)):                                                                                  15a0b00dc(*)
#        14178b6c0 48 89 54        MOV        qword ptr [RSP + 0x10],RDX
#                  24 10
#        14178b6c5 55              PUSH       RBP
#        14178b6c6 56              PUSH       RSI
#        14178b6c7 57              PUSH       RDI
#        14178b6c8 48 81 ec        SUB        RSP,0xf0
#                  f0 00 00 00
#        14178b6cf 48 c7 44        MOV        qword ptr [RSP + 0x78],-0x2
#                  24 78 fe
#                  ff ff ff
new_and_improved_bytes = (
    b'\xB0\x01\xC3\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90\x90'
)
# Translated:
# mov al, 1
# ret
# or in pseudocode, `return True;`
# (The rest of the bytes won't be reached, but they are there to ensure that the length of the new_and_improved_bytes is the same as the original_bytes, so that we don't accidentally shift any other bytes in the executable.)

# Older releases of this patcher wrote the original tail bytes back after `B0 01 C3` instead of
# NOP padding. Detect that form too, so people who already patched with an older version don't
# get a misleading "couldn't detect the patch" error.
old_patched_bytes = b'\xB0\x01\xC3\x24\x10\x55\x56\x57\x48\x81\xEC\xF0\x00\x00\x00\x48\xC7\x44\x24\x78'

# Sanity check: the replacement must be exactly the same length as the original so we never shift
# any other bytes in the executable. This is a hard requirement, so use sys.exit rather than assert
# (assert statements are stripped when Python runs with -O, which would silently remove this guard).
if len(original_bytes) != len(new_and_improved_bytes):
    sys.exit("Internal error: original_bytes and new_and_improved_bytes differ in length. Refusing to patch.")

try:
    with open(path, 'rb') as file:
        data = file.read()
except FileNotFoundError:
    print(f"Error: The file '{path}' was not found.")
    print("Hint: specify the path to the .exe as a command line argument, like this:")
    print(f"python {sys.argv[0]} /path/to/trialsrising.exe")
    sys.exit(1)

# The signature is chosen to be unique to this one function. Enforce that at runtime: patching the
# wrong occurrence would corrupt the executable, so fail loudly instead of silently patching one.
matches = data.count(original_bytes)
if matches == 0:
    # Nothing to patch. Either it's already patched, or we don't recognize this build.
    if data.count(new_and_improved_bytes) or data.count(old_patched_bytes):
        print("The patch had already been applied to your .exe, no further changes were needed.")
        sys.exit(0)
    print("The original bytes were not found in the binary, and we couldn't detect the patch either. \n\n \
            Please try running the patch again with the full path to your .exe, like this:")
    print(f"{sys.argv[0]} /path/to/trialsrising.exe")
    print("")
    print("If that still doesn't work, open an issue to let me know: https://github.com/kylefmohr/Trials-Rising-Steam-Deck-Patch/issues")
    sys.exit(1)
elif matches > 1:
    sys.exit(
        f"Error: expected exactly 1 match for the original bytes, but found {matches}. "
        "Refusing to patch to avoid corrupting the executable. Please open an issue."
    )

# Keep a one-time backup of the original so the patch can be reverted (for example, after a game update).
backup_path = path + '.bak'
if not os.path.exists(backup_path):
    with open(backup_path, 'wb') as backup:
        backup.write(data)

# Write atomically: `open(path, 'wb')` truncates the file immediately, so an interrupted write would
# leave an unusable large executable. Write to a temp file in the same directory, flush it to disk,
# then atomically replace the original.
patched_data = data.replace(original_bytes, new_and_improved_bytes, 1)
temp_path = None
try:
    with tempfile.NamedTemporaryFile('wb', dir=os.path.dirname(os.path.abspath(path)), delete=False) as temp:
        temp_path = temp.name
        temp.write(patched_data)
        temp.flush()
        os.fsync(temp.fileno())
    # NamedTemporaryFile creates the file as 0600, so copy the original file's permissions over.
    os.chmod(temp_path, os.stat(path).st_mode)
    os.replace(temp_path, path)
finally:
    if temp_path is not None and os.path.exists(temp_path):
        os.remove(temp_path)

print(f"Success: Byte sequence found and replaced in {path}.")
print(f"A backup of the original file was saved to {backup_path}.")