#!/usr/bin/env python3

import os
import sys
import subprocess
import argparse
from pathlib import Path
import struct
import re
from shutil import rmtree


if __name__ == "__main__":

    patcher_args = sys.argv[1:sys.argv.index("--")]
    
    parser = argparse.ArgumentParser(description="Patch Elden Ring executable and launch it without EAC.")
    
    parser.add_argument("-r", "--rate", type=int, default=60, help="Modify the frame rate limit (e.g. 30, 120, 165 or whatever).")
    parser.add_argument("--fix-camera", action='store_true', help="Disable camera auto-rotation.")
    parser.add_argument("--all", action='store_true', help="Enable all options except rate adjustment and gamplay changes like `--fix-camera`.")
    parser.add_argument("-u", "--ultrawide", action='store_true', help="Removes black bars when using a resolution with an aspect ratio other than 16:9.")
    parser.add_argument("-v", "--disable-vigniette", action='store_true', help="Disables the vigniette overlay.")
    parser.add_argument("-c", "--disable-ca", action='store_true', help="Disables chromatic abberation.")
    parser.add_argument("-a", "--increase-animation-distance", action='store_true', help="Increase animation distance.")
    parser.add_argument("-s", "--skip-intro", action='store_true', help="Skip intro logos.")
    parser.add_argument("-f", "--remove-60hz-fullscreen", action='store_true', help="Remove 60hz lock in fullscreen.")
    patch = parser.parse_args(patcher_args)

    game_dir = Path(".")
    with open(game_dir / "eldenring.exe", "rb") as f:
        exe_hex = f.read().hex()

    if patch.rate != 60 and patch.rate > 0:
        exe_hex = exe_hex.replace(
            "c743208988883ceb43897318ebca897318",
            "c743208988883ceb43897318ebca897318".replace(
                "8988883c", struct.pack('<f', 1 / patch.rate).hex()
            )
        )

    if patch.fix_camera:
        cf_pattern = '0f 29 a6 .. .. .. .. 41 0f 28 cf'.replace(" ", "")
        cf_addr = re.search(cf_pattern, exe_hex).span()[0]
        cf_offset = 0
        cf_patch = "90 90 90 90 90 90 90".replace(" ", "")
        exe_hex = exe_hex[:cf_addr + cf_offset] + cf_patch + exe_hex[cf_addr + cf_offset + len(cf_patch):]

    if patch.ultrawide or patch.all:
        exe_hex = exe_hex.replace(
            "8b0185c07442448b5904",
            "8b0185c0eb42448b5904"
        )

    if patch.disable_vigniette or patch.all:
        v_pattern = 'f3 0f 10 .. .. f3 0f 59 .. .. .. .. .. e8 .. .. .. .. f3 41 0f .. .. f3 45 0f .. .. 4c 8d .. .. .. .. .. .. 48'.replace(" ", "")
        v_addr = re.search(v_pattern, exe_hex).span()[0]
        v_offset = 46
        v_patch = "f3 0f 5c c0 90".replace(" ", "")  # SUBSS XMM0,XMM0; NOP;  all NOP does work too
        exe_hex = exe_hex[:v_addr + v_offset] + v_patch + exe_hex[v_addr + v_offset + len(v_patch):]

    if patch.disable_ca or patch.all:
        ca_addr = 94 + exe_hex.index("0f114360488d8b800000000f1087a00000000f1141f0488d87b00000000f10080f1109")
        if exe_hex[ca_addr:ca_addr + 8] == "0f114920":
            exe_hex = exe_hex[:ca_addr] + "660fefc9" + exe_hex[ca_addr + 8:]  # PXOR XMM1,XMM1

    if patch.increase_animation_distance or patch.all:
        iad_pattern = "e8 .. .. .. .. 0f 28 .. 0f 28 .. e8 .. .. .. .. f3 0f .. .. 0f 28 .. f3 41 0f 5e".replace(" ", "")
        iad_addr = re.search(iad_pattern, exe_hex).span()[0]
        iad_offset = 46
        iad_patch = "0f 57 c9 66 0f ef c9".replace(" ", "") # DIVSS XMM1,dword ptr [R12 + 0x54] -> XORPS XMM1,XMM1; PXOR XMM1,XMM1
        if exe_hex[iad_addr + iad_offset:iad_addr + iad_offset + len(iad_patch)] == "f3 41 0f 5e 4c 24 54".replace(" ", ""):
            exe_hex = exe_hex[:iad_addr] + iad_patch + exe_hex[iad_addr + len(iad_patch):]

    if patch.skip_intro or patch.all:
        exe_hex = exe_hex.replace(
            "80 bf b8 00 00 00 00 74 53 48".replace(" ", ""),
            "80 bf b8 00 00 00 00 90 90 48".replace(" ", "")
        )
    
    if patch.remove_60hz_fullscreen or patch.all:
        exe_hex = exe_hex.replace(
            "c745ef3c000000",
            "c745ef00000000"
        )
        
    game_dir_patched = Path("er-patcher-tmp")
    if not game_dir_patched.is_dir():
        game_dir_patched.mkdir()

    with open(game_dir_patched / "eldenring.exe", "wb") as f:
        f.write(bytes.fromhex(exe_hex))

    del exe_hex

    # recreate game directory tree in game_dir_patched
    game_dirs = [d for d in game_dir.rglob("*") if d.is_dir()]
    for d in game_dirs:
        if d == game_dir_patched:
            continue
        if not (game_dir_patched / d).is_dir():
             (game_dir_patched / d).mkdir(parents=True)

    # hard link game files to game_dir_patched; symbolic links would be easier
    # to handle but by default windows 10 doesn't allow them
    game_files = [f for f in game_dir.rglob("*") if f.is_file()]
    for f in game_files:
        if f.name in ["eldenring.exe", "start_protected_game.exe", "er-patcher"]:
            continue 
        if not (game_dir_patched / f).is_file():
             f.link_to(game_dir_patched / f)

    # start patched exe directly to avoid EAC
    steam_cmd = sys.argv[1 + sys.argv.index("--"):]
    steam_cmd[-1] = Path(steam_cmd[-1]).parent.absolute() / game_dir_patched / "eldenring.exe"
    subprocess.run(steam_cmd, cwd=steam_cmd[-1].parent.absolute())

    # cleanup
    rmtree(game_dir_patched)
