websocket_handler.py

"""
WebSocket Handler for Real-time STT/TTS with Barge-in Support
"""
from fastapi import WebSocket, WebSocketDisconnect
from typing import Dict, Any, Optional
import json
import asyncio
import base64
from datetime import datetime
from collections import deque
from enum import Enum
import numpy as np
import traceback

from session import Session, session_store
from config_provider import ConfigProvider
from chat_handler import handle_new_message, handle_parameter_followup
from stt_factory import STTFactory
from tts_factory import TTSFactory
from logger import log_info, log_error, log_debug, log_warning

# ========================= CONSTANTS =========================
# Default values - will be overridden by config
DEFAULT_SILENCE_THRESHOLD_MS = 2000
DEFAULT_AUDIO_CHUNK_SIZE = 4096
DEFAULT_ENERGY_THRESHOLD = 0.01
DEFAULT_AUDIO_BUFFER_MAX_SIZE = 1000

# ========================= ENUMS =========================
class ConversationState(Enum):
    IDLE = "idle"
    LISTENING = "listening"
    PROCESSING_STT = "processing_stt"
    PROCESSING_LLM = "processing_llm"
    PROCESSING_TTS = "processing_tts"
    PLAYING_AUDIO = "playing_audio"

# ========================= CLASSES =========================
class AudioBuffer:
    """Thread-safe circular buffer for audio chunks"""
    def __init__(self, max_size: int = DEFAULT_AUDIO_BUFFER_MAX_SIZE):
        self.buffer = deque(maxlen=max_size)
        self.lock = asyncio.Lock()
        
    async def add_chunk(self, chunk_data: str):
        """Add base64 encoded audio chunk"""
        async with self.lock:
            decoded = base64.b64decode(chunk_data)
            self.buffer.append(decoded)
            
    async def get_all_audio(self) -> bytes:
        """Get all audio data concatenated"""
        async with self.lock:
            return b''.join(self.buffer)
    
    async def clear(self):
        """Clear buffer"""
        async with self.lock:
            self.buffer.clear()
    
    def size(self) -> int:
        """Get current buffer size"""
        return len(self.buffer)


class SilenceDetector:
    """Detect silence in audio stream"""
    def __init__(self, threshold_ms: int = DEFAULT_SILENCE_THRESHOLD_MS, energy_threshold: float = DEFAULT_ENERGY_THRESHOLD):
        self.threshold_ms = threshold_ms
        self.energy_threshold = energy_threshold
        self.silence_start = None
        self.sample_rate = 16000
        
    def update(self, audio_chunk: bytes) -> int:
        """Update with new audio chunk and return silence duration in ms"""
        if self.is_silence(audio_chunk):
            if self.silence_start is None:
                self.silence_start = datetime.now()
            silence_duration = (datetime.now() - self.silence_start).total_seconds() * 1000
            return int(silence_duration)
        else:
            self.silence_start = None
            return 0
    
    def is_silence(self, audio_chunk: bytes) -> bool:
        """Check if audio chunk is silence"""
        try:
            # Convert bytes to numpy array (assuming 16-bit PCM)
            audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
            
            # Calculate RMS energy
            if len(audio_data) == 0:
                return True
                
            rms = np.sqrt(np.mean(audio_data.astype(float) ** 2))
            normalized_rms = rms / 32768.0  # Normalize for 16-bit audio
            
            return normalized_rms < self.energy_threshold
        except Exception as e:
            log_warning(f"Silence detection error: {e}")
            return False
    
    def reset(self):
        """Reset silence detection"""
        self.silence_start = None


