#!/usr/bin/env python3
import argparse
import zipfile
from pathlib import Path


SKIP_DIRS = {
    ".git",
    ".hf-push.git",
    ".hf-lfs-cache",
    "portable_runtime",
}


def iter_safetensors(root: Path):
    for path in root.rglob("*.safetensors"):
        rel_parts = path.relative_to(root).parts
        if any(part in SKIP_DIRS for part in rel_parts):
            continue
        yield path


def zip_one(root: Path, path: Path, keep: bool) -> Path:
    rel = path.relative_to(root)
    out_path = path.with_suffix(path.suffix + ".zip")
    tmp_path = path.with_suffix(path.suffix + ".zip.tmp")

    if out_path.exists():
        raise FileExistsError(f"Refusing to overwrite existing archive: {out_path}")

    with zipfile.ZipFile(tmp_path, "w", compression=zipfile.ZIP_STORED, allowZip64=True) as archive:
        archive.write(path, arcname=str(rel))
    tmp_path.replace(out_path)

    if not keep:
        path.unlink()

    return out_path


def main():
    parser = argparse.ArgumentParser(
        description="Store every project .safetensors file as .safetensors.zip."
    )
    parser.add_argument("--root", type=Path, default=Path.cwd())
    parser.add_argument("--keep", action="store_true", help="Keep original .safetensors files.")
    args = parser.parse_args()

    root = args.root.resolve()
    files = list(iter_safetensors(root))
    print(f"Found {len(files)} .safetensors files under {root}")

    for idx, path in enumerate(files, start=1):
        rel = path.relative_to(root)
        size_gib = path.stat().st_size / (1024 ** 3)
        print(f"[{idx}/{len(files)}] zipping {rel} ({size_gib:.2f} GiB)", flush=True)
        out_path = zip_one(root, path, keep=args.keep)
        print(f"    wrote {out_path.relative_to(root)}", flush=True)


if __name__ == "__main__":
    main()