backend/tasks/unified_evaluation_task.py

"""
UNIFIED EVALUATION CELERY TASK: Background execution using unified pipeline.

This task provides asynchronous execution of the unified pipeline
through Celery workers for production scalability.
"""

import logging
from celery import Celery
from typing import Dict, Any

from workers.celery_worker import celery_app
from services.evaluation_service import EvaluationService
from core.database import get_db

logger = logging.getLogger(__name__)


@celery_app.task(bind=True)
def run_unified_evaluation_task(self, evaluation_id: str, user_id: int) -> Dict[str, Any]:
    """
    Celery task for unified evaluation execution.
    
    This task runs the complete unified pipeline in the background:
    - Dataset loading and preparation
    - Learning engine insights
    - Red team pipeline execution
    - Scoring and audit trails
    - Analytics data flow
    - Report generation
    
    Args:
        evaluation_id: Evaluation ID to execute
        user_id: User ID requesting execution
        
    Returns:
        Dict with execution results
    """
    logger.info(f"🚀 Starting unified evaluation task: {evaluation_id}")
    
    try:
        # Create async session for database operations
        async def execute_evaluation():
            async for db in get_db():
                # Initialize evaluation service
                evaluation_service = EvaluationService(db)
                
                # Execute unified evaluation
                result = await evaluation_service.execute_unified_evaluation(
                    evaluation_id=evaluation_id,
                    user_id=user_id
                )
                
                logger.info(f"✅ Unified evaluation task completed: {evaluation_id}")
                return result
        
        # Run async function in sync context
        import asyncio
        return asyncio.run(execute_evaluation())
        
    except Exception as e:
        logger.error(f"❌ Unified evaluation task failed: {evaluation_id} - {str(e)}")
        
        # Update task status and return error
        self.update_state(
            state='FAILURE',
            meta={'error': str(e), 'evaluation_id': evaluation_id}
        )
        
        raise


@celery_app.task
def cleanup_unified_evaluation_resources(evaluation_id: str) -> Dict[str, Any]:
    """
    Cleanup task for unified evaluation resources.
    
    This task cleans up temporary resources after evaluation completion:
    - Temporary files
    - Memory caches
    - Connection pools
    
    Args:
        evaluation_id: Evaluation ID to cleanup
        
    Returns:
        Dict with cleanup results
    """
    logger.info(f"🧹 Starting cleanup for unified evaluation: {evaluation_id}")
    
    try:
        cleanup_results = {
            "evaluation_id": evaluation_id,
            "temp_files_cleaned": 0,
            "memory_cleared": True,
            "connections_closed": True,
            "cleanup_completed_at": None
        }
        
        # Implement cleanup logic here
        # This would integrate with the unified pipeline cleanup methods
        
        logger.info(f"✅ Cleanup completed for unified evaluation: {evaluation_id}")
        return cleanup_results
        
    except Exception as e:
        logger.error(f"❌ Cleanup failed for unified evaluation: {evaluation_id} - {str(e)}")
        raise


# Task chaining for complete workflow
def create_unified_evaluation_workflow(evaluation_id: str, user_id: int):
    """
    Create a complete workflow for unified evaluation.
    
    This chains the execution and cleanup tasks for a complete workflow.
    
    Args:
        evaluation_id: Evaluation ID
        user_id: User ID
        
    Returns:
        Celery chain result
    """
    from celery import chain
    
    workflow = chain(
        run_unified_evaluation_task.s(evaluation_id, user_id),
        cleanup_unified_evaluation_resources.s(evaluation_id)
    )
    
    return workflow()