#!/usr/bin/env python3
"""
Cleanup script to manage disk space and prevent service issues
"""
import os
import shutil
import glob
from datetime import datetime, timedelta

def cleanup_logs():
    """Clean up old log files"""
    log_dirs = ["/data/logs", "/var/log"]
    
    for log_dir in log_dirs:
        if os.path.exists(log_dir):
            # Remove log files older than 7 days
            cutoff_date = datetime.now() - timedelta(days=7)
            
            for log_file in glob.glob(os.path.join(log_dir, "*.log*")):
                try:
                    file_time = datetime.fromtimestamp(os.path.getmtime(log_file))
                    if file_time < cutoff_date:
                        os.remove(log_file)
                        print(f"[Cleanup] Removed old log: {log_file}")
                except Exception as e:
                    print(f"[Cleanup] Error removing {log_file}: {e}")

def cleanup_temp_files():
    """Clean up temporary files"""
    temp_dirs = ["/tmp", "/data/merged/temp"]
    
    for temp_dir in temp_dirs:
        if os.path.exists(temp_dir):
            # Remove files older than 1 day
            cutoff_date = datetime.now() - timedelta(days=1)
            
            for temp_file in glob.glob(os.path.join(temp_dir, "*")):
                try:
                    if os.path.isfile(temp_file):
                        file_time = datetime.fromtimestamp(os.path.getmtime(temp_file))
                        if file_time < cutoff_date:
                            os.remove(temp_file)
                            print(f"[Cleanup] Removed temp file: {temp_file}")
                except Exception as e:
                    print(f"[Cleanup] Error removing {temp_file}: {e}")

def cleanup_old_data():
    """Clean up old data files to save space"""
    # Keep only last 30 days of archived data
    archive_dir = "/data/merged/archive"
    if os.path.exists(archive_dir):
        cutoff_date = datetime.now() - timedelta(days=30)
        
        for archive_folder in os.listdir(archive_dir):
            folder_path = os.path.join(archive_dir, archive_folder)
            if os.path.isdir(folder_path):
                try:
                    folder_time = datetime.fromtimestamp(os.path.getmtime(folder_path))
                    if folder_time < cutoff_date:
                        shutil.rmtree(folder_path)
                        print(f"[Cleanup] Removed old archive: {folder_path}")
                except Exception as e:
                    print(f"[Cleanup] Error removing {folder_path}: {e}")

def get_disk_usage():
    """Get current disk usage"""
    try:
        import psutil
        # Check disk usage for the data mount if present
        disk_usage = psutil.disk_usage('/data' if os.path.exists('/data') else '/')
        free_gb = disk_usage.free / (1024**3)
        used_percent = (disk_usage.used / disk_usage.total) * 100
        return free_gb, used_percent
    except Exception:
        return None, None

def main():
    """Main cleanup function"""
    print(f"[Cleanup] Starting cleanup at {datetime.now()}")
    
    # Check disk usage before cleanup
    free_before, used_before = get_disk_usage()
    if free_before:
        print(f"[Cleanup] Disk usage before: {used_before:.1f}% used, {free_before:.1f}GB free")
    
    # Run cleanup tasks
    cleanup_logs()
    cleanup_temp_files()
    cleanup_old_data()
    
    # Check disk usage after cleanup
    free_after, used_after = get_disk_usage()
    if free_after and free_before:
        freed_space = free_after - free_before
        print(f"[Cleanup] Disk usage after: {used_after:.1f}% used, {free_after:.1f}GB free")
        if freed_space > 0:
            print(f"[Cleanup] Freed {freed_space:.2f}GB of disk space")
    
    print(f"[Cleanup] Cleanup completed at {datetime.now()}")

if __name__ == "__main__":
    main()