Trixy3/trixy_core/ml/events.py

"""
Event System Integration for Trixy ML Components.

This module provides comprehensive event system integration for all ML components
including wakeword detection, voice recognition, and workflow management. It defines
ML-specific event data structures and provides event handlers that integrate with
the main Trixy event system.

Key Features:
- ML-specific event data structures
- Event handler registration and management
- Integration with main Trixy event system
- Performance event monitoring
- Debug and diagnostic events
- Event correlation and tracking
- Real-time event streaming
- Event persistence and replay

Usage:
    from trixy_core.ml.events import MLEventManager, MLEventData
    
    # Create ML event manager
    ml_events = MLEventManager(main_event_handler)
    
    # Register ML component
    ml_events.register_ml_component(ml_manager)
    
    # Trigger ML events
    ml_events.trigger_wakeword_detected(detection_data)
"""

import logging
import threading
import time
import queue
from typing import Dict, List, Optional, Any, Callable, Union, Set
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
import uuid
from collections import defaultdict, deque

# Import event system components
try:
    from ..events import (
        EventHandler, TrixyEvent, EventType, TrixyEventData,
        WakewordReceivedEventData, RawAudioInputReceivedEventData,
        SpeakerInfo, EventDataFactory
    )
    EVENTS_AVAILABLE = True
except ImportError:
    # Create mock classes if event system not available
    EVENTS_AVAILABLE = False
    EventHandler = object
    TrixyEvent = lambda events: lambda func: func
    EventType = object
    TrixyEventData = object
    WakewordReceivedEventData = object
    RawAudioInputReceivedEventData = object
    SpeakerInfo = object
    EventDataFactory = object

# Configure logging
logger = logging.getLogger(__name__)


class MLEventType(Enum):
    """ML-specific event types."""
    # Core ML events
    ML_SYSTEM_STARTED = "ml_system_started"
    ML_SYSTEM_STOPPED = "ml_system_stopped"
    ML_SYSTEM_ERROR = "ml_system_error"
    
    # Wakeword events
    WAKEWORD_DETECTED = "wakeword_detected"
    WAKEWORD_ENHANCED = "wakeword_enhanced"
    WAKEWORD_FALSE_POSITIVE = "wakeword_false_positive"
    
    # Speaker recognition events
    SPEAKER_IDENTIFIED = "speaker_identified"
    SPEAKER_VERIFIED = "speaker_verified"
    SPEAKER_ENROLLED = "speaker_enrolled"
    SPEAKER_REMOVED = "speaker_removed"
    SPEAKER_UPDATED = "speaker_updated"
    
    # Audio processing events
    AUDIO_PROCESSING_STARTED = "audio_processing_started"
    AUDIO_PROCESSING_STOPPED = "audio_processing_stopped"
    AUDIO_QUALITY_WARNING = "audio_quality_warning"
    AUDIO_BUFFER_OVERRUN = "audio_buffer_overrun"
    SILENCE_DETECTED = "silence_detected"
    VOICE_ACTIVITY_DETECTED = "voice_activity_detected"
    
    # Workflow events
    WORKFLOW_STARTED = "workflow_started"
    WORKFLOW_STEP_COMPLETED = "workflow_step_completed"
    WORKFLOW_COMPLETED = "workflow_completed"
    WORKFLOW_FAILED = "workflow_failed"
    WORKFLOW_TIMEOUT = "workflow_timeout"
    
    # Performance events
    PERFORMANCE_METRICS_UPDATED = "performance_metrics_updated"
    LATENCY_WARNING = "latency_warning"
    MEMORY_WARNING = "memory_warning"
    GPU_ERROR = "gpu_error"
    
    # Model events
    MODEL_LOADED = "model_loaded"
    MODEL_UNLOADED = "model_unloaded"
    MODEL_ERROR = "model_error"
    MODEL_UPDATED = "model_updated"
    
    # Network events
    SERVER_SELECTION_RECEIVED = "server_selection_received"
    AUDIO_STREAMING_STARTED = "audio_streaming_started"
    AUDIO_STREAMING_STOPPED = "audio_streaming_stopped"
    NETWORK_ERROR = "network_error"