class BargeInHandler:
    """Handle user interruptions during TTS playback"""
    def __init__(self):
        self.active_tts_task: Optional[asyncio.Task] = None
        self.is_interrupting = False
        self.lock = asyncio.Lock()
    
    async def start_tts_task(self, coro):
        """Start a cancellable TTS task"""
        async with self.lock:
            # Cancel any existing task
            if self.active_tts_task and not self.active_tts_task.done():
                self.active_tts_task.cancel()
                try:
                    await self.active_tts_task
                except asyncio.CancelledError:
                    pass
            
            # Start new task
            self.active_tts_task = asyncio.create_task(coro)
            return self.active_tts_task
    
    async def handle_interruption(self, current_state: ConversationState):
        """Handle barge-in interruption"""
        async with self.lock:
            self.is_interrupting = True
            
            # Cancel TTS if active
            if self.active_tts_task and not self.active_tts_task.done():
                log_info("Barge-in: Cancelling active TTS")
                self.active_tts_task.cancel()
                try:
                    await self.active_tts_task
                except asyncio.CancelledError:
                    pass
            
            # Reset flag after short delay
            await asyncio.sleep(0.5)
            self.is_interrupting = False


class RealtimeSession:
    """Manage a real-time conversation session"""
    def __init__(self, session: Session):
        self.session = session
        self.state = ConversationState.IDLE
        
        # Get settings from config
        config = ConfigProvider.get().global_config.stt_provider.settings
        
        # Initialize with config values or defaults
        silence_threshold = config.get("speech_timeout_ms", DEFAULT_SILENCE_THRESHOLD_MS)
        energy_threshold = config.get("energy_threshold", DEFAULT_ENERGY_THRESHOLD)
        buffer_max_size = config.get("audio_buffer_max_size", DEFAULT_AUDIO_BUFFER_MAX_SIZE)
        
        self.audio_buffer = AudioBuffer(max_size=buffer_max_size)
        self.silence_detector = SilenceDetector(
            threshold_ms=silence_threshold,
            energy_threshold=energy_threshold
        )
        self.barge_in_handler = BargeInHandler()
        self.stt_manager = None
        self.current_transcription = ""
        self.is_streaming = False
        self.lock = asyncio.Lock()
        
        # Store config for later use
        self.audio_chunk_size = config.get("audio_chunk_size", DEFAULT_AUDIO_CHUNK_SIZE)
        self.silence_threshold_ms = silence_threshold
        
    async def initialize_stt(self):
        """Initialize STT provider"""
        try:
            self.stt_manager = STTFactory.create_provider()
            if self.stt_manager and self.stt_manager.supports_realtime():
                config = ConfigProvider.get().global_config.stt_provider.settings
                
                # STTConfig nesnesi oluştur, dict değil!
                from stt_interface import STTConfig
                stt_config = STTConfig(
                    language=config.get("language", "tr-TR"),
                    interim_results=config.get("interim_results", True),
                    single_utterance=False,
                    enable_punctuation=config.get("enable_punctuation", True),
                    sample_rate=16000,
                    encoding="WEBM_OPUS",
                    model=config.get("model", "latest_long"),
                    use_enhanced=config.get("use_enhanced", True),
                    # Voice Activity Detection
                    vad_enabled=True,
                    speech_timeout_ms=config.get("speech_timeout_ms", 2000),
                    # Noise reduction
                    noise_reduction_enabled=True,
                    noise_reduction_level=config.get("noise_reduction_level", 2)
                )
                
                await self.stt_manager.start_streaming(stt_config)
                log_info("STT manager initialized", session_id=self.session.session_id)
                return True
            else:
                log_warning("STT provider does not support realtime", session_id=self.session.session_id)
                return False
        except Exception as e:
            log_error(f"Failed to initialize STT", error=str(e), session_id=self.session.session_id)
            # STT başarısız oldu, manager'ı None yap
            self.stt_manager = None
            return False
    
    async def change_state(self, new_state: ConversationState):
        """Change conversation state"""
        async with self.lock:
            old_state = self.state
            self.state = new_state
            log_debug(
                f"State change: {old_state.value} → {new_state.value}",
                session_id=self.session.session_id
            )
    
    async def handle_barge_in(self):
        """Handle user interruption"""
        await self.barge_in_handler.handle_interruption(self.state)
        await self.change_state(ConversationState.LISTENING)
    
    async def reset_for_new_utterance(self):
        """Reset for new user utterance"""
        await self.audio_buffer.clear()
        self.silence_detector.reset()
        self.current_transcription = ""
    
    async def cleanup(self):
        """Clean up resources"""
        try:
            if self.stt_manager:
                await self.stt_manager.stop_streaming()
            log_info(f"Cleaned up realtime session", session_id=self.session.session_id)
        except Exception as e:
            log_warning(f"Cleanup error", error=str(e), session_id=self.session.session_id)


