"""
DICOM file cleanup utility to ensure temporary DICOM files are removed even on script termination.
"""
import os
import sys
import shutil
import atexit
import signal
from config import settings
from utils.logger import main_logger as logger

# Global registry of directories to clean up
_cleanup_dirs = set()

def register_cleanup_dir(directory):
    """
    Register a directory for cleanup when the script exits.
    
    Args:
        directory (str): Path to the directory to be cleaned up
    """
    global _cleanup_dirs
    if os.path.exists(directory) and os.path.isdir(directory):
        _cleanup_dirs.add(directory)
        logger.debug(f"Registered directory for cleanup: {directory}")

def cleanup_dicom_files():
    """
    Remove all registered DICOM directories.
    This function is called when the script exits, ensuring cleanup even on abnormal termination.
    """
    global _cleanup_dirs
    logger.info(f"Running cleanup for {len(_cleanup_dirs)} directories")
    
    for directory in _cleanup_dirs:
        try:
            if os.path.exists(directory):
                shutil.rmtree(directory)
                logger.info(f"Cleanup: Removed DICOM files from {directory}")
        except Exception as e:
            logger.error(f"Cleanup: Failed to remove DICOM files from {directory}: {str(e)}")
    
    # Clear the registry after cleanup
    _cleanup_dirs.clear()

def cleanup_study_directory(study_uid):
    """
    Immediately remove the directory for a specific study.
    
    Args:
        study_uid (str): Study Instance UID of the study to clean up
        
    Returns:
        bool: True if cleanup was successful, False otherwise
    """
    study_dir = os.path.join(settings.DICOM_STORE_DIR, study_uid)
    
    if not os.path.exists(study_dir):
        logger.debug(f"Cleanup: Study directory does not exist: {study_dir}")
        return True
        
    try:
        shutil.rmtree(study_dir)
        logger.info(f"Cleanup: Removed DICOM files from {study_dir}")
        
        # Also remove from the registry if it's there
        global _cleanup_dirs
        if study_dir in _cleanup_dirs:
            _cleanup_dirs.remove(study_dir)
            
        return True
    except Exception as e:
        logger.error(f"Cleanup: Failed to remove DICOM files from {study_dir}: {str(e)}")
        return False

def register_exit_handlers():
    """
    Register cleanup handlers for various exit scenarios.
    """
    # Register for normal exit
    atexit.register(cleanup_dicom_files)
    
    # Register for signals
    signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
    signal.signal(signal.SIGTERM, signal_handler)  # Termination request
    
    logger.info("Registered cleanup handlers for script exit")

def signal_handler(sig, frame):
    """
    Handle termination signals by performing cleanup before exit.
    
    Args:
        sig: Signal number
        frame: Current stack frame
    """
    signal_name = {
        signal.SIGINT: 'SIGINT (Ctrl+C)',
        signal.SIGTERM: 'SIGTERM'
    }.get(sig, f'Signal {sig}')
    
    logger.info(f"Received {signal_name}, cleaning up...")
    cleanup_dicom_files()
    # Don't call exit() directly in a web application context
    # Let the calling application handle the exit process
    # This fix makes it compatible with FastAPI and other web frameworks
    # Raise SystemExit instead of calling exit() directly
    if hasattr(frame, '_is_web_app') and frame._is_web_app:
        return
    else:
        # For command-line applications, we can exit directly
        sys.exit(0)