class EventPriority(Enum):
    """Event priority levels."""
    LOW = "low"
    NORMAL = "normal"
    HIGH = "high"
    CRITICAL = "critical"


@dataclass
class MLEventData(TrixyEventData if EVENTS_AVAILABLE else object):
    """Base class for ML event data."""
    event_type: str = ""
    correlation_id: Optional[str] = None
    session_id: Optional[str] = None
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    source_component: str = "ml_system"
    priority: EventPriority = EventPriority.NORMAL
    processing_time_ms: float = 0.0
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class WakewordEventData(MLEventData):
    """Event data for wakeword detection events."""
    wakeword_type: str = ""
    confidence: float = 0.0
    audio_chunk_position: int = 0
    speaker_id: Optional[str] = None
    speaker_name: Optional[str] = None
    speaker_confidence: float = 0.0
    enhancement_data: Dict[str, Any] = field(default_factory=dict)


@dataclass
class SpeakerEventData(MLEventData):
    """Event data for speaker recognition events."""
    speaker_id: str = ""
    speaker_name: str = ""
    confidence: float = 0.0
    similarity_score: float = 0.0
    embedding_quality: float = 0.0
    verification_result: Optional[bool] = None
    enrollment_status: Optional[str] = None


@dataclass
class AudioEventData(MLEventData):
    """Event data for audio processing events."""
    audio_metrics: Dict[str, float] = field(default_factory=dict)
    buffer_status: Dict[str, Any] = field(default_factory=dict)
    quality_indicators: Dict[str, float] = field(default_factory=dict)
    vad_result: Optional[Dict[str, Any]] = None
    silence_duration: float = 0.0
    voice_duration: float = 0.0


@dataclass
class WorkflowEventData(MLEventData):
    """Event data for workflow events."""
    workflow_state: str = ""
    workflow_step: str = ""
    step_duration_ms: float = 0.0
    workflow_duration_ms: float = 0.0
    success: bool = True
    error_message: Optional[str] = None
    conversation_session_id: Optional[str] = None


@dataclass
class PerformanceEventData(MLEventData):
    """Event data for performance monitoring events."""
    cpu_usage_percent: float = 0.0
    memory_usage_mb: float = 0.0
    gpu_usage_percent: float = 0.0
    gpu_memory_usage_mb: float = 0.0
    latency_ms: float = 0.0
    throughput_samples_per_second: float = 0.0
    error_rate: float = 0.0
    queue_depth: int = 0


@dataclass
class ModelEventData(MLEventData):
    """Event data for model events."""
    model_type: str = ""
    model_path: str = ""
    model_name: str = ""
    model_version: str = ""
    model_size_mb: float = 0.0
    load_time_ms: float = 0.0
    validation_result: bool = True
    error_details: Optional[str] = None


@dataclass
class NetworkEventData(MLEventData):
    """Event data for network events."""
    command_type: str = ""
    server_response: str = ""
    bytes_transmitted: int = 0
    transmission_time_ms: float = 0.0
    connection_status: str = ""
    error_details: Optional[str] = None


class EventCorrelation:
    """Event correlation and tracking."""
    
    def __init__(self, max_correlation_time: float = 30.0):
        """
        Initialize event correlation.
        
        Args:
            max_correlation_time: Maximum time to keep correlations (seconds)
        """
        self.max_correlation_time = max_correlation_time
        self._correlations: Dict[str, List[MLEventData]] = defaultdict(list)
        self._correlation_times: Dict[str, float] = {}
        self._lock = threading.RLock()
    
    def add_event(self, event_data: MLEventData) -> None:
        """Add event to correlation tracking."""
        if not event_data.correlation_id:
            return
        
        with self._lock:
            current_time = time.time()
            
            # Clean up old correlations
            self._cleanup_old_correlations(current_time)
            
            # Add event to correlation
            self._correlations[event_data.correlation_id].append(event_data)
            self._correlation_times[event_data.correlation_id] = current_time
    
    def get_correlated_events(self, correlation_id: str) -> List[MLEventData]:
        """Get all events with the same correlation ID."""
        with self._lock:
            return self._correlations.get(correlation_id, []).copy()
    
    def _cleanup_old_correlations(self, current_time: float) -> None:
        """Clean up old correlation data."""
        expired_ids = [
            correlation_id for correlation_id, timestamp in self._correlation_times.items()
            if current_time - timestamp > self.max_correlation_time
        ]
        
        for correlation_id in expired_ids:
            del self._correlations[correlation_id]
            del self._correlation_times[correlation_id]


