Trixy3/plugins/stt_plugin/main.py

"""
Speech-to-Text (STT) Plugin for Trixy Application

This plugin demonstrates STT integration with the Trixy plugin system:
- Audio processing and speech recognition
- Integration with various STT engines (OpenAI Whisper, Google Speech, Azure, etc.)
- Real-time audio stream processing
- Confidence scoring and result filtering
- Language detection and multi-language support
- Audio preprocessing (noise reduction, normalization)
- Batch processing and streaming modes
- Result caching and optimization
- Error handling and fallback mechanisms

This serves as a reference for implementing STT functionality in Trixy.

Features:
- Multiple STT engine support
- Real-time audio processing
- Language detection
- Confidence filtering
- Audio preprocessing
- Result caching
- Performance monitoring
"""

import os
import time
import threading
import queue
import tempfile
import json
from typing import Any, Dict, Optional, List, Tuple, Union
from pathlib import Path
import numpy as np

# Import plugin base class and event decorator
from trixy_core.plugins import TrixyPlugin
from trixy_core.events import TrixyEvent


def pprint(message: str) -> None:
    """STT plugin logging function."""
    print(f"[STT_PLUGIN] {message}")


class AudioFormat:
    """Audio format specification."""
    def __init__(self, sample_rate: int = 16000, channels: int = 1, bit_depth: int = 16):
        self.sample_rate = sample_rate
        self.channels = channels
        self.bit_depth = bit_depth


class STTResult:
    """Speech-to-text recognition result."""
    def __init__(
        self, 
        text: str, 
        confidence: float, 
        language: Optional[str] = None,
        processing_time: float = 0.0,
        engine: str = "unknown",
        alternatives: List[Dict] = None
    ):
        self.text = text
        self.confidence = confidence
        self.language = language
        self.processing_time = processing_time
        self.engine = engine
        self.alternatives = alternatives or []
        self.timestamp = time.time()


