"""
Unified Celery task entrypoint for all AI functions
"""
import logging
from celery import shared_task
from igny8_core.ai.engine import AIEngine
from igny8_core.ai.registry import get_function_instance

logger = logging.getLogger(__name__)


@shared_task(bind=True, max_retries=3)
def run_ai_task(self, function_name: str, payload: dict, account_id: int = None):
    """
    Single Celery entrypoint for all AI functions.
    Dynamically loads and executes the requested function.
    
    Args:
        function_name: Name of the AI function (e.g., 'auto_cluster')
        payload: Function-specific payload
        account_id: Account ID for account isolation
    """
    logger.info("=" * 80)
    logger.info(f"run_ai_task STARTED: {function_name}")
    logger.info(f"  - Task ID: {self.request.id}")
    logger.info(f"  - Function: {function_name}")
    logger.info(f"  - Account ID: {account_id}")
    logger.info(f"  - Payload keys: {list(payload.keys())}")
    logger.info("=" * 80)
    
    try:
        # Get account
        account = None
        if account_id:
            from igny8_core.auth.models import Account
            try:
                account = Account.objects.get(id=account_id)
            except Account.DoesNotExist:
                logger.warning(f"Account {account_id} not found")
        
        # Get function from registry
        fn = get_function_instance(function_name)
        if not fn:
            error_msg = f'Function {function_name} not found in registry'
            logger.error(error_msg)
            return {
                'success': False,
                'error': error_msg
            }
        
        # Create engine and execute
        engine = AIEngine(celery_task=self, account=account)
        result = engine.execute(fn, payload)
        
        logger.info("=" * 80)
        logger.info(f"run_ai_task COMPLETED: {function_name}")
        logger.info(f"  - Success: {result.get('success')}")
        if not result.get('success'):
            logger.error(f"  - Error: {result.get('error')}")
        logger.info("=" * 80)
        
        # If execution failed, update state and return error (don't raise to avoid serialization issues)
        if not result.get('success'):
            error_msg = result.get('error', 'Task execution failed')
            error_type = result.get('error_type', 'ExecutionError')
            # Update task state with error details
            error_meta = {
                'error': error_msg,
                'error_type': error_type,
                'function_name': function_name,
                'phase': result.get('phase', 'ERROR'),
                'percentage': 0,
                'message': f'Error: {error_msg}',
                'request_steps': result.get('request_steps', []),
                'response_steps': result.get('response_steps', [])
            }
            try:
                self.update_state(
                    state='FAILURE',
                    meta=error_meta
                )
            except Exception as update_err:
                logger.warning(f"Failed to update task state: {update_err}")
            
            # Return error result - Celery will mark as FAILURE based on state
            # Don't raise exception to avoid serialization issues
            return {
                'success': False,
                'error': error_msg,
                'error_type': error_type,
                **error_meta
            }
        
        return result
        
    except Exception as e:
        error_type = type(e).__name__
        error_msg = str(e)
        
        logger.error("=" * 80)
        logger.error(f"run_ai_task FAILED: {function_name}")
        logger.error(f"  - Error: {error_type}: {error_msg}")
        logger.error("=" * 80, exc_info=True)
        
        # Update task state with error details (don't raise to avoid serialization issues)
        error_meta = {
            'error': error_msg,
            'error_type': error_type,
            'function_name': function_name,
            'phase': 'ERROR',
            'percentage': 0,
            'message': f'Error: {error_msg}'
        }
        try:
            self.update_state(
                state='FAILURE',
                meta=error_meta
            )
        except Exception as update_err:
            logger.warning(f"Failed to update task state: {update_err}")
        
        # Return error result - don't raise to avoid Celery serialization issues
        return {
            'success': False,
            'error': error_msg,
            'error_type': error_type,
            'function_name': function_name,
            **error_meta
        }

