#!/usr/bin/env python3
"""
Cleanup script for Hugging Face Spaces to remove cached models and free up storage.
Run this if you encounter storage limit errors.
"""

import os
import shutil
from pathlib import Path

def get_size(path):
    """Get size of directory in GB."""
    total = 0
    try:
        for entry in os.scandir(path):
            if entry.is_file(follow_symlinks=False):
                total += entry.stat().st_size
            elif entry.is_dir(follow_symlinks=False):
                total += get_size(entry.path)
    except PermissionError:
        pass
    return total / (1024**3)  # Convert to GB

def cleanup():
    """Remove cache directories to free up space."""
    cache_dirs = [
        Path.home() / ".cache" / "huggingface",
        Path.home() / ".cache" / "torch",
        Path("/tmp/huggingface_cache"),
        Path("/tmp/torch_cache"),
    ]
    
    total_freed = 0
    for cache_dir in cache_dirs:
        if cache_dir.exists():
            size = get_size(str(cache_dir))
            print(f"Removing {cache_dir} ({size:.2f} GB)...")
            try:
                shutil.rmtree(cache_dir)
                total_freed += size
                print(f"  ✓ Removed")
            except Exception as e:
                print(f"  ✗ Failed: {e}")
    
    print(f"\nTotal space freed: {total_freed:.2f} GB")

if __name__ == "__main__":
    cleanup()