#!/usr/bin/env python3
"""Create a ZIP archive from a directory, with an optional extra directory.

Usage: create_zip.py <source_dir> <output.zip> <archive_dir_name> [extra_dir]

Files from source_dir and (if given) extra_dir are both stored under
archive_dir_name/ in the ZIP, regardless of the actual directory names on disk.
"""
import sys
import zipfile
from pathlib import Path

if len(sys.argv) < 4:
    print(f'Usage: {sys.argv[0]} <source_dir> <output.zip> <archive_dir_name> [extra_dir]', file=sys.stderr)
    sys.exit(1)

source_dir = Path(sys.argv[1])
zip_path = sys.argv[2]
archive_name = sys.argv[3]
extra_dir = Path(sys.argv[4]) if len(sys.argv) >= 5 else None

with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED, compresslevel=6) as zf:
    for file_path in sorted(source_dir.rglob('*')):
        if file_path.is_file():
            arcname = Path(archive_name) / file_path.relative_to(source_dir)
            zf.write(file_path, arcname)

    if extra_dir and extra_dir.is_dir():
        for file_path in sorted(extra_dir.rglob('*')):
            if file_path.is_file():
                arcname = Path(archive_name) / file_path.relative_to(extra_dir)
                zf.write(file_path, arcname)

print(f'Created: {zip_path}')