analytics/instrumentation.py

import datetime
import functools
import logging
from advanced_workflow_orchestrator import AdvancedWorkflowOrchestrator
from analytics.collector import AsyncAnalyticsCollector

logger = logging.getLogger(__name__)

def activate_analytics(orchestrator_instance: AdvancedWorkflowOrchestrator):
    """
    Monkey-patches the AdvancedWorkflowOrchestrator's execution method
    to automatically record execution metrics.
    """
    logger.info("🧬 Activating Workflow DNA (Analytics Instrumentation) for AdvancedWorkflowOrchestrator...")
    
    # Capture the original method
    # Note: simple assignment like `original = inst.method` captures a bound method.
    # We need to be careful not to create infinite recursion.
    original_method = orchestrator_instance._execute_workflow_step
    
    # Create the wrapper
    async def instrumented_execute_step(self, workflow, step_id, context):
        start_time = datetime.datetime.now()
        
        # Resolve step details
        step = next((s for s in workflow.steps if s.step_id == step_id), None)
        step_type = step.step_type.value if step and hasattr(step.step_type, 'value') else "unknown"
        if step_type == "unknown" and step:
             step_type = str(step.step_type)

        status = "COMPLETED"
        error = None
        
        try:
            # Call original method
            # Since original_method is already bound to the instance (if we captured it from instance),
            # we might not need to pass 'self' again if we just call it?
            # actually, if we replace the method on the instance, 'self' will be passed 
            # to our wrapper.
            # But 'original_method' is the OLD bound method.
            # So we should call `original_method(workflow, step_id, context)` directly depending on how it was captured.
            
            # If we captured `inst._execute_workflow_step`, it IS a bound method.
            # So we don't pass self.
            await original_method(workflow, step_id, context)
            
            # Check context for status (it might have failed inside)
            result = context.results.get(step_id)
            if result and result.get("status") == "failed":
                status = "FAILED"
                error = result.get("error")
                
        except Exception as e:
            status = "FAILED"
            error = str(e)
            raise e
            
        finally:
            end_time = datetime.datetime.now()
            
            # Log to sidecar
            # MAPPING FIX: 
            # 1. WorkflowContext.workflow_id IS the execution_id (e.g. exec_123)
            # 2. Real Workflow Definition ID is usually in input_data OR we treat the execution ID as the workflow ID if missing.
            
            execution_id = getattr(context, 'workflow_id', 'unknown')
            
            # Try to find the Definition ID
            workflow_def_id = context.input_data.get("_ui_workflow_id") if context.input_data else None
            # Fallback: if we can't find a definition ID, use the execution ID or 'ad-hoc'
            if not workflow_def_id:
                 workflow_def_id = "ad-hoc" 

            if execution_id != "unknown":
                await AsyncAnalyticsCollector.get_instance().log_step(
                    execution_id=execution_id,
                    workflow_id=workflow_def_id,
                    step_id=step_id,
                    step_type=step_type,
                    start_time=start_time,
                    end_time=end_time,
                    status=status,
                    error=error,
                    results=context.results.get(step_id) if hasattr(context, 'results') else None
                )
            
    # Apply the patch
    # We are replacing a BOUND method on the instance.
    # The wrapper function `instrumented_execute_step` expects `self` as first arg.
    # We need to bind it to the instance manually or use partial.
    orchestrator_instance._execute_workflow_step = functools.partial(instrumented_execute_step, orchestrator_instance)
    
    logger.info("✅ Workflow DNA Active: Instrumentation applied to AdvancedWorkflowOrchestrator.")