#!/usr/bin/env python3
"""
Bakery IA Disk Space Cleanup Script
===================================

This script performs comprehensive cleanup of Docker and Kubernetes resources
to prevent disk space exhaustion during development.

Features:
- Automatic cleanup based on disk space thresholds
- Manual cleanup on demand
- Comprehensive resource cleanup (images, containers, volumes, etc.)
- Detailed reporting and logging

Usage:
  ./scripts/cleanup_disk_space.py [--manual] [--threshold GB] [--verbose]

Environment Variables:
  TILT_DISK_THRESHOLD_GB - Minimum free space required (default: 10GB)
  TILT_CLEANUP_VERBOSE - Set to "true" for verbose output
"""

import subprocess
import sys
import os
import argparse
import time
from datetime import datetime

def get_disk_space():
    """Get available disk space in GB"""
    try:
        result = subprocess.run(['df', '/', '--output=avail', '-h'],
                              capture_output=True, text=True, check=True)
        # Extract numeric value from output like "15G"
        output = result.stdout.strip().split('\n')[-1].strip()
        if 'G' in output:
            return float(output.replace('G', ''))
        elif 'M' in output:
            return float(output.replace('M', '')) / 1024
        else:
            return 0
    except Exception as e:
        print(f"⚠️  Could not check disk space: {e}")
        return 999  # Assume plenty of space if we can't check

def cleanup_docker_images(verbose=False):
    """Clean up old and unused Docker images"""
    if verbose:
        print("🧹 Cleaning up Docker images...")
    
    try:
        # Remove dangling images
        if verbose:
            print("  Removing dangling images...")
        subprocess.run(['docker', 'image', 'prune', '-f'],
                      capture_output=True, text=True)
        
        # Remove unused images (not referenced by any container)
        if verbose:
            print("  Removing unused images...")
        subprocess.run(['docker', 'image', 'prune', '-a', '-f'],
                      capture_output=True, text=True)
        
        # Remove old images (older than 2 hours)
        if verbose:
            print("  Removing old images (>2 hours)...")
        subprocess.run(['docker', 'image', 'prune', '-a', '-f',
                       '--filter', 'until=2h'],
                      capture_output=True, text=True)
        
        if verbose:
            print("✅ Docker image cleanup completed")
        return True
    except Exception as e:
        print(f"⚠️  Docker image cleanup failed: {e}")
        return False

def cleanup_docker_containers(verbose=False):
    """Clean up stopped containers"""
    if verbose:
        print("🧹 Cleaning up Docker containers...")
    
    try:
        # Remove stopped containers
        if verbose:
            print("  Removing stopped containers...")
        subprocess.run(['docker', 'container', 'prune', '-f'],
                      capture_output=True, text=True)
        
        # Remove old containers (older than 1 hour)
        if verbose:
            print("  Removing old containers (>1 hour)...")
        subprocess.run(['docker', 'container', 'prune', '-f',
                       '--filter', 'until=1h'],
                      capture_output=True, text=True)
        
        if verbose:
            print("✅ Docker container cleanup completed")
        return True
    except Exception as e:
        print(f"⚠️  Docker container cleanup failed: {e}")
        return False

def cleanup_docker_volumes(verbose=False):
    """Clean up unused volumes"""
    if verbose:
        print("🧹 Cleaning up Docker volumes...")
    
    try:
        # Remove unused volumes
        if verbose:
            print("  Removing unused volumes...")
        subprocess.run(['docker', 'volume', 'prune', '-f'],
                      capture_output=True, text=True)
        
        if verbose:
            print("✅ Docker volume cleanup completed")
        return True
    except Exception as e:
        print(f"⚠️  Docker volume cleanup failed: {e}")
        return False

def cleanup_docker_system(verbose=False):
    """Clean up Docker system (build cache, networks, etc.)"""
    if verbose:
        print("🧹 Cleaning up Docker system...")
    
    try:
        # Remove build cache
        if verbose:
            print("  Removing build cache...")
        subprocess.run(['docker', 'builder', 'prune', '-f'],
                      capture_output=True, text=True)
        
        # Remove unused networks
        if verbose:
            print("  Removing unused networks...")
        subprocess.run(['docker', 'network', 'prune', '-f'],
                      capture_output=True, text=True)
        
        if verbose:
            print("✅ Docker system cleanup completed")
        return True
    except Exception as e:
        print(f"⚠️  Docker system cleanup failed: {e}")
        return False

