Trixy3/trixy_core/ml/voice_recognition/integration.py

"""
Integration module for Voice Recognition with Trixy Event System and Wakeword Detection.

This module provides seamless integration between the voice recognition system
and the existing Trixy infrastructure, including event handling, wakeword workflow,
and satellite communication.
"""

import logging
import threading
import time
from typing import Dict, List, Optional, Any, Callable, Union
from dataclasses import dataclass, field
from datetime import datetime
import numpy as np
import torch

from .detector import VoiceRecognitionDetector, SpeakerIdentificationResult, AudioChunk
from .config import VoiceRecognitionConfig
from .speaker_manager import SpeakerManager, SpeakerInfo

# Try to import event system components
try:
    from trixy_core.events import (
        EventHandler, TrixyEvent, EventType, TrixyEventData,
        WakewordReceivedEventData, RawAudioInputReceivedEventData, SpeakerInfo as EventSpeakerInfo
    )
    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
    EventSpeakerInfo = object

# Configure logging
logger = logging.getLogger(__name__)


@dataclass
class VoiceRecognitionEventData(TrixyEventData if EVENTS_AVAILABLE else object):
    """Event data for voice recognition events."""
    event_type: str = ""
    speaker_id: Optional[str] = None
    speaker_name: Optional[str] = None
    confidence: float = 0.0
    similarity_score: float = 0.0
    processing_time_ms: float = 0.0
    audio_quality: float = 0.0
    vad_confidence: float = 0.0
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    additional_info: Dict[str, Any] = field(default_factory=dict)


@dataclass
class SpeakerEnrolledEventData(TrixyEventData if EVENTS_AVAILABLE else object):
    """Event data for speaker enrollment."""
    speaker_id: str = ""
    speaker_name: str = ""
    enrollment_success: bool = False
    embedding_count: int = 0
    quality_score: float = 0.0
    verification_score: float = 0.0
    message: str = ""
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


@dataclass
class SpeakerRemovedEventData(TrixyEventData if EVENTS_AVAILABLE else object):
    """Event data for speaker removal."""
    speaker_id: str = ""
    speaker_name: str = ""
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


