audiobook/backend/services/unified_tts_service.py

"""
Unified TTS service that supports multiple TTS engines (Kokoro, Piper).
"""
import logging
import os
from typing import Dict, List, Tuple, Optional
from enum import Enum

logger = logging.getLogger(__name__)


class TTSEngine(str, Enum):
    """Available TTS engines."""
    KOKORO = "kokoro"
    PIPER = "piper"


class UnifiedTTSService:
    """
    Unified service that delegates to specific TTS engines.
    Allows runtime selection between Kokoro and Piper TTS.
    """
    
    def __init__(self, engine: Optional[str] = None, ollama_service=None):
        """
        Initialize unified TTS service.
        
        Args:
            engine: TTS engine to use ('kokoro' or 'piper'). 
                   Defaults to TTS_ENGINE env var or 'kokoro'
            ollama_service: Optional OllamaService for AI segmentation
        """
        self.engine_name = engine or os.getenv("TTS_ENGINE", "kokoro")
        self.ollama_service = ollama_service
        
        # Lazy load engines
        self._kokoro_service = None
        self._piper_service = None
        
        logger.info(f"UnifiedTTSService initialized with engine: {self.engine_name}")
    
    @property
    def kokoro_service(self):
        """Lazy load Kokoro service."""
        if self._kokoro_service is None:
            from .tts_service import TTSService
            self._kokoro_service = TTSService(self.ollama_service)
        return self._kokoro_service
    
    @property
    def piper_service(self):
        """Lazy load Piper service."""
        if self._piper_service is None:
            from .piper_tts_service import PiperTTSService
            self._piper_service = PiperTTSService()
        return self._piper_service
    
    def set_engine(self, engine: str):
        """
        Change the active TTS engine.
        
        Args:
            engine: Engine name ('kokoro' or 'piper')
        """
        if engine not in [e.value for e in TTSEngine]:
            raise ValueError(f"Unknown engine: {engine}. Use 'kokoro' or 'piper'")
        
        self.engine_name = engine
        logger.info(f"Switched TTS engine to: {engine}")
    
    def get_available_engines(self) -> List[str]:
        """
        Get list of available TTS engines.
        
        Returns:
            List of engine names
        """
        return [e.value for e in TTSEngine]
    
    def get_available_voices(self) -> Dict[str, List[str]]:
        """
        Get available voices for all engines.
        
        Returns:
            Dict mapping engine names to lists of voice IDs
        """
        voices = {}
        
        # Kokoro voices (from documentation)
        voices["kokoro"] = [
            "af_bella", "af_nicole", "af_sarah", "af_sky",
            "am_adam", "am_michael", "bf_emma", "bf_isabella",
            "bm_george", "bm_lewis"
        ]
        
        # Piper voices
        try:
            piper_voices = self.piper_service.get_available_voices()
            voices["piper"] = list(piper_voices.keys())
        except Exception as e:
            logger.warning(f"Could not load Piper voices: {e}")
            voices["piper"] = []
        
        return voices
    
    async def generate_script_segments(
        self,
        text: str,
        roles: List[str]
    ) -> List[Tuple[str, str]]:
        """
        Generate script segments with speaker assignments.
        Uses AI when available, falls back to pattern matching.
        
        Args:
            text: Input text to segment
            roles: List of detected role IDs/names
            
        Returns:
            List of tuples: (speaker_name, text_content)
        """
        # This is engine-agnostic, use Kokoro's implementation
        return await self.kokoro_service.generate_script_segments(text, roles)
    
    async def generate_audiobook(
        self,
        text: str,
        lang: str,
        role_voice_map: Dict[str, str],
        speed: float = 1.0,
        pitch_adjustments: Optional[Dict[str, int]] = None
    ) -> bytes:
        """
        Generate audiobook using the selected TTS engine.
        
        Args:
            text: Input text
            lang: Language code
            role_voice_map: Role to voice ID mapping
            speed: Playback speed
            pitch_adjustments: Optional dict of role to pitch shift (semitones)
            
        Returns:
            Audio data as bytes (WAV format)
        """
        if self.engine_name == TTSEngine.KOKORO:
            # Kokoro doesn't support pitch adjustment
            return await self.kokoro_service.generate_audiobook(
                text, lang, role_voice_map, speed
            )
        
        elif self.engine_name == TTSEngine.PIPER:
            # Use Piper with pitch support
            return await self._generate_with_piper(
                text, lang, role_voice_map, speed, pitch_adjustments or {}
            )
        
        else:
            raise ValueError(f"Unknown engine: {self.engine_name}")
    
    async def generate_audiobook_from_segments(
        self,
        segments: List[Tuple[str, str]],
        lang: str,
        role_voice_map: Dict[str, str],
        speed: float = 1.0,
        pitch_adjustments: Optional[Dict[str, int]] = None
    ) -> bytes:
        """
        Generate audiobook from pre-segmented script.
        
        Args:
            segments: List of (speaker_name, text_content) tuples
            lang: Language code
            role_voice_map: Role to voice ID mapping
            speed: Playback speed
            pitch_adjustments: Optional dict of role to pitch shift (semitones)
            
        Returns:
            Audio data as bytes (WAV format)
        """
        if self.engine_name == TTSEngine.KOKORO:
            return await self.kokoro_service.generate_audiobook_from_segments(
                segments, lang, role_voice_map, speed
            )
        
        elif self.engine_name == TTSEngine.PIPER:
            return await self._generate_from_segments_with_piper(
                segments, lang, role_voice_map, speed, pitch_adjustments or {}
            )
        
        else:
            raise ValueError(f"Unknown engine: {self.engine_name}")
    
    async def _generate_with_piper(
        self,
        text: str,
        lang: str,
        role_voice_map: Dict[str, str],
        speed: float,
        pitch_adjustments: Dict[str, int]
    ) -> bytes:
        """Generate audiobook using Piper TTS with pitch control."""
        import io
        import numpy as np
        import soundfile as sf
        
        # Segment text using Kokoro's logic
        segments = await self.kokoro_service.segment_text_with_ai(text, role_voice_map)
        
        if segments is None:
            segments = self.kokoro_service.segment_text(text, role_voice_map)
        
        if not segments:
            raise ValueError("No segments to synthesize")
        
        # Synthesize each segment with Piper
        audio_segments = []
        for role_id, voice, content in segments:
            pitch_shift = pitch_adjustments.get(role_id, 0)
            logger.debug(
                f"Synthesizing with Piper: role='{role_id}', voice='{voice}', "
                f"pitch={pitch_shift}"
            )
            
            audio = self.piper_service.synthesize_with_pitch(
                content, voice, speed, pitch_shift
            )
            audio_segments.append(audio)
        
        # Concatenate
        final_audio = np.concatenate([s for s in audio_segments if len(s) > 0])
        
        # Convert to WAV bytes
        buffer = io.BytesIO()
        sf.write(buffer, final_audio, self.piper_service.SAMPLE_RATE, format='WAV')
        buffer.seek(0)
        
        logger.info(f"Generated Piper audiobook: {len(buffer.getvalue())} bytes")
        return buffer.read()
    
    async def _generate_from_segments_with_piper(
        self,
        segments: List[Tuple[str, str]],
        lang: str,
        role_voice_map: Dict[str, str],
        speed: float,
        pitch_adjustments: Dict[str, int]
    ) -> bytes:
        """Generate audiobook from segments using Piper TTS."""
        import io
        import numpy as np
        import soundfile as sf
        
        if not segments:
            raise ValueError("No segments provided")
        
        # Synthesize each segment
        audio_segments = []
        for speaker, text in segments:
            voice = role_voice_map.get(speaker, self.piper_service.DEFAULT_VOICE)
            pitch_shift = pitch_adjustments.get(speaker, 0)
            
            logger.debug(
                f"Synthesizing segment: speaker='{speaker}', voice='{voice}', "
                f"pitch={pitch_shift}"
            )
            
            audio = self.piper_service.synthesize_with_pitch(
                text, voice, speed, pitch_shift
            )
            audio_segments.append(audio)
        
        # Concatenate
        final_audio = np.concatenate([s for s in audio_segments if len(s) > 0])
        
        # Convert to WAV bytes
        buffer = io.BytesIO()
        sf.write(buffer, final_audio, self.piper_service.SAMPLE_RATE, format='WAV')
        buffer.seek(0)
        
        logger.info(f"Generated Piper audiobook from segments: {len(buffer.getvalue())} bytes")
        return buffer.read()