"""
Cleanup script to remove extraneous files from the root directory.

This script removes files that have already been organized into the src directory structure.
"""

import os
import glob
import sys

def confirm_deletion(files):
    """
    Ask for confirmation before deleting files.
    
    Parameters:
        files (list): List of files to delete
        
    Returns:
        bool: True if user confirms deletion, False otherwise
    """
    if not files:
        print("No files to delete.")
        return False
    
    print("The following files will be deleted:")
    for file in files:
        print(f"  - {file}")
    
    response = input("\nAre you sure you want to delete these files? (y/n): ")
    return response.lower() == 'y'


def cleanup_root_directory():
    """
    Remove extraneous files from the root directory.
    """
    # Files to remove
    python_files = [
        "lab2_part1.py",
        "lab2_part2.py",
        "lab2_part1_animated.py",
        "lab2_part2_animated.py",
        "animate_bouncing_balls.py"
    ]
    
    # Data files
    data_files = glob.glob("*.csv")
    
    # Image files
    image_files = glob.glob("*.png")
    
    # Combine all files
    all_files = python_files + data_files + image_files
    
    # Filter out files that don't exist
    files_to_delete = [file for file in all_files if os.path.exists(file)]
    
    # Ask for confirmation
    if confirm_deletion(files_to_delete):
        # Delete files
        for file in files_to_delete:
            try:
                os.remove(file)
                print(f"Deleted: {file}")
            except Exception as e:
                print(f"Error deleting {file}: {e}")
        
        print("\nCleanup complete!")
    else:
        print("\nCleanup cancelled.")


if __name__ == "__main__":
    # Check if script is run from the root directory
    if not os.path.exists("src/utils/cleanup.py"):
        print("Error: This script must be run from the project root directory.")
        print("Please run: python src/utils/cleanup.py")
        sys.exit(1)
    
    cleanup_root_directory() 