#!/usr/bin/env python3
import os
import shutil
import subprocess
import sys
import zipfile

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PATCHES_DIR = SCRIPT_DIR
INSTANCE_DIR = os.path.abspath(os.path.join(SCRIPT_DIR, "..", ".."))
RUNTIME_DIR = os.path.join(INSTANCE_DIR, "runtime", "data")

TARGET_JAR_NAME = "defaultoptions-forge-1.20.1-18.0.5.jar"
ORIGINAL_JAR_PATH = os.path.join(RUNTIME_DIR, "automodpack", "host-modpack", "main", "mods", TARGET_JAR_NAME)
OUTPUT_JAR_PATH = os.path.join(PATCHES_DIR, TARGET_JAR_NAME)
CLIENT_SRG_JAR = os.path.join(PATCHES_DIR, "client-1.20.1-srg.jar")

CONTAINER_NAME = "mc_forge_server"
BUILD_TEMP_DIR = os.path.join(PATCHES_DIR, "_build_temp")

def main():
    if not os.path.exists(ORIGINAL_JAR_PATH):
        print(f"[!] Original JAR not found at: {ORIGINAL_JAR_PATH}")
        sys.exit(1)

    if not os.path.exists(CLIENT_SRG_JAR):
        print(f"[!] Required client SRG JAR missing at: {CLIENT_SRG_JAR}")
        print("Please run the scp command from your laptop to copy client-1.20.1-srg.jar!")
        sys.exit(1)

    print(f"[*] Preparing workspace for {TARGET_JAR_NAME}...")
    if os.path.exists(BUILD_TEMP_DIR):
        shutil.rmtree(BUILD_TEMP_DIR)
    os.makedirs(BUILD_TEMP_DIR, exist_ok=True)

    shutil.copy2(ORIGINAL_JAR_PATH, OUTPUT_JAR_PATH)
    shutil.copy2(CLIENT_SRG_JAR, os.path.join(BUILD_TEMP_DIR, "client-1.20.1-srg.jar"))

    extract_dir = os.path.join(BUILD_TEMP_DIR, "jar_extracted")
    with zipfile.ZipFile(ORIGINAL_JAR_PATH, 'r') as z:
        z.extractall(extract_dir)

    src_java = os.path.join(PATCHES_DIR, "KeyMappingDefaultsHandler.java")
    if not os.path.exists(src_java):
        print(f"[!] Source file missing: {src_java}")
        sys.exit(1)

    shutil.copy2(src_java, os.path.join(BUILD_TEMP_DIR, "KeyMappingDefaultsHandler.java"))

    cp_entries = [
        "/build/client-1.20.1-srg.jar",
        "/build/jar_extracted",
        "/data/automodpack/host-modpack/main/mods/balm-forge-1.20.1-7.3.42.jar",
        "/data/libraries/net/minecraftforge/forge/1.20.1-47.4.10/forge-1.20.1-47.4.10-universal.jar",
        "/data/libraries/org/apache/logging/log4j/log4j-api/2.19.0/log4j-api-2.19.0.jar"
    ]
    cp_str = ":".join(cp_entries)

    javac_args_file = os.path.join(BUILD_TEMP_DIR, "javac_args.txt")
    with open(javac_args_file, "w") as f:
        f.write(f"-cp {cp_str}\n")
        f.write("-d /build/out\n")
        f.write("/build/KeyMappingDefaultsHandler.java\n")

    print("[*] Compiling KeyMappingDefaultsHandler.java via JDK container...")
    docker_compile = [
        "docker", "run", "--rm",
        "--volumes-from", CONTAINER_NAME,
        "-v", f"{BUILD_TEMP_DIR}:/build",
        "eclipse-temurin:17-jdk",
        "sh", "-c",
        "mkdir -p /build/out && javac @/build/javac_args.txt"
    ]

    res = subprocess.run(docker_compile, capture_output=True, text=True)
    if res.returncode != 0:
        print("[!] Compilation failed:")
        print(res.stderr)
        print(res.stdout)
        sys.exit(1)

    print("[+] Compilation succeeded!")

    target_class_path = "net/blay09/mods/defaultoptions/keys/KeyMappingDefaultsHandler.class"
    compiled_class_file = os.path.join(BUILD_TEMP_DIR, "out", target_class_path)

    if not os.path.exists(compiled_class_file):
        print(f"[!] Compiled class file not found at: {compiled_class_file}")
        sys.exit(1)

    print(f"[*] Injecting patched class into {OUTPUT_JAR_PATH}...")
    temp_jar_path = OUTPUT_JAR_PATH + ".tmp"
    with zipfile.ZipFile(OUTPUT_JAR_PATH, 'r') as zin, zipfile.ZipFile(temp_jar_path, 'w', compression=zipfile.ZIP_DEFLATED) as zout:
        for item in zin.infolist():
            if item.filename != target_class_path:
                zout.writestr(item, zin.read(item.filename))
        with open(compiled_class_file, 'rb') as f:
            zout.writestr(target_class_path, f.read())

    os.replace(temp_jar_path, OUTPUT_JAR_PATH)

    shutil.rmtree(BUILD_TEMP_DIR, ignore_errors=True)

    print(f"[+] Successfully built patched JAR at: {OUTPUT_JAR_PATH}")

if __name__ == "__main__":
    main()