class STTPlugin(TrixyPlugin):
    """
    Speech-to-Text plugin for Trixy application.
    
    This plugin provides STT functionality with support for:
    - Multiple STT engines (OpenAI Whisper, Google Speech, etc.)
    - Real-time and batch audio processing
    - Language detection and multi-language support
    - Audio preprocessing and optimization
    - Result caching and confidence filtering
    - Performance monitoring and statistics
    """
    
    def initialize(self) -> None:
        """Initialize the STT plugin."""
        pprint("Initializing STT Plugin...")
        
        # Plugin state
        self._processing_queue = queue.Queue()
        self._processing_thread: Optional[threading.Thread] = None
        self._processing_shutdown = threading.Event()
        self._active_sessions: Dict[str, Dict] = {}
        
        # Statistics
        self._stats = {
            'total_requests': 0,
            'successful_recognitions': 0,
            'failed_recognitions': 0,
            'total_processing_time': 0.0,
            'average_processing_time': 0.0,
            'cache_hits': 0,
            'languages_detected': set(),
            'confidence_distribution': []
        }
        
        # Setup configuration
        self._setup_configuration()
        
        # Initialize STT engines
        self._initialize_engines()
        
        # Setup audio format
        self._audio_format = AudioFormat(
            sample_rate=self.sample_rate,
            channels=self.channels,
            bit_depth=self.bit_depth
        )
        
        # Start processing thread
        self._start_processing_thread()
        
        pprint(f"STT Plugin initialized with engine: {self.primary_engine}")
    
    def _setup_configuration(self) -> None:
        """Setup plugin configuration."""
        # Engine configuration
        self.primary_engine = self.get_config_value("primary_engine", "whisper", str)
        self.fallback_engine = self.get_config_value("fallback_engine", "google", str)
        self.engines_enabled = self.get_config_value("engines_enabled", ["whisper", "google"], list)
        
        # Audio configuration
        self.sample_rate = self.get_config_value("sample_rate", 16000, int)
        self.channels = self.get_config_value("channels", 1, int)
        self.bit_depth = self.get_config_value("bit_depth", 16, int)
        
        # Processing configuration
        self.min_confidence = self.get_config_value("min_confidence", 0.7, float)
        self.max_audio_length = self.get_config_value("max_audio_length", 60.0, float)
        self.enable_preprocessing = self.get_config_value("enable_preprocessing", True, bool)
        self.enable_caching = self.get_config_value("enable_caching", True, bool)
        
        # Language configuration
        self.default_language = self.get_config_value("default_language", "en", str)
        self.auto_detect_language = self.get_config_value("auto_detect_language", True, bool)
        self.supported_languages = self.get_config_value(
            "supported_languages", 
            ["en", "es", "fr", "de", "it", "pt", "ru", "ja", "ko", "zh"], 
            list
        )
        
        # Performance configuration
        self.processing_timeout = self.get_config_value("processing_timeout", 30.0, float)
        self.max_concurrent_requests = self.get_config_value("max_concurrent_requests", 3, int)
        self.enable_streaming = self.get_config_value("enable_streaming", False, bool)
        
        # Validation
        if self.min_confidence < 0 or self.min_confidence > 1:
            pprint("Warning: min_confidence should be between 0 and 1")
            self.min_confidence = max(0, min(1, self.min_confidence))
        
        if self.sample_rate not in [8000, 16000, 22050, 44100, 48000]:
            pprint(f"Warning: Unusual sample rate {self.sample_rate}, may cause issues")
    
    def _initialize_engines(self) -> None:
        """Initialize available STT engines."""
        self._engines = {}
        
        # Initialize Whisper engine
        if "whisper" in self.engines_enabled:
            self._engines["whisper"] = self._init_whisper_engine()
        
        # Initialize Google Speech engine
        if "google" in self.engines_enabled:
            self._engines["google"] = self._init_google_engine()
        
        # Initialize Azure Speech engine
        if "azure" in self.engines_enabled:
            self._engines["azure"] = self._init_azure_engine()
        
        # Initialize custom/local engines
        if "custom" in self.engines_enabled:
            self._engines["custom"] = self._init_custom_engine()
        
        pprint(f"Initialized {len(self._engines)} STT engines: {list(self._engines.keys())}")
    
    def _init_whisper_engine(self) -> Dict[str, Any]:
        """Initialize OpenAI Whisper engine."""
        try:
            # Try to import whisper (placeholder - actual implementation would import real whisper)
            # import whisper
            
            # Whisper configuration
            whisper_config = {
                'model_size': self.get_config_value("whisper_model_size", "base", str),
                'device': self.get_config_value("whisper_device", "cpu", str),
                'compute_type': self.get_config_value("whisper_compute_type", "float32", str),
                'beam_size': self.get_config_value("whisper_beam_size", 5, int),
                'temperature': self.get_config_value("whisper_temperature", 0.0, float),
            }
            
            pprint(f"Whisper engine configured with model: {whisper_config['model_size']}")
            
            return {
                'name': 'whisper',
                'available': True,
                'config': whisper_config,
                'model': None,  # Placeholder - would load actual model
                'supported_formats': ['wav', 'mp3', 'flac', 'm4a'],
                'max_duration': 600  # 10 minutes
            }
        except ImportError:
            pprint("Whisper not available - install openai-whisper package")
            return {'name': 'whisper', 'available': False}
    
    def _init_google_engine(self) -> Dict[str, Any]:
        """Initialize Google Speech-to-Text engine."""
        try:
            # Try to import google cloud speech (placeholder)
            # from google.cloud import speech
            
            google_config = {
                'credentials_path': self.get_config_value("google_credentials_path", "", str),
                'language_code': self.get_config_value("google_language_code", "en-US", str),
                'alternative_language_codes': self.get_config_value(
                    "google_alternative_languages", 
                    ["es-ES", "fr-FR", "de-DE"], 
                    list
                ),
                'enable_automatic_punctuation': self.get_config_value(
                    "google_auto_punctuation", 
                    True, 
                    bool
                ),
                'enable_speaker_diarization': self.get_config_value(
                    "google_speaker_diarization", 
                    False, 
                    bool
                )
            }
            
            pprint("Google Speech engine configured")
            
            return {
                'name': 'google',
                'available': True,
                'config': google_config,
                'client': None,  # Placeholder - would initialize actual client
                'supported_formats': ['wav', 'flac'],
                'max_duration': 480  # 8 minutes
            }
        except ImportError:
            pprint("Google Speech not available - install google-cloud-speech package")
            return {'name': 'google', 'available': False}
    
    def _init_azure_engine(self) -> Dict[str, Any]:
        """Initialize Azure Speech Services engine."""
        try:
            # Try to import azure speech (placeholder)
            # import azure.cognitiveservices.speech as speechsdk
            
            azure_config = {
                'subscription_key': self.get_config_value("azure_subscription_key", "", str),
                'region': self.get_config_value("azure_region", "eastus", str),
                'language': self.get_config_value("azure_language", "en-US", str),
                'endpoint_id': self.get_config_value("azure_endpoint_id", "", str)
            }
            
            pprint("Azure Speech engine configured")
            
            return {
                'name': 'azure',
                'available': True,
                'config': azure_config,
                'recognizer': None,  # Placeholder
                'supported_formats': ['wav', 'flac'],
                'max_duration': 600
            }
        except ImportError:
            pprint("Azure Speech not available - install azure-cognitiveservices-speech package")
            return {'name': 'azure', 'available': False}
    
    def _init_custom_engine(self) -> Dict[str, Any]:
        """Initialize custom/local STT engine."""
        custom_config = {
            'model_path': self.get_config_value("custom_model_path", "", str),
            'vocab_path': self.get_config_value("custom_vocab_path", "", str),
            'use_gpu': self.get_config_value("custom_use_gpu", False, bool)
        }
        
        return {
            'name': 'custom',
            'available': bool(custom_config['model_path']),
            'config': custom_config,
            'model': None,
            'supported_formats': ['wav'],
            'max_duration': 300
        }
    
    def _start_processing_thread(self) -> None:
        """Start the audio processing thread."""
        if self._processing_thread is None or not self._processing_thread.is_alive():
            self._processing_thread = threading.Thread(
                target=self._processing_loop,
                name="STTProcessingThread",
                daemon=True
            )
            self._processing_thread.start()
            pprint("STT processing thread started")
    
    def _processing_loop(self) -> None:
        """Main processing loop for STT requests."""
        while not self._processing_shutdown.wait(0.1):
            try:
                # Get request from queue with timeout
                try:
                    request = self._processing_queue.get(timeout=1.0)
                except queue.Empty:
                    continue
                
                # Process the audio
                self._process_audio_request(request)
                self._processing_queue.task_done()
                
            except Exception as e:
                pprint(f"Error in processing loop: {e}")
    
    def _process_audio_request(self, request: Dict[str, Any]) -> None:
        """Process a single audio STT request."""
        try:
            audio_data = request['audio_data']
            session_id = request.get('session_id', 'unknown')
            source_info = request.get('source_info', {})
            
            pprint(f"Processing STT request for session {session_id}")
            
            # Preprocess audio if enabled
            if self.enable_preprocessing:
                audio_data = self._preprocess_audio(audio_data)
            
            # Check cache if enabled
            if self.enable_caching:
                cached_result = self._check_cache(audio_data)
                if cached_result:
                    self._stats['cache_hits'] += 1
                    self._send_stt_result(cached_result, session_id, source_info)
                    return
            
            # Perform speech recognition
            result = self._recognize_speech(audio_data)
            
            # Cache result if successful and caching enabled
            if result and result.confidence >= self.min_confidence and self.enable_caching:
                self._cache_result(audio_data, result)
            
            # Send result
            if result:
                self._send_stt_result(result, session_id, source_info)
            
            # Update statistics
            self._update_stats(result)
            
        except Exception as e:
            pprint(f"Error processing audio request: {e}")
            self._stats['failed_recognitions'] += 1
    
    def _recognize_speech(self, audio_data: bytes) -> Optional[STTResult]:
        """Recognize speech from audio data."""
        start_time = time.time()
        
        try:
            # Try primary engine first
            if self.primary_engine in self._engines and self._engines[self.primary_engine]['available']:
                result = self._recognize_with_engine(audio_data, self.primary_engine)
                if result and result.confidence >= self.min_confidence:
                    return result
            
            # Fallback to secondary engine
            if (self.fallback_engine in self._engines and 
                self._engines[self.fallback_engine]['available'] and
                self.fallback_engine != self.primary_engine):
                
                pprint(f"Primary engine failed, trying fallback: {self.fallback_engine}")
                result = self._recognize_with_engine(audio_data, self.fallback_engine)
                if result and result.confidence >= self.min_confidence:
                    return result
            
            # Try other available engines
            for engine_name, engine in self._engines.items():
                if (engine['available'] and 
                    engine_name not in [self.primary_engine, self.fallback_engine]):
                    
                    pprint(f"Trying additional engine: {engine_name}")
                    result = self._recognize_with_engine(audio_data, engine_name)
                    if result and result.confidence >= self.min_confidence:
                        return result
            
            return None
            
        except Exception as e:
            pprint(f"Speech recognition failed: {e}")
            return None
        finally:
            processing_time = time.time() - start_time
            pprint(f"Speech recognition completed in {processing_time:.2f}s")
    
    def _recognize_with_engine(self, audio_data: bytes, engine_name: str) -> Optional[STTResult]:
        """Recognize speech using a specific engine."""
        if engine_name == "whisper":
            return self._recognize_whisper(audio_data)
        elif engine_name == "google":
            return self._recognize_google(audio_data)
        elif engine_name == "azure":
            return self._recognize_azure(audio_data)
        elif engine_name == "custom":
            return self._recognize_custom(audio_data)
        else:
            pprint(f"Unknown engine: {engine_name}")
            return None
    
    def _recognize_whisper(self, audio_data: bytes) -> Optional[STTResult]:
        """Recognize speech using Whisper engine."""
        try:
            # Placeholder implementation - actual Whisper integration
            pprint("Processing with Whisper engine...")
            
            # Simulate processing delay
            time.sleep(0.5)
            
            # Simulate result (actual implementation would call Whisper API)
            text = "This is a simulated Whisper transcription"
            confidence = 0.85
            language = "en"
            
            return STTResult(
                text=text,
                confidence=confidence,
                language=language,
                engine="whisper",
                processing_time=0.5
            )
        except Exception as e:
            pprint(f"Whisper recognition failed: {e}")
            return None
    
    def _recognize_google(self, audio_data: bytes) -> Optional[STTResult]:
        """Recognize speech using Google Speech engine."""
        try:
            pprint("Processing with Google Speech engine...")
            
            # Simulate processing
            time.sleep(0.3)
            
            # Simulate result
            text = "This is a simulated Google Speech transcription"
            confidence = 0.92
            
            return STTResult(
                text=text,
                confidence=confidence,
                language="en-US",
                engine="google",
                processing_time=0.3
            )
        except Exception as e:
            pprint(f"Google Speech recognition failed: {e}")
            return None
    
    def _recognize_azure(self, audio_data: bytes) -> Optional[STTResult]:
        """Recognize speech using Azure Speech engine."""
        try:
            pprint("Processing with Azure Speech engine...")
            
            # Simulate processing
            time.sleep(0.4)
            
            # Simulate result
            text = "This is a simulated Azure Speech transcription"
            confidence = 0.88
            
            return STTResult(
                text=text,
                confidence=confidence,
                language="en-US",
                engine="azure",
                processing_time=0.4
            )
        except Exception as e:
            pprint(f"Azure Speech recognition failed: {e}")
            return None
    
    def _recognize_custom(self, audio_data: bytes) -> Optional[STTResult]:
        """Recognize speech using custom engine."""
        try:
            pprint("Processing with custom engine...")
            
            # Simulate processing
            time.sleep(0.6)
            
            # Simulate result
            text = "This is a simulated custom engine transcription"
            confidence = 0.78
            
            return STTResult(
                text=text,
                confidence=confidence,
                language=self.default_language,
                engine="custom",
                processing_time=0.6
            )
        except Exception as e:
            pprint(f"Custom engine recognition failed: {e}")
            return None
    
    def _preprocess_audio(self, audio_data: bytes) -> bytes:
        """Preprocess audio for better recognition."""
        try:
            pprint("Preprocessing audio...")
            
            # Placeholder - actual implementation would:
            # - Normalize volume
            # - Apply noise reduction
            # - Convert sample rate if needed
            # - Apply audio filters
            
            # For now, return original data
            return audio_data
        except Exception as e:
            pprint(f"Audio preprocessing failed: {e}")
            return audio_data
    
    def _check_cache(self, audio_data: bytes) -> Optional[STTResult]:
        """Check if audio result is cached."""
        # Placeholder - actual implementation would use audio hash
        return None
    
    def _cache_result(self, audio_data: bytes, result: STTResult) -> None:
        """Cache STT result."""
        # Placeholder - actual implementation would cache to disk/memory
        pass
    
    def _send_stt_result(self, result: STTResult, session_id: str, source_info: Dict) -> None:
        """Send STT result via event system."""
        try:
            event_handler = self.application.get_event_handler()
            if event_handler:
                event_data = {
                    'conversation_id': session_id,
                    'text': result.text,
                    'confidence': result.confidence,
                    'language': result.language,
                    'engine': result.engine,
                    'processing_time': result.processing_time,
                    'timestamp': result.timestamp,
                    'speaker_info': source_info,
                    'alternatives': result.alternatives
                }
                
                event_handler.trigger_event("text_received", event_data)
                pprint(f"STT result sent: '{result.text}' (confidence: {result.confidence:.2f})")
        except Exception as e:
            pprint(f"Error sending STT result: {e}")
    
    def _update_stats(self, result: Optional[STTResult]) -> None:
        """Update processing statistics."""
        self._stats['total_requests'] += 1
        
        if result:
            self._stats['successful_recognitions'] += 1
            self._stats['total_processing_time'] += result.processing_time
            self._stats['average_processing_time'] = (
                self._stats['total_processing_time'] / self._stats['successful_recognitions']
            )
            self._stats['confidence_distribution'].append(result.confidence)
            
            if result.language:
                self._stats['languages_detected'].add(result.language)
        else:
            self._stats['failed_recognitions'] += 1
    
    # Event Handlers
    
    @TrixyEvent("raw_audio_input_received")
    def on_audio_input(self, event_name: str, event_data: Any) -> None:
        """Handle raw audio input for STT processing."""
        if not self.is_enabled():
            return
        
        try:
            conversation_id = event_data.get('conversation_id', 'unknown')
            audio_data = event_data.get('audio_data')
            speaker_info = event_data.get('speaker_info', {})
            
            if not audio_data:
                pprint("No audio data received, skipping STT processing")
                return
            
            pprint(f"Received audio input for conversation {conversation_id}")
            
            # Create processing request
            request = {
                'audio_data': audio_data,
                'session_id': conversation_id,
                'source_info': speaker_info,
                'timestamp': time.time()
            }
            
            # Queue for processing
            self._processing_queue.put(request)
            pprint(f"Audio queued for STT processing (queue size: {self._processing_queue.qsize()})")
            
        except Exception as e:
            pprint(f"Error handling audio input: {e}")
    
    @TrixyEvent("system_startup")
    def on_system_startup(self, event_name: str, event_data: Any) -> None:
        """Handle system startup."""
        if not self.is_enabled():
            return
        
        pprint("STT Plugin system startup")
        
        # Verify engines
        available_engines = [name for name, engine in self._engines.items() if engine['available']]
        if not available_engines:
            pprint("WARNING: No STT engines available!")
        else:
            pprint(f"Available STT engines: {available_engines}")
    
    # Lifecycle Hooks
    
    def on_enable(self) -> None:
        """Called when plugin is enabled."""
        pprint("STT Plugin enabled")
        self._start_processing_thread()
    
    def on_disable(self) -> None:
        """Called when plugin is disabled."""
        pprint("STT Plugin disabled")
        # Clear processing queue
        while not self._processing_queue.empty():
            try:
                self._processing_queue.get_nowait()
            except queue.Empty:
                break
    
    def cleanup(self) -> None:
        """Clean up STT plugin resources."""
        pprint("STT Plugin cleanup")
        
        # Stop processing thread
        self._processing_shutdown.set()
        if self._processing_thread and self._processing_thread.is_alive():
            self._processing_thread.join(timeout=5.0)
        
        # Clear queue
        while not self._processing_queue.empty():
            try:
                self._processing_queue.get_nowait()
            except queue.Empty:
                break
        
        # Save final statistics
        self.set_config_value("final_stats", self._stats.copy())
    
    # Plugin-specific Methods
    
    def get_stt_stats(self) -> Dict[str, Any]:
        """Get STT processing statistics."""
        stats = self._stats.copy()
        
        # Convert set to list for JSON serialization
        stats['languages_detected'] = list(stats['languages_detected'])
        
        # Add current status
        stats.update({
            'queue_size': self._processing_queue.qsize(),
            'active_sessions': len(self._active_sessions),
            'available_engines': [name for name, engine in self._engines.items() if engine['available']],
            'primary_engine': self.primary_engine,
            'fallback_engine': self.fallback_engine
        })
        
        return stats
    
    def test_engine(self, engine_name: str) -> Dict[str, Any]:
        """Test a specific STT engine."""
        if engine_name not in self._engines:
            return {'success': False, 'error': f'Engine {engine_name} not found'}
        
        if not self._engines[engine_name]['available']:
            return {'success': False, 'error': f'Engine {engine_name} not available'}
        
        try:
            # Generate test audio (silence)
            test_audio = b'\x00' * (self.sample_rate * 2)  # 1 second of silence
            
            result = self._recognize_with_engine(test_audio, engine_name)
            
            return {
                'success': True,
                'engine': engine_name,
                'result': {
                    'text': result.text if result else None,
                    'confidence': result.confidence if result else 0,
                    'processing_time': result.processing_time if result else 0
                }
            }
        except Exception as e:
            return {'success': False, 'error': str(e)}