# ========================= MAIN HANDLER =========================
async def websocket_endpoint(websocket: WebSocket, session_id: str):
    """Main WebSocket endpoint for real-time conversation"""
    await websocket.accept()
    log_info(f"WebSocket connected", session_id=session_id)
    
    # Get session
    session = session_store.get_session(session_id)
    if not session:
        await websocket.send_json({
            "type": "error",
            "message": "Session not found"
        })
        await websocket.close()
        return
    
    # Mark as realtime session
    session.is_realtime = True
    session_store.update_session(session)
    
    # Initialize conversation
    realtime_session = RealtimeSession(session)
    
    # Initialize STT
    stt_initialized = await realtime_session.initialize_stt()
    if not stt_initialized:
        await websocket.send_json({
            "type": "error",
            "message": "Speech-to-Text service initialization failed. Please check your configuration.",
            "error_type": "stt_init_failed"
        })
        await realtime_session.cleanup()
        await websocket.close()
        return
    
    # Send welcome message if exists
    if session.chat_history and len(session.chat_history) > 0:
        first_message = session.chat_history[0]
        if first_message.get('role') == 'assistant':
            # Send welcome message
            await websocket.send_json({
                "type": "assistant_response",
                "text": first_message.get('content', ''),
                "is_welcome": True
            })
            
            # Generate TTS for welcome message
            tts_provider = TTSFactory.create_provider()
            if tts_provider:
                try:
                    welcome_text = first_message.get('content', '')
                    audio_data = await tts_provider.synthesize(welcome_text)
                    
                    # Send audio in chunks
                    chunk_size = realtime_session.audio_chunk_size
                    total_chunks = (len(audio_data) + chunk_size - 1) // chunk_size
                    
                    for i in range(0, len(audio_data), chunk_size):
                        chunk = audio_data[i:i + chunk_size]
                        chunk_index = i // chunk_size
                        
                        await websocket.send_json({
                            "type": "tts_audio",
                            "data": base64.b64encode(chunk).decode('utf-8'),
                            "chunk_index": chunk_index,
                            "total_chunks": total_chunks,
                            "is_last": chunk_index == total_chunks - 1,
                            "is_welcome": True
                        })
                        
                        await asyncio.sleep(0.01)
                    
                    log_info(f"✅ Welcome message TTS sent", session_id=session_id)
                except Exception as e:
                    log_error(f"Failed to generate welcome TTS", error=str(e))

    # STT başarılı, devam et
    try:
        while True:
            # Receive message
            message = await websocket.receive_json()
            message_type = message.get("type")
            
            if message_type == "audio_chunk":
                await handle_audio_chunk(websocket, realtime_session, message)
                
            elif message_type == "control":
                await handle_control_message(websocket, realtime_session, message)
                
            elif message_type == "ping":
                # Keep-alive ping
                await websocket.send_json({"type": "pong"})
                
    except WebSocketDisconnect:
        log_info(f"WebSocket disconnected", session_id=session_id)
    except Exception as e:
        log_error(
            f"WebSocket error",
            error=str(e),
            traceback=traceback.format_exc(),
            session_id=session_id
        )
        await websocket.send_json({
            "type": "error",
            "message": str(e)
        })
    finally:
        await realtime_session.cleanup()


