#!/usr/bin/env python3
"""
Immediate Model Cleanup Script

This script will clean up all existing model files and prepare the system
for fresh training with the new model management system.
"""

import logging
import sys
from model_manager import ModelManager

def main():
    """Run the model cleanup"""
    
    # Configure logging for better output
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    print("=" * 60)
    print("GOGO2 MODEL CLEANUP SYSTEM")
    print("=" * 60)
    print()
    print("This script will:")
    print("1. Delete ALL existing model files (.pt, .pth)")
    print("2. Remove ALL checkpoint directories")
    print("3. Clear model backup directories")
    print("4. Reset the model registry")
    print("5. Create clean directory structure")
    print()
    print("WARNING: This action cannot be undone!")
    print()
    
    # Calculate current space usage first
    try:
        manager = ModelManager()
        storage_stats = manager.get_storage_stats()
        print(f"Current storage usage:")
        print(f"- Models: {storage_stats['total_models']}")
        print(f"- Size: {storage_stats['actual_size_mb']:.1f}MB")
        print()
    except Exception as e:
        print(f"Error checking current storage: {e}")
        print()
    
    # Ask for confirmation
    print("Type 'CLEANUP' to proceed with the cleanup:")
    user_input = input("> ").strip()
    
    if user_input != "CLEANUP":
        print("Cleanup cancelled. No changes made.")
        return
    
    print()
    print("Starting cleanup...")
    print("-" * 40)
    
    try:
        # Create manager and run cleanup
        manager = ModelManager()
        cleanup_result = manager.cleanup_all_existing_models(confirm=True)
        
        print()
        print("=" * 60)
        print("CLEANUP COMPLETE")
        print("=" * 60)
        print(f"Files deleted: {cleanup_result['deleted_files']}")
        print(f"Space freed: {cleanup_result['freed_space_mb']:.1f} MB")
        print(f"Directories cleaned: {len(cleanup_result['deleted_directories'])}")
        
        if cleanup_result['errors']:
            print(f"Errors encountered: {len(cleanup_result['errors'])}")
            print("Errors:")
            for error in cleanup_result['errors'][:5]:  # Show first 5 errors
                print(f"  - {error}")
            if len(cleanup_result['errors']) > 5:
                print(f"  ... and {len(cleanup_result['errors']) - 5} more")
        
        print()
        print("System is now ready for fresh model training!")
        print("The following directories have been created:")
        print("- models/best_models/")
        print("- models/cnn/")
        print("- models/rl/")
        print("- models/checkpoints/")
        print("- NN/models/saved/")
        print()
        print("New models will be automatically managed by the ModelManager.")
        
    except Exception as e:
        print(f"Error during cleanup: {e}")
        logging.exception("Cleanup failed")
        sys.exit(1)

if __name__ == "__main__":
    main()