import UnityPy
import sys
import os

def extract_assets(folder_path):
    meshes = []
    textures = []
    for root, dirs, files in os.walk(folder_path):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            try:
                env = UnityPy.load(file_path)
                for obj in env.objects:
                    if obj.type.name == "Mesh":
                        meshes.append(obj.get_raw_data())
                    elif obj.type.name == "Texture2D":
                        textures.append(obj.get_raw_data())
            except Exception:
                continue
    return meshes, textures

def inject_assets(folder_path, hp_meshes, hp_textures):
    modified_any = False
    for root, dirs, files in os.walk(folder_path):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            try:
                env = UnityPy.load(file_path)
                file_modified = False
                m_idx = 0
                t_idx = 0
                
                for obj in env.objects:
                    if obj.type.name == "Mesh" and m_idx < len(hp_meshes):
                        obj.set_raw_data(hp_meshes[m_idx])
                        m_idx += 1
                        file_modified = True
                        modified_any = True
                    elif obj.type.name == "Texture2D" and t_idx < len(hp_textures):
                        obj.set_raw_data(hp_textures[t_idx])
                        t_idx += 1
                        file_modified = True
                        modified_any = True
                        
                if file_modified:
                    # SIMPAN FILE MENIMPA ASLINYA (UABE METHOD) AGAR CAB TETAP VALID
                    with open(file_path, "wb") as f:
                        f.write(env.file.save())
                    print(f"[UABE-PYTHON] 🎯 Biner {file_name} ditimpa! CAB aman!")
            except Exception:
                continue
    return modified_any

if __name__ == "__main__":
    if len(sys.argv) < 3:
        sys.exit(1)

    dummy_dir = sys.argv[1]
    hp_dir = sys.argv[2]

    print("[UABE-PYTHON] Mengambil data biner murni dari Raksasa...")
    meshes, textures = extract_assets(hp_dir)

    print(f"[UABE-PYTHON] Menemukan {len(meshes)} Mesh dan {len(textures)} Texture. Menyuntikkan ke Dummy...")
    if inject_assets(dummy_dir, meshes, textures):
        print("[SUCCESS] Operasi UABE In-Place berhasil mutlak!")
    else:
        print("[FAIL] Gagal melakukan modifikasi biner.")
        sys.exit(1)