# ========================= MESSAGE HANDLERS =========================
async def handle_audio_chunk(websocket: WebSocket, session: RealtimeSession, message: Dict[str, Any]):
    """Handle incoming audio chunk with barge-in support"""
    try:
        audio_data = message.get("data")
        if not audio_data:
            return

        # STT manager kontrolü
        if not session.stt_manager:
            log_warning("No STT manager available, ignoring audio chunk", session_id=session.session.session_id)
            await websocket.send_json({
                "type": "error",
                "message": "Speech recognition not available",
                "error_type": "stt_unavailable"
            })
            return
            
        # Check for barge-in during TTS/audio playback
        if session.state in [ConversationState.PLAYING_AUDIO, ConversationState.PROCESSING_TTS]:
            await session.handle_barge_in()
            await websocket.send_json({
                "type": "control",
                "action": "stop_playback"
            })
            log_info(f"Barge-in detected", session_id=session.session.session_id, state=session.state.value)
        
        # Change state to listening if idle
        if session.state == ConversationState.IDLE:
            await session.change_state(ConversationState.LISTENING)
            await websocket.send_json({
                "type": "state_change",
                "from": "idle",
                "to": "listening"
            })
        
        # Add to buffer - don't lose any audio
        await session.audio_buffer.add_chunk(audio_data)
        
        # Decode for processing
        decoded_audio = base64.b64decode(audio_data)
        
        # Check silence
        silence_duration = session.silence_detector.update(decoded_audio)
        
        # Stream to STT if available
        if session.stt_manager and session.state == ConversationState.LISTENING:
            async for result in session.stt_manager.stream_audio(decoded_audio):
                # Send transcription updates
                await websocket.send_json({
                    "type": "transcription",
                    "text": result.text,
                    "is_final": result.is_final,
                    "confidence": result.confidence
                })
                
                if result.is_final:
                    session.current_transcription = result.text
        
        # Process if silence detected and we have transcription
        if silence_duration > session.silence_threshold_ms and session.current_transcription:
            log_info(
                f"User stopped speaking",
                session_id=session.session.session_id,
                silence_ms=silence_duration,
                text=session.current_transcription
            )
            await process_user_input(websocket, session)
            
    except Exception as e:
        log_error(
            f"Audio chunk handling error",
            error=str(e),
            traceback=traceback.format_exc(),
            session_id=session.session.session_id
        )
        await websocket.send_json({
            "type": "error",
            "message": f"Audio processing error: {str(e)}"
        })


async def handle_control_message(websocket: WebSocket, session: RealtimeSession, message: Dict[str, Any]):
    """Handle control messages"""
    action = message.get("action")
    config = message.get("config", {})
    
    log_debug(f"Control message", action=action, session_id=session.session.session_id)
    
    if action == "start_session":
        # Session configuration
        await websocket.send_json({
            "type": "session_started",
            "session_id": session.session.session_id,
            "config": {
                "silence_threshold_ms": session.silence_threshold_ms,
                "audio_chunk_size": session.audio_chunk_size,
                "supports_barge_in": True
            }
        })
        
        # Send welcome message and TTS if available
        chat_history = session.session.chat_history
        if chat_history and len(chat_history) > 0:
            last_message = chat_history[-1]
            if last_message["role"] == "assistant":
                welcome_text = last_message["content"]
                
                # Send text message
                await websocket.send_json({
                    "type": "assistant_response",
                    "text": welcome_text
                })
                
                # Generate TTS if enabled
                tts_provider = TTSFactory.create_provider()
                if tts_provider:
                    await session.change_state(ConversationState.PROCESSING_TTS)
                    await websocket.send_json({
                        "type": "state_change",
                        "from": "idle",
                        "to": "processing_tts"
                    })
                    
                    # Generate and stream TTS
                    tts_task = session.barge_in_handler.start_tts_task(
                        generate_and_stream_tts(websocket, session, tts_provider, welcome_text)
                    )
                    
                    try:
                        await tts_task
                    except asyncio.CancelledError:
                        log_info("Welcome TTS cancelled", session_id=session.session.session_id)