def cleanup_kubernetes_resources(verbose=False):
    """Clean up Kubernetes resources"""
    if verbose:
        print("🧹 Cleaning up Kubernetes resources...")
    
    try:
        # Remove completed jobs older than 1 hour
        if verbose:
            print("  Removing completed jobs (>1 hour)...")
        subprocess.run(['kubectl', 'delete', 'jobs', '-n', 'bakery-ia', 
                       '--field-selector=status.successful=1'],
                      capture_output=True, text=True)
        
        # Remove failed jobs older than 1 hour
        if verbose:
            print("  Removing failed jobs (>1 hour)...")
        subprocess.run(['kubectl', 'delete', 'jobs', '-n', 'bakery-ia',
                       '--field-selector=status.failed>0'],
                      capture_output=True, text=True)
        
        if verbose:
            print("✅ Kubernetes resource cleanup completed")
        return True
    except Exception as e:
        print(f"⚠️  Kubernetes resource cleanup failed: {e}")
        return False

def perform_cleanup(manual=False, threshold_gb=10, verbose=False):
    """Perform comprehensive cleanup"""
    
    print("\n" + "="*60)
    print("🚀 STARTING COMPREHENSIVE CLEANUP")
    print("="*60)
    
    if manual:
        print("🎛️  Mode: MANUAL (forced cleanup)")
    else:
        print("🎛️  Mode: AUTOMATIC (threshold-based)")
    
    print(f"📊 Threshold: {threshold_gb}GB free space")
    
    # Check disk space before cleanup
    free_space_before = get_disk_space()
    print(f"📊 Disk space before cleanup: {free_space_before:.1f}GB free")
    
    # Check if cleanup is needed (unless manual)
    if not manual and free_space_before >= threshold_gb:
        print("✅ Sufficient disk space available, skipping cleanup")
        return True
    
    cleanup_results = []
    
    # Perform all cleanup operations
    cleanup_results.append(("Docker Images", cleanup_docker_images(verbose)))
    cleanup_results.append(("Docker Containers", cleanup_docker_containers(verbose)))
    cleanup_results.append(("Docker Volumes", cleanup_docker_volumes(verbose)))
    cleanup_results.append(("Docker System", cleanup_docker_system(verbose)))
    cleanup_results.append(("Kubernetes Resources", cleanup_kubernetes_resources(verbose)))
    
    # Check disk space after cleanup
    free_space_after = get_disk_space()
    space_reclaimed = free_space_after - free_space_before
    
    print(f"\n📊 Disk space after cleanup: {free_space_after:.1f}GB free")
    print(f"🎯 Space reclaimed: {space_reclaimed:.1f}GB")
    
    # Summary
    print("\n📋 CLEANUP SUMMARY:")
    for name, success in cleanup_results:
        status = "✅ SUCCESS" if success else "❌ FAILED"
        print(f"  {name}: {status}")
    
    print("="*60)
    print("✅ CLEANUP COMPLETED")
    print("="*60 + "\n")
    
    return True

def main():
    parser = argparse.ArgumentParser(
        description='Bakery IA Disk Space Cleanup Script',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  ./cleanup_disk_space.py              # Automatic cleanup (checks threshold)
  ./cleanup_disk_space.py --manual     # Force cleanup regardless of threshold
  ./cleanup_disk_space.py --threshold 5  # Use 5GB threshold
  ./cleanup_disk_space.py --verbose    # Verbose output
        """
    )
    
    parser.add_argument('--manual', action='store_true',
                       help='Force cleanup regardless of disk space threshold')
    parser.add_argument('--threshold', type=int, default=10,
                       help='Minimum free space required in GB (default: 10)')
    parser.add_argument('--verbose', action='store_true',
                       help='Enable verbose output')
    
    args = parser.parse_args()
    
    # Get threshold from environment variable if set
    env_threshold = os.getenv('TILT_DISK_THRESHOLD_GB')
    if env_threshold:
        try:
            args.threshold = int(env_threshold)
        except ValueError:
            pass
    
    # Get verbose from environment variable if set
    env_verbose = os.getenv('TILT_CLEANUP_VERBOSE', 'false').lower()
    if env_verbose == 'true':
        args.verbose = True
    
    return perform_cleanup(
        manual=args.manual,
        threshold_gb=args.threshold,
        verbose=args.verbose
    )

if __name__ == '__main__':
    success = main()
    sys.exit(0 if success else 1)