from __future__ import annotations

import argparse
from pathlib import Path
import zipfile

EXCLUDE_DIRS = {
    ".git",
    ".cache",
    ".venv",
    "venv",
    "artifacts",
    "checkpoints",
    "runs",
    "wandb",
    "raw_data",
    "processed_data",
    "__pycache__",
    ".ipynb_checkpoints",
}
EXCLUDE_SUFFIXES = {
    ".pt", ".pth", ".ckpt", ".safetensors", ".bin", ".gguf",
    ".db", ".sqlite", ".sqlite3", ".key", ".pem", ".token",
}
EXCLUDE_NAMES = {".env", "kaggle.json", ".netrc"}


def should_include(path: Path, root: Path) -> bool:
    rel = path.relative_to(root)
    if set(rel.parts) & EXCLUDE_DIRS:
        return False
    if path.name in EXCLUDE_NAMES:
        return False
    if path.suffix.lower() in EXCLUDE_SUFFIXES:
        return False
    return path.is_file()


def package(root: Path, out: Path) -> int:
    out.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with zipfile.ZipFile(out, "w", compression=zipfile.ZIP_DEFLATED) as z:
        for path in sorted(root.rglob("*")):
            if should_include(path, root):
                arcname = root.name + "/" + str(path.relative_to(root))
                z.write(path, arcname)
                count += 1
    return count


def main() -> None:
    parser = argparse.ArgumentParser(description="Create a public-safe Ares Static Space zip.")
    parser.add_argument("--root", default=".")
    parser.add_argument("--out", default="../ares-static-space-public.zip")
    args = parser.parse_args()
    root = Path(args.root).resolve()
    out = Path(args.out).resolve()
    count = package(root, out)
    print(f"Wrote {out} with {count} files")


if __name__ == "__main__":
    main()