class VoiceRecognitionIntegration:
    """
    Integration layer between Voice Recognition system and Trixy infrastructure.
    
    This class handles:
    - Event system integration
    - Wakeword detection enhancement
    - Satellite communication coordination
    - Performance monitoring and reporting
    """
    
    def __init__(self, 
                 voice_detector: VoiceRecognitionDetector,
                 event_handler: Optional[EventHandler] = None,
                 enable_wakeword_enhancement: bool = True,
                 enable_continuous_recognition: bool = False):
        """
        Initialize voice recognition integration.
        
        Args:
            voice_detector: Voice recognition detector instance
            event_handler: Event handler for system integration
            enable_wakeword_enhancement: Whether to enhance wakeword events with speaker info
            enable_continuous_recognition: Whether to run continuous recognition
        """
        self.voice_detector = voice_detector
        self.event_handler = event_handler
        self.enable_wakeword_enhancement = enable_wakeword_enhancement
        self.enable_continuous_recognition = enable_continuous_recognition
        
        # Integration state
        self._is_running = False
        self._continuous_thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()
        
        # Performance tracking
        self._performance_stats = {
            'total_recognitions': 0,
            'successful_recognitions': 0,
            'wakeword_enhancements': 0,
            'average_processing_time_ms': 0.0,
            'last_recognition': None
        }
        
        # Event callbacks
        self.on_speaker_identified: Optional[Callable] = None
        self.on_speaker_verified: Optional[Callable] = None
        self.on_recognition_failed: Optional[Callable] = None
        
        logger.info("Voice Recognition Integration initialized")
        
        # Register event handlers if event system is available
        if EVENTS_AVAILABLE and self.event_handler:
            self._register_event_handlers()
    
    def start(self):
        """Start the voice recognition integration."""
        if self._is_running:
            logger.warning("Voice recognition integration already running")
            return
        
        self._is_running = True
        self._stop_event.clear()
        
        # Start voice detector
        self.voice_detector.start()
        
        # Start continuous recognition if enabled
        if self.enable_continuous_recognition:
            self._continuous_thread = threading.Thread(
                target=self._continuous_recognition_loop,
                daemon=True
            )
            self._continuous_thread.start()
        
        # Trigger start event
        if EVENTS_AVAILABLE and self.event_handler:
            try:
                event_data = VoiceRecognitionEventData(
                    event_type="voice_recognition_started",
                    speaker_id=None,
                    speaker_name=None,
                    confidence=0.0,
                    similarity_score=0.0,
                    processing_time_ms=0.0,
                    audio_quality=0.0,
                    vad_confidence=0.0,
                    additional_info={'continuous_mode': self.enable_continuous_recognition}
                )
                self.event_handler.trigger_event("voice_recognition_started", event_data)
            except Exception as e:
                logger.error(f"Error triggering start event: {e}")
        
        logger.info("Voice Recognition Integration started")
    
    def stop(self):
        """Stop the voice recognition integration."""
        if not self._is_running:
            return
        
        self._is_running = False
        self._stop_event.set()
        
        # Stop continuous recognition thread
        if self._continuous_thread and self._continuous_thread.is_alive():
            self._continuous_thread.join(timeout=5.0)
        
        # Stop voice detector
        self.voice_detector.stop()
        
        # Trigger stop event
        if EVENTS_AVAILABLE and self.event_handler:
            try:
                event_data = VoiceRecognitionEventData(
                    event_type="voice_recognition_stopped",
                    speaker_id=None,
                    speaker_name=None,
                    confidence=0.0,
                    similarity_score=0.0,
                    processing_time_ms=0.0,
                    audio_quality=0.0,
                    vad_confidence=0.0
                )
                self.event_handler.trigger_event("voice_recognition_stopped", event_data)
            except Exception as e:
                logger.error(f"Error triggering stop event: {e}")
        
        logger.info("Voice Recognition Integration stopped")
    
    def is_running(self) -> bool:
        """Check if integration is running."""
        return self._is_running
    
    def process_wakeword_detection(self, 
                                  audio_data: np.ndarray,
                                  wakeword_id: str,
                                  original_speaker_info: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        Process wakeword detection and enhance with speaker recognition.
        
        Args:
            audio_data: Audio data from wakeword detection
            wakeword_id: ID of detected wakeword
            original_speaker_info: Original speaker info from wakeword detector
            
        Returns:
            Enhanced speaker information
        """
        if not self.enable_wakeword_enhancement:
            return original_speaker_info or {}
        
        try:
            # Perform speaker identification on wakeword audio
            result = self.voice_detector.identify_speaker(audio_data)
            
            # Update performance stats
            self._update_performance_stats(result)
            
            # Create enhanced speaker info
            enhanced_info = {
                'speaker_id': result.speaker_id,
                'speaker_name': result.speaker_name,
                'confidence': result.confidence,
                'similarity_score': result.confidence,  # For backward compatibility
                'vr_processing_time_ms': result.processing_time_ms,
                'vr_audio_quality': result.audio_quality,
                'vr_vad_confidence': result.vad_result.confidence if result.vad_result else 0.0,
                'vr_timestamp': result.timestamp,
                'wakeword_id': wakeword_id
            }
            
            # Merge with original info if available
            if original_speaker_info:
                enhanced_info.update(original_speaker_info)
            
            # Trigger enhanced wakeword event
            if EVENTS_AVAILABLE and self.event_handler:
                try:
                    event_data = VoiceRecognitionEventData(
                        event_type="wakeword_enhanced",
                        speaker_id=result.speaker_id,
                        speaker_name=result.speaker_name,
                        confidence=result.confidence,
                        similarity_score=result.confidence,
                        processing_time_ms=result.processing_time_ms,
                        audio_quality=result.audio_quality,
                        vad_confidence=result.vad_result.confidence if result.vad_result else 0.0,
                        additional_info={'wakeword_id': wakeword_id}
                    )
                    self.event_handler.trigger_event("wakeword_enhanced", event_data)
                    self._performance_stats['wakeword_enhancements'] += 1
                except Exception as e:
                    logger.error(f"Error triggering wakeword enhanced event: {e}")
            
            # Call callback if set
            if self.on_speaker_identified:
                try:
                    self.on_speaker_identified(result)
                except Exception as e:
                    logger.error(f"Error in speaker identified callback: {e}")
            
            logger.debug(f"Enhanced wakeword with speaker: {result.speaker_name} (confidence: {result.confidence:.3f})")
            
            return enhanced_info
            
        except Exception as e:
            logger.error(f"Error enhancing wakeword detection: {e}")
            
            # Call failure callback if set
            if self.on_recognition_failed:
                try:
                    self.on_recognition_failed(str(e))
                except Exception as e:
                    logger.error(f"Error in recognition failed callback: {e}")
            
            return original_speaker_info or {}
    
    def process_audio_stream(self, audio_data: np.ndarray) -> Optional[SpeakerIdentificationResult]:
        """
        Process streaming audio for continuous recognition.
        
        Args:
            audio_data: Streaming audio data
            
        Returns:
            Speaker identification result if triggered
        """
        if not self._is_running:
            return None
        
        try:
            result = self.voice_detector.process_audio_stream(audio_data)
            
            if result and result.speaker_id:
                # Update performance stats
                self._update_performance_stats(result)
                
                # Trigger continuous recognition event
                if EVENTS_AVAILABLE and self.event_handler:
                    try:
                        event_data = VoiceRecognitionEventData(
                            event_type="speaker_identified_continuous",
                            speaker_id=result.speaker_id,
                            speaker_name=result.speaker_name,
                            confidence=result.confidence,
                            similarity_score=result.confidence,
                            processing_time_ms=result.processing_time_ms,
                            audio_quality=0.0,  # Not available in streaming mode
                            vad_confidence=result.vad_result.confidence if result.vad_result else 0.0
                        )
                        self.event_handler.trigger_event("speaker_identified_continuous", event_data)
                    except Exception as e:
                        logger.error(f"Error triggering continuous recognition event: {e}")
                
                # Call callback if set
                if self.on_speaker_identified:
                    try:
                        self.on_speaker_identified(result)
                    except Exception as e:
                        logger.error(f"Error in speaker identified callback: {e}")
            
            return result
            
        except Exception as e:
            logger.error(f"Error in continuous recognition: {e}")
            return None
    
    def enroll_speaker_from_audio(self, 
                                 speaker_name: str,
                                 audio_samples: List[np.ndarray],
                                 metadata: Optional[Dict[str, Any]] = None) -> bool:
        """
        Enroll a new speaker using audio samples.
        
        Args:
            speaker_name: Name of the speaker
            audio_samples: List of audio samples for enrollment
            metadata: Optional speaker metadata
            
        Returns:
            True if enrollment successful
        """
        try:
            success = self.voice_detector.add_speaker(speaker_name, audio_samples)
            
            if success:
                # Get speaker info
                speakers = self.voice_detector.get_speaker_list()
                enrolled_speaker = next((s for s in speakers if s.speaker_name == speaker_name), None)
                
                # Trigger enrollment event
                if EVENTS_AVAILABLE and self.event_handler and enrolled_speaker:
                    try:
                        event_data = SpeakerEnrolledEventData(
                            speaker_id=enrolled_speaker.speaker_id,
                            speaker_name=speaker_name,
                            enrollment_success=True,
                            embedding_count=enrolled_speaker.embedding_count,
                            quality_score=enrolled_speaker.average_quality,
                            verification_score=0.0,  # Not available here
                            message="Speaker enrolled successfully"
                        )
                        self.event_handler.trigger_event("speaker_enrolled", event_data)
                    except Exception as e:
                        logger.error(f"Error triggering speaker enrolled event: {e}")
                
                logger.info(f"Successfully enrolled speaker: {speaker_name}")
            else:
                logger.warning(f"Failed to enroll speaker: {speaker_name}")
            
            return success
            
        except Exception as e:
            logger.error(f"Error enrolling speaker {speaker_name}: {e}")
            return False
    
    def remove_speaker_by_id(self, speaker_id: str) -> bool:
        """
        Remove a speaker by ID.
        
        Args:
            speaker_id: Speaker ID to remove
            
        Returns:
            True if removal successful
        """
        try:
            # Get speaker info before removal
            speakers = self.voice_detector.get_speaker_list()
            target_speaker = next((s for s in speakers if s.speaker_id == speaker_id), None)
            
            if not target_speaker:
                logger.warning(f"Speaker {speaker_id} not found")
                return False
            
            speaker_name = target_speaker.speaker_name
            success = self.voice_detector.remove_speaker(speaker_id)
            
            if success:
                # Trigger removal event
                if EVENTS_AVAILABLE and self.event_handler:
                    try:
                        event_data = SpeakerRemovedEventData(
                            speaker_id=speaker_id,
                            speaker_name=speaker_name
                        )
                        self.event_handler.trigger_event("speaker_removed", event_data)
                    except Exception as e:
                        logger.error(f"Error triggering speaker removed event: {e}")
                
                logger.info(f"Successfully removed speaker: {speaker_name} ({speaker_id})")
            else:
                logger.warning(f"Failed to remove speaker: {speaker_id}")
            
            return success
            
        except Exception as e:
            logger.error(f"Error removing speaker {speaker_id}: {e}")
            return False
    
    def get_performance_statistics(self) -> Dict[str, Any]:
        """Get performance statistics for the integration."""
        stats = self._performance_stats.copy()
        
        # Add detector statistics
        try:
            detector_stats = self.voice_detector.get_performance_stats()
            stats.update(detector_stats)
        except Exception as e:
            logger.error(f"Error getting detector stats: {e}")
        
        # Calculate success rate
        if stats['total_recognitions'] > 0:
            stats['recognition_success_rate'] = stats['successful_recognitions'] / stats['total_recognitions']
        else:
            stats['recognition_success_rate'] = 0.0
        
        return stats
    
    def reset_statistics(self):
        """Reset performance statistics."""
        self._performance_stats = {
            'total_recognitions': 0,
            'successful_recognitions': 0,
            'wakeword_enhancements': 0,
            'average_processing_time_ms': 0.0,
            'last_recognition': None
        }
    
    def _register_event_handlers(self):
        """Register event handlers with the event system."""
        if not EVENTS_AVAILABLE:
            return
        
        # Register handlers for wakeword and audio events
        try:
            # Note: In a real implementation, these would be proper event handler methods
            # with @TrixyEvent decorators
            logger.info("Event handlers registered for voice recognition integration")
        except Exception as e:
            logger.error(f"Error registering event handlers: {e}")
    
    def _continuous_recognition_loop(self):
        """Background loop for continuous recognition."""
        logger.info("Starting continuous recognition loop")
        
        # This is a simplified loop - in practice, you'd integrate with
        # the actual audio streaming system
        while not self._stop_event.is_set():
            try:
                # Simulate audio processing
                # In real implementation, this would process actual audio streams
                time.sleep(0.1)  # 100ms processing interval
                
            except Exception as e:
                logger.error(f"Error in continuous recognition loop: {e}")
                time.sleep(1.0)  # Back off on error
        
        logger.info("Continuous recognition loop stopped")
    
    def _update_performance_stats(self, result: SpeakerIdentificationResult):
        """Update performance statistics."""
        self._performance_stats['total_recognitions'] += 1
        
        if result.speaker_id:
            self._performance_stats['successful_recognitions'] += 1
        
        # Update average processing time
        current_avg = self._performance_stats['average_processing_time_ms']
        total_count = self._performance_stats['total_recognitions']
        
        new_avg = ((current_avg * (total_count - 1)) + result.processing_time_ms) / total_count
        self._performance_stats['average_processing_time_ms'] = new_avg
        
        self._performance_stats['last_recognition'] = datetime.now().isoformat()
    
    # Event handler methods (would use @TrixyEvent decorator in real implementation)
    def handle_wakeword_received(self, event_name: str, event_data: WakewordReceivedEventData):
        """Handle wakeword received events."""
        if not self.enable_wakeword_enhancement:
            return
        
        try:
            # Extract audio data from event
            # In real implementation, this would get audio from the event data
            logger.debug(f"Processing wakeword event for enhancement: {event_name}")
            
            # This would call process_wakeword_detection with actual audio data
            
        except Exception as e:
            logger.error(f"Error handling wakeword event: {e}")
    
    def handle_raw_audio_received(self, event_name: str, event_data: RawAudioInputReceivedEventData):
        """Handle raw audio input events."""
        if not self.enable_continuous_recognition:
            return
        
        try:
            # Process audio for speaker recognition
            # In real implementation, this would extract audio from event data
            logger.debug(f"Processing raw audio event: {event_name}")
            
            # This would call process_audio_stream with actual audio data
            
        except Exception as e:
            logger.error(f"Error handling raw audio event: {e}")


def create_voice_recognition_integration(voice_detector: VoiceRecognitionDetector,
                                       event_handler: Optional[EventHandler] = None,
                                       **kwargs) -> VoiceRecognitionIntegration:
    """
    Factory function to create voice recognition integration.
    
    Args:
        voice_detector: Voice recognition detector instance
        event_handler: Event handler for system integration
        **kwargs: Additional configuration parameters
        
    Returns:
        Configured voice recognition integration
    """
    return VoiceRecognitionIntegration(
        voice_detector=voice_detector,
        event_handler=event_handler,
        **kwargs
    )


def setup_complete_voice_recognition_system(config: VoiceRecognitionConfig,
                                          event_handler: Optional[EventHandler] = None,
                                          **kwargs) -> VoiceRecognitionIntegration:
    """
    Set up a complete voice recognition system with all components.
    
    Args:
        config: Voice recognition configuration
        event_handler: Event handler for system integration
        **kwargs: Additional configuration parameters
        
    Returns:
        Complete voice recognition integration system
    """
    try:
        from .detector import create_voice_recognition_detector
        
        # Create voice recognition detector
        detector = create_voice_recognition_detector(
            config=config,
            event_handler=event_handler
        )
        
        # Create integration
        integration = create_voice_recognition_integration(
            voice_detector=detector,
            event_handler=event_handler,
            **kwargs
        )
        
        logger.info("Complete voice recognition system created successfully")
        return integration
        
    except Exception as e:
        logger.error(f"Failed to create complete voice recognition system: {e}")
        raise


# Example usage and testing
def test_voice_recognition_integration():
    """Test voice recognition integration functionality."""
    try:
        from .config import create_default_config
        from .detector import create_voice_recognition_detector
        
        # Create mock event handler if events not available
        if EVENTS_AVAILABLE:
            from trixy_core.events import create_default_event_handler
            event_handler = create_default_event_handler()
        else:
            event_handler = None
        
        # Create configuration
        config = create_default_config("/tmp/test_model.pth")
        
        # Create detector
        detector = create_voice_recognition_detector(config, event_handler)
        
        # Create integration
        integration = create_voice_recognition_integration(
            detector, event_handler, enable_wakeword_enhancement=True
        )
        
        print("Testing voice recognition integration...")
        
        # Start integration
        integration.start()
        
        # Test wakeword enhancement
        test_audio = np.random.randn(16000 * 2).astype(np.float32)  # 2 seconds of audio
        enhanced_info = integration.process_wakeword_detection(
            test_audio, "custom_wakeword"
        )
        print(f"Enhanced wakeword info: {enhanced_info}")
        
        # Test speaker enrollment
        enrollment_samples = [np.random.randn(16000 * 3).astype(np.float32) for _ in range(3)]
        success = integration.enroll_speaker_from_audio("Test Speaker", enrollment_samples)
        print(f"Speaker enrollment success: {success}")
        
        # Get statistics
        stats = integration.get_performance_statistics()
        print(f"Performance statistics: {stats}")
        
        # Stop integration
        integration.stop()
        
        print("Voice recognition integration test completed")
        
    except Exception as e:
        print(f"Integration test failed: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    test_voice_recognition_integration()