# ========================= PROCESSING FUNCTIONS =========================
async def process_user_input(websocket: WebSocket, session: RealtimeSession):
    """Process complete user input"""
    try:
        user_text = session.current_transcription
        if not user_text:
            await session.reset_for_new_utterance()
            await session.change_state(ConversationState.IDLE)
            return
        
        log_info(f"Processing user input", text=user_text, session_id=session.session.session_id)
        
        # State: STT Processing
        await session.change_state(ConversationState.PROCESSING_STT)
        await websocket.send_json({
            "type": "state_change",
            "from": "listening",
            "to": "processing_stt"
        })
        
        # Send final transcription
        await websocket.send_json({
            "type": "transcription",
            "text": user_text,
            "is_final": True,
            "confidence": 0.95
        })
        
        # State: LLM Processing
        await session.change_state(ConversationState.PROCESSING_LLM)
        await websocket.send_json({
            "type": "state_change",
            "from": "processing_stt",
            "to": "processing_llm"
        })
        
        # Add to chat history
        session.session.add_message("user", user_text)
        
        # Get LLM response based on session state
        if session.session.state == "collect_params":
            response_text = await handle_parameter_followup(session.session, user_text)
        else:
            response_text = await handle_new_message(session.session, user_text)
        
        # Add response to history
        session.session.add_message("assistant", response_text)
        
        # Send text response
        await websocket.send_json({
            "type": "assistant_response",
            "text": response_text
        })
        
        # Generate TTS if enabled
        tts_provider = TTSFactory.create_provider()
        if tts_provider:
            await session.change_state(ConversationState.PROCESSING_TTS)
            await websocket.send_json({
                "type": "state_change",
                "from": "processing_llm",
                "to": "processing_tts"
            })
            
            # Generate TTS with barge-in support
            tts_task = session.barge_in_handler.start_tts_task(
                generate_and_stream_tts(websocket, session, tts_provider, response_text)
            )
            
            try:
                await tts_task
            except asyncio.CancelledError:
                log_info("TTS cancelled due to barge-in", session_id=session.session.session_id)
        else:
            # No TTS, go back to idle
            await session.change_state(ConversationState.IDLE)
            await websocket.send_json({
                "type": "state_change",
                "from": "processing_llm",
                "to": "idle"
            })
        
        # Reset for next input
        await session.reset_for_new_utterance()
        
    except Exception as e:
        log_error(
            f"Error processing user input",
            error=str(e),
            traceback=traceback.format_exc(),
            session_id=session.session.session_id
        )
        await websocket.send_json({
            "type": "error",
            "message": f"Processing error: {str(e)}"
        })
        await session.reset_for_new_utterance()
        await session.change_state(ConversationState.IDLE)


async def generate_and_stream_tts(
    websocket: WebSocket,
    session: RealtimeSession,
    tts_provider,
    text: str
):
    """Generate and stream TTS audio with cancellation support"""
    try:
        # Generate audio
        audio_data = await tts_provider.synthesize(text)
        
        # Change state to playing
        await session.change_state(ConversationState.PLAYING_AUDIO)
        await websocket.send_json({
            "type": "state_change",
            "from": "processing_tts",
            "to": "playing_audio"
        })
        
        # Convert entire audio to base64 for transmission
        import base64
        audio_base64 = base64.b64encode(audio_data).decode('utf-8')
        
        # Stream audio in chunks
        chunk_size = 16384  # Larger chunk size for base64
        total_length = len(audio_base64)
        total_chunks = (total_length + chunk_size - 1) // chunk_size
        
        log_info(f"Streaming TTS audio: {len(audio_data)} bytes as {total_length} base64 chars in {total_chunks} chunks")
        
        for i in range(0, total_length, chunk_size):
            # Check for cancellation
            if asyncio.current_task().cancelled():
                break
                
            chunk = audio_base64[i:i + chunk_size]
            chunk_index = i // chunk_size
            
            await websocket.send_json({
                "type": "tts_audio",
                "data": chunk,
                "chunk_index": chunk_index,
                "total_chunks": total_chunks,
                "is_last": chunk_index == total_chunks - 1,
                "mime_type": "audio/mpeg"
            })
            
            # Small delay to prevent overwhelming the client
            await asyncio.sleep(0.01)
        
        log_info(
            f"TTS streaming completed",
            session_id=session.session.session_id,
            text_length=len(text),
            audio_size=len(audio_data)
        )
        
    except asyncio.CancelledError:
        log_info("TTS streaming cancelled", session_id=session.session.session_id)
        raise
    except Exception as e:
        log_error(
            f"TTS generation error",
            error=str(e),
            session_id=session.session.session_id
        )
        await websocket.send_json({
            "type": "error",
            "message": f"TTS error: {str(e)}"
        })