#!/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("--with-eac", action='store_true', help="Run game with EAC (Use at own your risk)")
    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:
        r_pattern = "c7 43 20 89 88 88 3c eb 43 89 73 18 eb ca 89 73 18".replace(" ", "")
        if (res := re.search(r_pattern, exe_hex)) is not None:
            r_addr = res.span()[0] + 6
            r_patch = struct.pack('<f', 1 / patch.rate).hex()
            exe_hex = exe_hex[:r_addr] + r_patch + exe_hex[r_addr + len(r_patch):]
        else:
            print("er-patcher: rate pattern scan failed")

    if patch.fix_camera:
        cf_pattern = '0f 29 a6 .. .. .. .. 41 0f 28 cf'.replace(" ", "")
        if (res := re.search(cf_pattern, exe_hex)) is not None:
            cf_addr = res.span()[0]
            cf_patch = "90 90 90 90 90 90 90".replace(" ", "")
            exe_hex = exe_hex[:cf_addr] + cf_patch + exe_hex[cf_addr + len(cf_patch):]
        else:
            print("er-patcher: fix_camera pattern scan failed")

    if patch.ultrawide or patch.all:
        uw_pattern = "74 50 .. 8b .. .. dc 03 00 00 .. 85 .. 74 .. .. 8b .. .. 0f af".replace(" ", "")
        if (res := re.search(uw_pattern, exe_hex)) is not None:
            uw_addr = res.span()[0]
            uw_patch = "eb"
            exe_hex = exe_hex[:uw_addr] + uw_patch + exe_hex[uw_addr + len(uw_patch):]
        else:
            print("er-patcher: ultrawide pattern scan failed")

    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(" ", "")
        if (res := re.search(v_pattern, exe_hex)) is not None:
            v_addr = res.span()[0] + 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_patch + exe_hex[v_addr + len(v_patch):]
        else:
            print("er-patcher: disable_vigniette pattern scan failed")

    if patch.disable_ca or patch.all:
        ca_pattern = "0f 11 43 60 48 8d 8b 80 00 00 00 0f 10 87 a0 00 00 00 0f 11 41 f0 48 8d 87 b0 00 00 00 0f 10 08 0f 11 09".replace(" ", "")
        if (res := re.search(ca_pattern, exe_hex)) is not None:
            ca_addr = res.span()[0] + 94
            ca_orig = "0f 11 49 20".replace(" ", "")
            ca_patch = "66 0f ef c9".replace(" ", "")  # PXOR XMM1,XMM1
            if exe_hex[ca_addr:ca_addr + len(ca_patch)] == ca_orig:
                exe_hex = exe_hex[:ca_addr] + ca_patch + exe_hex[ca_addr + len(ca_patch):]
        else:
            print("er-patcher: disable_ca pattern scan failed")

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

    if patch.skip_intro or patch.all:
        si_pattern = "80 bf b8 00 00 00 00 74 53 48".replace(" ", "")
        if (res := re.search(si_pattern, exe_hex)) is not None:
            si_addr = res.span()[0] + 14
            si_patch = "90 90".replace(" ", "")
            exe_hex = exe_hex[:si_addr] + si_patch + exe_hex[si_addr + len(si_patch):]
        else:
            print("er-patcher: skip_intro pattern scan failed")
    
    if patch.remove_60hz_fullscreen or patch.all:
        fs_pattern = "c7 45 ef 3c 00 00 00".replace(" ", "")
        if (res := re.search(fs_pattern, exe_hex)) is not None:
            fs_addr = res.span()[0] + 6
            fs_patch = "00"
            exe_hex = exe_hex[:fs_addr] + fs_patch + exe_hex[fs_addr + len(fs_patch):]
        else:
            print("er-patcher: remove_60hz_fullscreen pattern scan failed")

    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", "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 / ("start_protected_game.exe" if patch.with_eac else "eldenring.exe")
    subprocess.run(steam_cmd, cwd=steam_cmd[-1].parent.absolute())

    # cleanup
    rmtree(game_dir_patched)