import bpy
import os
import shutil
import zipfile


def zip_files(paths, output_zip_path):
    with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for path in paths:
            if os.path.isfile(path):
                # If the path is a file, add it to the zip
                zipf.write(path, arcname=os.path.basename(path))
            elif os.path.isdir(path):
                # If the path is a directory, add all its files and subdirectories
                for root, dirs, files in os.walk(path):
                    for file in files:
                        full_path = os.path.join(root, file)
                        zipf.write(full_path, arcname=os.path.join(os.path.basename(path), os.path.relpath(full_path, path)))


# Get File path
project_path = str(bpy.data.filepath)

# Pack Files
bpy.ops.file.pack_all()
bpy.ops.file.make_paths_absolute()

# Temp dir
tmp_dir = os.path.join(os.path.dirname(project_path), 'tmp')
asset_dir = os.path.join(tmp_dir, 'assets')
os.makedirs(tmp_dir, exist_ok=True)

# Find images we could not pack - usually videos
for img in bpy.data.images:
    if not img.packed_file and img.filepath and img.users:
        os.makedirs(asset_dir, exist_ok=True)
        shutil.copy2(img.filepath, os.path.join(asset_dir, os.path.basename(img.filepath)))
        print(f"Copied {os.path.basename(img.filepath)} to tmp directory")
        img.filepath = '//' + os.path.join('assets', os.path.basename(img.filepath))

# Save Output
bpy.ops.wm.save_as_mainfile(filepath=os.path.join(tmp_dir, os.path.basename(project_path)), compress=True, relative_remap=False)

# Save Zip
zip_path = os.path.join(os.path.dirname(project_path), f"{os.path.basename(project_path).split('.')[0]}.zip")
zip_files([os.path.join(tmp_dir, os.path.basename(project_path)), asset_dir], zip_path)
if os.path.exists(zip_path):
    print(f'Saved to: {zip_path}')
else:
    print("Error saving zip file!")

# Cleanup
shutil.rmtree(tmp_dir, ignore_errors=True)
