#!/usr/bin/env python
import os
import shutil
import sys

def clean_pycache(root_dir):
    """
    Recursively deletes all __pycache__ directories and .pyc files.
    """
    print(f"Cleaning __pycache__ in: {root_dir}")
    for root, dirs, files in os.walk(root_dir):
        # Delete __pycache__ directories
        if "__pycache__" in dirs:
            pycache_path = os.path.join(root, "__pycache__")
            print(f"Removing directory: {pycache_path}")
            shutil.rmtree(pycache_path)
            dirs.remove("__pycache__")  # Don't visit the deleted directory

        # Delete orphan .pyc files
        for file in files:
            if file.endswith(".pyc"):
                pyc_path = os.path.join(root, file)
                print(f"Removing file: {pyc_path}")
                os.remove(pyc_path)

def clean_extensions(root_dir):
    """
    Recursively deletes all compiled .so and .pyd files.
    """
    print(f"Cleaning compiled extensions in: {root_dir}")
    for root, dirs, files in os.walk(root_dir):
        for file in files:
            if file.endswith(".so") or file.endswith(".pyd"):
                ext_path = os.path.join(root, file)
                print(f"Removing extension: {ext_path}")
                os.remove(ext_path)

if __name__ == "__main__":
    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    
    # Clean pycache
    clean_pycache(project_root)
    
    # Optionally clean extensions if requested via argument
    if len(sys.argv) > 1 and sys.argv[1] == "--all":
        clean_extensions(project_root)
    
    print("Cleanup complete.")