audiobook/backend/services/tts_service.py

"""
TTS service for generating multi-voice audiobooks using Kokoro TTS.
"""
import logging
import re
import io
from typing import Dict, List, Tuple, Optional
import numpy as np
import soundfile as sf
from kokoro import KPipeline
from .ollama_service import OllamaService

logger = logging.getLogger(__name__)


class TTSService:
    """
    Service for text-to-speech synthesis using Kokoro TTS.
    
    Handles pipeline initialization, text segmentation, voice assignment,
    and audio concatenation for multi-voice audiobook generation.
    """
    
    # Default voice for unmapped roles
    DEFAULT_VOICE = "af_bella"
    
    # Sample rate for Kokoro TTS
    SAMPLE_RATE = 24000
    
    # Role marker patterns
    COLON_PATTERN = r'^([A-Z][a-zA-Z\s]+?):\s*(.+)$'
    BRACKET_PATTERN = r'^\[([A-Z][a-zA-Z\s]+?)\]\s*(.+)$'
    
    def __init__(self, ollama_service: Optional[OllamaService] = None):
        """Initializes TTS service with empty pipeline cache and optional Ollama service."""
        self.pipelines: Dict[str, KPipeline] = {}
        self.ollama_service = ollama_service or OllamaService()
        logger.info("TTSService initialized")
    
    def get_pipeline(self, lang_code: str) -> KPipeline:
        """
        Returns or creates KPipeline for given language.
        Caches pipelines to avoid reloading.
        
        Args:
            lang_code: Language code (a, b, e, f, h, i, j, p, z)
            
        Returns:
            KPipeline instance for the language
            
        Raises:
            ValueError: If language code is not supported
        """
        supported_langs = ['a', 'b', 'e', 'f', 'h', 'i', 'j', 'p', 'z']
        if lang_code not in supported_langs:
            raise ValueError(
                f"Unsupported language code: {lang_code}. "
                f"Supported codes: {supported_langs}"
            )
        
        # Return cached pipeline if available
        if lang_code in self.pipelines:
            logger.debug(f"Using cached pipeline for language: {lang_code}")
            return self.pipelines[lang_code]
        
        # Create new pipeline
        logger.info(f"Initializing KPipeline for language: {lang_code}")
        try:
            pipeline = KPipeline(lang_code=lang_code)
            self.pipelines[lang_code] = pipeline
            logger.info(f"Successfully initialized pipeline for language: {lang_code}")
            return pipeline
        except Exception as e:
            logger.error(f"Failed to initialize pipeline for language {lang_code}: {e}")
            raise
    
    async def segment_text_with_ai(
        self,
        text: str,
        role_voice_map: Dict[str, str]
    ) -> Optional[List[Tuple[str, str, str]]]:
        """
        Uses AI to segment text and assign speakers intelligently.
        
        Args:
            text: Input text to segment
            role_voice_map: Mapping of role IDs to voice IDs
            
        Returns:
            List of tuples: (role_id, voice_id, text_content) or None if AI unavailable
        """
        try:
            # Check if Ollama is available
            is_available = await self.ollama_service.is_available()
            if not is_available:
                logger.info("Ollama not available for AI segmentation")
                return None
            
            # Get list of speakers from role_voice_map
            speakers = list(role_voice_map.keys())
            
            # Request AI segmentation
            ai_segments = await self.ollama_service.segment_text(text, speakers)
            
            if not ai_segments:
                logger.warning("AI segmentation returned no segments")
                return None
            
            # Convert to format with voice assignments
            segments = []
            for speaker, content in ai_segments:
                voice = role_voice_map.get(speaker, self.DEFAULT_VOICE)
                segments.append((speaker, voice, content))
            
            logger.info(f"AI segmented text into {len(segments)} segments")
            return segments
            
        except Exception as e:
            logger.warning(f"AI segmentation failed: {e}")
            return None
    
    def segment_text(
        self, 
        text: str, 
        role_voice_map: Dict[str, str]
    ) -> List[Tuple[str, str, str]]:
        """
        Segments text into (role_id, voice, text_content) tuples using pattern matching.
        Identifies role markers and assigns segments accordingly.
        
        This is the fallback method when AI segmentation is not available.
        
        Args:
            text: Input text to segment
            role_voice_map: Mapping of role IDs to voice IDs
            
        Returns:
            List of tuples: (role_id, voice_id, text_content)
        """
        segments = []
        lines = text.split('\n')
        
        for line in lines:
            line = line.strip()
            if not line:
                continue
            
            # Try colon pattern first
            colon_match = re.match(self.COLON_PATTERN, line, re.MULTILINE)
            if colon_match:
                role_name = colon_match.group(1).strip()
                content = colon_match.group(2).strip()
                voice = role_voice_map.get(role_name, self.DEFAULT_VOICE)
                segments.append((role_name, voice, content))
                continue
            
            # Try bracket pattern
            bracket_match = re.match(self.BRACKET_PATTERN, line)
            if bracket_match:
                role_name = bracket_match.group(1).strip()
                content = bracket_match.group(2).strip()
                voice = role_voice_map.get(role_name, self.DEFAULT_VOICE)
                segments.append((role_name, voice, content))
                continue
            
            # No role marker found - use default narrator
            voice = role_voice_map.get("Narrator", self.DEFAULT_VOICE)
            segments.append(("Narrator", voice, line))
        
        logger.info(f"Pattern-based segmentation produced {len(segments)} segments")
        return segments
    
    def _chunk_text(self, text: str, max_length: int = 500) -> List[str]:
        """
        Splits long text into manageable chunks.
        Splits on sentence boundaries when possible.
        
        Args:
            text: Text to chunk
            max_length: Maximum characters per chunk
            
        Returns:
            List of text chunks
        """
        if len(text) <= max_length:
            return [text]
        
        chunks = []
        # Split on sentence boundaries
        sentences = re.split(r'([.!?]+\s+)', text)
        
        current_chunk = ""
        for i in range(0, len(sentences), 2):
            sentence = sentences[i]
            punctuation = sentences[i + 1] if i + 1 < len(sentences) else ""
            full_sentence = sentence + punctuation
            
            if len(current_chunk) + len(full_sentence) <= max_length:
                current_chunk += full_sentence
            else:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = full_sentence
        
        if current_chunk:
            chunks.append(current_chunk.strip())
        
        return chunks if chunks else [text]
    
    def synthesize_segment(
        self,
        pipeline: KPipeline,
        text: str,
        voice: str,
        speed: float = 1.0
    ) -> np.ndarray:
        """
        Synthesizes a single text segment with specified voice.
        Handles long text by chunking.
        
        Args:
            pipeline: KPipeline instance
            text: Text to synthesize
            voice: Kokoro voice ID
            speed: Playback speed (0.5-1.5)
            
        Returns:
            Audio samples as numpy array
        """
        # Chunk long text
        chunks = self._chunk_text(text)
        
        audio_chunks = []
        for chunk in chunks:
            try:
                # Generate audio for chunk
                result = pipeline(chunk, voice=voice, speed=speed)
                # Kokoro returns a generator that yields Result objects
                # Each Result object has an .audio attribute (torch.Tensor)
                audio_parts = []
                for part in result:
                    # Extract audio tensor and convert to numpy
                    if hasattr(part, 'audio'):
                        audio_tensor = part.audio
                        # Convert torch tensor to numpy and flatten
                        audio_np = audio_tensor.cpu().numpy().flatten()
                        audio_parts.append(audio_np)
                
                # Concatenate all parts
                if audio_parts:
                    audio = np.concatenate(audio_parts)
                    audio_chunks.append(audio)
            except Exception as e:
                logger.error(f"Error synthesizing chunk with voice {voice}: {e}")
                raise
        
        # Concatenate chunks if multiple
        if len(audio_chunks) > 1:
            return np.concatenate(audio_chunks)
        elif audio_chunks:
            return audio_chunks[0]
        else:
            return np.array([], dtype=np.float32)
    
    def concatenate_audio(self, segments: List[np.ndarray]) -> np.ndarray:
        """
        Concatenates multiple audio segments into single array.
        Ensures consistent sample rate (24kHz).
        
        Args:
            segments: List of audio arrays
            
        Returns:
            Concatenated audio array
        """
        if not segments:
            return np.array([], dtype=np.float32)
        
        # Filter out empty segments
        valid_segments = [s for s in segments if len(s) > 0]
        
        if not valid_segments:
            return np.array([], dtype=np.float32)
        
        # Concatenate all segments
        concatenated = np.concatenate(valid_segments)
        logger.info(f"Concatenated {len(valid_segments)} audio segments")
        
        return concatenated
    
    async def generate_script_segments(
        self,
        text: str,
        roles: List[str]
    ) -> List[Tuple[str, str]]:
        """
        Generates script segments with speaker assignments using AI.
        
        This method is used for the script preview feature, allowing users to
        review and edit speaker assignments before generating audio.
        
        Args:
            text: Input text to segment
            roles: List of detected role IDs/names
            
        Returns:
            List of tuples: (speaker_name, text_content)
            
        Raises:
            Exception: If both AI and pattern-based segmentation fail
        """
        logger.info(f"Generating script segments for {len(text)} chars with {len(roles)} roles")
        
        # Try AI segmentation first
        try:
            is_available = await self.ollama_service.is_available()
            if is_available:
                logger.info("Using AI for script segmentation")
                segments = await self.ollama_service.segment_text(text, roles)
                
                if segments:
                    logger.info(f"AI segmentation produced {len(segments)} segments")
                    return segments
                else:
                    logger.warning("AI segmentation returned empty result")
            else:
                logger.info("Ollama not available for script segmentation")
        except Exception as e:
            logger.warning(f"AI segmentation failed: {e}")
        
        # Fall back to pattern-based segmentation
        logger.info("Using pattern-based segmentation as fallback")
        
        # Create a dummy role_voice_map for pattern-based segmentation
        # We don't need actual voices here, just role assignments
        role_voice_map = {role: self.DEFAULT_VOICE for role in roles}
        
        # Use existing segment_text method
        segments_with_voice = self.segment_text(text, role_voice_map)
        
        # Convert from (role, voice, text) to (role, text)
        segments = [(role, text) for role, voice, text in segments_with_voice]
        
        logger.info(f"Pattern-based segmentation produced {len(segments)} segments")
        return segments
    
    async def generate_audiobook(
        self,
        text: str,
        lang: str,
        role_voice_map: Dict[str, str],
        speed: float = 1.0
    ) -> bytes:
        """
        Main orchestration method for audiobook generation.
        Segments text, synthesizes each segment, concatenates, and returns audio bytes.
        
        Uses AI segmentation when available, falls back to pattern matching.
        
        Args:
            text: Input text
            lang: Language code
            role_voice_map: Role to voice ID mapping
            speed: Playback speed
            
        Returns:
            Audio data as bytes (WAV format)
        """
        logger.info(
            f"Generating audiobook: lang={lang}, speed={speed}, "
            f"roles={len(role_voice_map)}"
        )
        
        # Get or create pipeline for language
        pipeline = self.get_pipeline(lang)
        
        # Try AI segmentation first, fall back to pattern matching
        segments = await self.segment_text_with_ai(text, role_voice_map)
        
        if segments is None:
            logger.info("Using pattern-based segmentation as fallback")
            segments = self.segment_text(text, role_voice_map)
        else:
            logger.info("Using AI-powered segmentation")
        
        if not segments:
            raise ValueError("No segments to synthesize")
        
        # Synthesize each segment
        audio_segments = []
        for role_id, voice, content in segments:
            logger.debug(f"Synthesizing segment for role '{role_id}' with voice '{voice}'")
            audio = self.synthesize_segment(pipeline, content, voice, speed)
            audio_segments.append(audio)
        
        # Concatenate all audio
        final_audio = self.concatenate_audio(audio_segments)
        
        # Convert to bytes (WAV format)
        buffer = io.BytesIO()
        sf.write(buffer, final_audio, self.SAMPLE_RATE, format='WAV')
        buffer.seek(0)
        audio_bytes = buffer.read()
        
        logger.info(f"Generated audiobook: {len(audio_bytes)} bytes")
        return audio_bytes