class MLEventManager:
    """
    Comprehensive event manager for ML components.
    
    This class provides centralized event management for all ML components,
    integrating with the main Trixy event system while providing ML-specific
    functionality.
    """
    
    def __init__(
        self,
        main_event_handler: Optional[EventHandler] = None,
        enable_correlation: bool = True,
        enable_persistence: bool = False,
        max_event_history: int = 1000
    ):
        """
        Initialize ML event manager.
        
        Args:
            main_event_handler: Main Trixy event handler
            enable_correlation: Enable event correlation tracking
            enable_persistence: Enable event persistence
            max_event_history: Maximum events to keep in history
        """
        self.main_event_handler = main_event_handler
        self.enable_correlation = enable_correlation
        self.enable_persistence = enable_persistence
        self.max_event_history = max_event_history
        
        # Event processing
        self._event_queue = queue.Queue()
        self._processing_thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()
        self._is_running = False
        
        # Event handlers
        self._ml_event_handlers: Dict[MLEventType, List[Callable]] = defaultdict(list)
        self._registered_components: Set[Any] = set()
        
        # Event history and correlation
        self._event_history = deque(maxlen=max_event_history)
        self._correlation = EventCorrelation() if enable_correlation else None
        
        # Performance tracking
        self._event_counts: Dict[str, int] = defaultdict(int)
        self._last_event_time: Dict[str, float] = {}
        self._processing_times: Dict[str, List[float]] = defaultdict(list)
        
        # Event filtering
        self._event_filters: List[Callable[[MLEventData], bool]] = []
        self._debug_mode = False
        
        logger.info("ML Event Manager initialized")
    
    def start(self) -> None:
        """Start the ML event manager."""
        if self._is_running:
            logger.warning("ML event manager already running")
            return
        
        logger.info("Starting ML event manager...")
        self._is_running = True
        self._stop_event.clear()
        
        # Start event processing thread
        self._processing_thread = threading.Thread(
            target=self._event_processing_loop,
            daemon=True,
            name="MLEventManager"
        )
        self._processing_thread.start()
        
        logger.info("ML event manager started")
    
    def stop(self) -> None:
        """Stop the ML event manager."""
        if not self._is_running:
            return
        
        logger.info("Stopping ML event manager...")
        self._is_running = False
        self._stop_event.set()
        
        # Stop processing thread
        if self._processing_thread and self._processing_thread.is_alive():
            self._processing_thread.join(timeout=5.0)
        
        logger.info("ML event manager stopped")
    
    def register_ml_component(self, component: Any) -> None:
        """
        Register an ML component for event handling.
        
        Args:
            component: ML component to register
        """
        if component in self._registered_components:
            logger.warning(f"Component {component} already registered")
            return
        
        self._registered_components.add(component)
        
        # Auto-register event handlers if the component has them
        self._auto_register_handlers(component)
        
        logger.info(f"Registered ML component: {type(component).__name__}")
    
    def unregister_ml_component(self, component: Any) -> None:
        """
        Unregister an ML component.
        
        Args:
            component: ML component to unregister
        """
        if component in self._registered_components:
            self._registered_components.remove(component)
            logger.info(f"Unregistered ML component: {type(component).__name__}")
    
    def add_event_handler(self, event_type: MLEventType, handler: Callable) -> None:
        """
        Add an event handler for a specific ML event type.
        
        Args:
            event_type: ML event type
            handler: Event handler function
        """
        self._ml_event_handlers[event_type].append(handler)
        logger.debug(f"Added event handler for {event_type.value}")
    
    def remove_event_handler(self, event_type: MLEventType, handler: Callable) -> None:
        """
        Remove an event handler.
        
        Args:
            event_type: ML event type
            handler: Event handler function to remove
        """
        if handler in self._ml_event_handlers[event_type]:
            self._ml_event_handlers[event_type].remove(handler)
            logger.debug(f"Removed event handler for {event_type.value}")
    
    def trigger_event(self, event_data: MLEventData) -> None:
        """
        Trigger an ML event.
        
        Args:
            event_data: Event data to trigger
        """
        if not self._is_running:
            logger.warning("Cannot trigger event - event manager not running")
            return
        
        # Apply event filters
        if not self._should_process_event(event_data):
            return
        
        # Add to processing queue
        try:
            if not self._event_queue.full():
                self._event_queue.put(event_data, block=False)
            else:
                logger.warning("Event queue full - dropping event")
        except Exception as e:
            logger.error(f"Error queuing event: {e}")
    
    def trigger_wakeword_detected(
        self,
        wakeword_type: str,
        confidence: float,
        speaker_id: Optional[str] = None,
        speaker_name: Optional[str] = None,
        speaker_confidence: float = 0.0,
        **kwargs
    ) -> None:
        """Trigger wakeword detected event."""
        event_data = WakewordEventData(
            event_type=MLEventType.WAKEWORD_DETECTED.value,
            wakeword_type=wakeword_type,
            confidence=confidence,
            speaker_id=speaker_id,
            speaker_name=speaker_name,
            speaker_confidence=speaker_confidence,
            **kwargs
        )
        self.trigger_event(event_data)
    
    def trigger_speaker_identified(
        self,
        speaker_id: str,
        speaker_name: str,
        confidence: float,
        similarity_score: float = 0.0,
        **kwargs
    ) -> None:
        """Trigger speaker identified event."""
        event_data = SpeakerEventData(
            event_type=MLEventType.SPEAKER_IDENTIFIED.value,
            speaker_id=speaker_id,
            speaker_name=speaker_name,
            confidence=confidence,
            similarity_score=similarity_score,
            **kwargs
        )
        self.trigger_event(event_data)
    
    def trigger_audio_quality_warning(
        self,
        audio_metrics: Dict[str, float],
        warning_details: str = "",
        **kwargs
    ) -> None:
        """Trigger audio quality warning event."""
        event_data = AudioEventData(
            event_type=MLEventType.AUDIO_QUALITY_WARNING.value,
            audio_metrics=audio_metrics,
            metadata={'warning_details': warning_details},
            priority=EventPriority.HIGH,
            **kwargs
        )
        self.trigger_event(event_data)
    
    def trigger_workflow_completed(
        self,
        workflow_state: str,
        success: bool = True,
        duration_ms: float = 0.0,
        session_id: Optional[str] = None,
        **kwargs
    ) -> None:
        """Trigger workflow completed event."""
        event_data = WorkflowEventData(
            event_type=MLEventType.WORKFLOW_COMPLETED.value,
            workflow_state=workflow_state,
            success=success,
            workflow_duration_ms=duration_ms,
            conversation_session_id=session_id,
            **kwargs
        )
        self.trigger_event(event_data)
    
    def trigger_performance_update(
        self,
        cpu_usage: float = 0.0,
        memory_usage: float = 0.0,
        latency_ms: float = 0.0,
        **kwargs
    ) -> None:
        """Trigger performance metrics update event."""
        event_data = PerformanceEventData(
            event_type=MLEventType.PERFORMANCE_METRICS_UPDATED.value,
            cpu_usage_percent=cpu_usage,
            memory_usage_mb=memory_usage,
            latency_ms=latency_ms,
            **kwargs
        )
        self.trigger_event(event_data)
    
    def add_event_filter(self, filter_func: Callable[[MLEventData], bool]) -> None:
        """
        Add an event filter function.
        
        Args:
            filter_func: Function that returns True if event should be processed
        """
        self._event_filters.append(filter_func)
        logger.debug("Added event filter")
    
    def enable_debug_mode(self, enabled: bool = True) -> None:
        """Enable or disable debug mode."""
        self._debug_mode = enabled
        logger.info(f"Debug mode {'enabled' if enabled else 'disabled'}")
    
    def get_event_statistics(self) -> Dict[str, Any]:
        """Get event processing statistics."""
        total_events = sum(self._event_counts.values())
        
        return {
            'total_events_processed': total_events,
            'events_by_type': dict(self._event_counts),
            'average_processing_times': {
                event_type: sum(times) / len(times) if times else 0.0
                for event_type, times in self._processing_times.items()
            },
            'event_history_size': len(self._event_history),
            'registered_components': len(self._registered_components),
            'active_handlers': sum(len(handlers) for handlers in self._ml_event_handlers.values()),
            'is_running': self._is_running
        }
    
    def get_event_history(self, event_type: Optional[str] = None, limit: int = 100) -> List[MLEventData]:
        """
        Get event history.
        
        Args:
            event_type: Filter by event type (optional)
            limit: Maximum number of events to return
            
        Returns:
            List of historical events
        """
        history = list(self._event_history)
        
        if event_type:
            history = [event for event in history if event.event_type == event_type]
        
        return history[-limit:] if limit > 0 else history
    
    def get_correlated_events(self, correlation_id: str) -> List[MLEventData]:
        """Get events with the same correlation ID."""
        if self._correlation:
            return self._correlation.get_correlated_events(correlation_id)
        return []
    
    def _event_processing_loop(self) -> None:
        """Main event processing loop."""
        logger.debug("ML event processing loop started")
        
        while not self._stop_event.is_set():
            try:
                # Get event from queue
                try:
                    event_data = self._event_queue.get(timeout=0.1)
                except queue.Empty:
                    continue
                
                # Process the event
                self._process_event(event_data)
                
                # Mark task done
                self._event_queue.task_done()
                
            except Exception as e:
                logger.error(f"Error in event processing loop: {e}")
                time.sleep(0.1)
        
        logger.debug("ML event processing loop stopped")
    
    def _process_event(self, event_data: MLEventData) -> None:
        """Process a single event."""
        start_time = time.time()
        
        try:
            # Add to history
            self._event_history.append(event_data)
            
            # Add to correlation tracking
            if self._correlation and event_data.correlation_id:
                self._correlation.add_event(event_data)
            
            # Update statistics
            self._event_counts[event_data.event_type] += 1
            self._last_event_time[event_data.event_type] = start_time
            
            # Debug logging
            if self._debug_mode:
                logger.debug(f"Processing ML event: {event_data.event_type}")
            
            # Call ML-specific handlers
            event_type = MLEventType(event_data.event_type) if hasattr(MLEventType, event_data.event_type.upper()) else None
            if event_type and event_type in self._ml_event_handlers:
                for handler in self._ml_event_handlers[event_type]:
                    try:
                        handler(event_data)
                    except Exception as e:
                        logger.error(f"Error in ML event handler: {e}")
            
            # Forward to main event system if available
            if EVENTS_AVAILABLE and self.main_event_handler:
                self._forward_to_main_event_system(event_data)
            
            # Update processing time statistics
            processing_time = (time.time() - start_time) * 1000
            self._processing_times[event_data.event_type].append(processing_time)
            
            # Keep only recent processing times
            if len(self._processing_times[event_data.event_type]) > 100:
                self._processing_times[event_data.event_type] = self._processing_times[event_data.event_type][-50:]
            
        except Exception as e:
            logger.error(f"Error processing event {event_data.event_type}: {e}")
    
    def _forward_to_main_event_system(self, event_data: MLEventData) -> None:
        """Forward ML event to main Trixy event system."""
        try:
            # Convert ML event to main event system format
            if event_data.event_type == MLEventType.WAKEWORD_DETECTED.value:
                self._forward_wakeword_event(event_data)
            elif event_data.event_type == MLEventType.SPEAKER_IDENTIFIED.value:
                self._forward_speaker_event(event_data)
            elif hasattr(event_data, 'priority') and event_data.priority == EventPriority.CRITICAL:
                # Forward critical events as system events
                self._forward_system_event(event_data)
            
        except Exception as e:
            logger.error(f"Error forwarding event to main system: {e}")
    
    def _forward_wakeword_event(self, event_data: MLEventData) -> None:
        """Forward wakeword event to main event system."""
        if not isinstance(event_data, WakewordEventData):
            return
        
        # Create speaker info if available
        speaker_info = None
        if event_data.speaker_id:
            speaker_info = SpeakerInfo(
                speaker_id=event_data.speaker_id,
                speaker_name=event_data.speaker_name or "Unknown",
                confidence=event_data.speaker_confidence
            )
        
        # Create main event data
        main_event_data = WakewordReceivedEventData(
            wakeword_type=event_data.wakeword_type,
            confidence=event_data.confidence,
            processing_time_ms=event_data.processing_time_ms,
            audio_chunk_position=event_data.audio_chunk_position,
            timestamp=event_data.timestamp,
            speaker_info=speaker_info
        )
        
        # Trigger in main event system
        self.main_event_handler.trigger_event("wakeword_received", main_event_data)
    
    def _forward_speaker_event(self, event_data: MLEventData) -> None:
        """Forward speaker event to main event system."""
        # Create generic event for speaker identification
        self.main_event_handler.trigger_event(
            "speaker_identified",
            event_data.metadata
        )
    
    def _forward_system_event(self, event_data: MLEventData) -> None:
        """Forward critical event as system event."""
        self.main_event_handler.trigger_event(
            "ml_system_event",
            {
                'ml_event_type': event_data.event_type,
                'priority': event_data.priority.value,
                'timestamp': event_data.timestamp,
                'metadata': event_data.metadata
            }
        )
    
    def _should_process_event(self, event_data: MLEventData) -> bool:
        """Check if event should be processed based on filters."""
        for filter_func in self._event_filters:
            try:
                if not filter_func(event_data):
                    return False
            except Exception as e:
                logger.error(f"Error in event filter: {e}")
        
        return True
    
    def _auto_register_handlers(self, component: Any) -> None:
        """Auto-register event handlers from component."""
        # Look for methods with ML event decorators
        for attr_name in dir(component):
            attr = getattr(component, attr_name)
            if callable(attr) and hasattr(attr, '_ml_events'):
                # Register handler for each event type
                for event_type in attr._ml_events:
                    if isinstance(event_type, MLEventType):
                        self.add_event_handler(event_type, attr)
                        logger.debug(f"Auto-registered handler {attr_name} for {event_type.value}")


def MLEvent(event_types: List[MLEventType]) -> Callable:
    """
    Decorator for ML event handlers.
    
    Args:
        event_types: List of ML event types to handle
        
    Returns:
        Decorated function
    """
    def decorator(func: Callable) -> Callable:
        func._ml_events = event_types
        return func
    return decorator


def create_ml_event_manager(
    main_event_handler: Optional[EventHandler] = None,
    **kwargs
) -> MLEventManager:
    """
    Factory function to create an ML event manager.
    
    Args:
        main_event_handler: Main Trixy event handler
        **kwargs: Additional configuration options
        
    Returns:
        Configured ML event manager
    """
    return MLEventManager(main_event_handler=main_event_handler, **kwargs)


# Export main classes and functions
__all__ = [
    'MLEventManager',
    'MLEventType',
    'EventPriority',
    'MLEventData',
    'WakewordEventData',
    'SpeakerEventData',
    'AudioEventData',
    'WorkflowEventData',
    'PerformanceEventData',
    'ModelEventData',
    'NetworkEventData',
    'EventCorrelation',
    'MLEvent',
    'create_ml_event_manager'
]