import os, shutil, UnityPy

def copy_and_replace(source_path, destination_path):
    if os.path.exists(destination_path):
        os.remove(destination_path)
    shutil.copy2(source_path, destination_path)
    return

def make_mipmaps(source_folder : str, destination_folder : str):
    for root, dirs, files in os.walk(source_folder):
        for file_name in files:
            changed = False
            file_path = os.path.join(root, file_name)
            env = UnityPy.load(file_path)
            for path,obj in env.container.items():
                if obj.type.name == "Texture2D":
                    data = obj.read()
                    is_sleeve = data.m_Name.startswith("ProtectorIcon")
                    if not is_sleeve: continue
                    data.set_image(data.image, mipmap_count = 1)
                    data.set_image(data.image, mipmap_count = 4)
                    data.save()
                    changed = True
            dest = os.path.join(destination_folder, file_path.removeprefix(source_folder).removeprefix("\\"))
            os.makedirs(os.path.dirname(dest), exist_ok = True)
            with open(dest, "wb") as f:
                f.write(env.file.save(packer="none"))
            if not changed:
                copy_and_replace(file_path, dest)
    return

def main_func():
    source_folder       = input("Enter source 0000/AssetBundle folder\n")
    destination_folder  = input("Enter destination 0000/AssetBundle folder\n")
    make_mipmaps(source_folder, destination_folder)
    print("Sleeves with mipmaps created in destination folder.")
    print("Optional: It is recommended that you compress the files with UABEA.")
    print("Press Enter to exit.")
    return

main_func()