tts.py

"""
TTS module — unified interface for text-to-speech synthesis.

Supports two backends:
  - Qwen3-TTS Base (1.7B): voice-cloned synthesis using a cached voice profile
  - Qwen3-TTS CustomVoice (0.6B): fast predefined speakers (default stock voice)

Features:
  - Audio cache: avoids re-synthesis for repeated text+voice combos
  - Pre-generation: background-generates entire story on book select
  - Transition chime: soft audio between paragraphs to mask gaps
  - Eager pre-buffering: generates next chunks while current plays

Usage:
    chunks = split_into_chunks(text)
    for sr, wav, i, n, err in generate_audio_stream(chunks, voice_profile_id="abc123"):
        ...
"""
import hashlib
import logging
import os
import queue
import re
import threading
from collections import OrderedDict

import numpy as np
import soundfile as sf

from runtime_config import AUDIO_CACHE_DIR, GPU_INFERENCE_LOCK, APP_ROOT

logger = logging.getLogger(__name__)

_SENTENCE_RE = re.compile(r'(?<=[.!?;])\s+')
_CLAUSE_RE = re.compile(r'(?<=[,;:\u2014])\s+|(?<=\.)\s+|(?<=[!?])\s+')
_SENTINEL = object()

# Target chunk length — shorter chunks = lower latency per chunk
_MAX_CHUNK_CHARS = 120

# Audio cache directory
_CACHE_DIR = str(AUDIO_CACHE_DIR)
os.makedirs(_CACHE_DIR, exist_ok=True)

# Copy pre-generated audio from app bundle to runtime cache (HF Space: repo → /data)
_BUNDLED_CACHE = APP_ROOT / "audio_cache"
if _BUNDLED_CACHE.is_dir() and str(_BUNDLED_CACHE) != _CACHE_DIR:
    import shutil
    _copied = 0
    for _wav in _BUNDLED_CACHE.glob("*.wav"):
        _dest = os.path.join(_CACHE_DIR, _wav.name)
        if not os.path.exists(_dest):
            shutil.copy2(str(_wav), _dest)
            _copied += 1
    if _copied:
        logger.info("Copied %d pre-generated audio files to runtime cache.", _copied)

# Transition chime (soft sine fade, 0.3s at 24kHz)
_CHIME_SR = 24000
_CHIME_DURATION = 0.3
_chime_t = np.linspace(0, _CHIME_DURATION, int(_CHIME_SR * _CHIME_DURATION), dtype=np.float32)
_TRANSITION_CHIME = 0.08 * np.sin(2 * np.pi * 440 * _chime_t) * np.linspace(1, 0, len(_chime_t))

# Background pre-generation state
_pregen_lock = threading.Lock()
_pregen_cache: OrderedDict[str, list[tuple[int, np.ndarray]]] = OrderedDict()
_pregen_in_progress: set[str] = set()
_pregen_progress: dict[str, dict[str, int | bool]] = {}
_pregen_cancel_events: dict[str, threading.Event] = {}

# Keep background work small so live Ask/playback can take the GPU quickly.
_PREGEN_CHUNK_LIMIT = int(os.environ.get("MOMSVOICE_PREGEN_CHUNKS", "3"))
_PREGEN_CACHE_MAX_STORIES = int(os.environ.get("MOMSVOICE_PREGEN_CACHE_STORIES", "4"))
_PREGEN_PROGRESS_MAX_STORIES = int(os.environ.get("MOMSVOICE_PREGEN_PROGRESS_STORIES", "16"))
_AUDIO_CACHE_MAX_FILES = int(os.environ.get("MOMSVOICE_AUDIO_CACHE_MAX_FILES", "400"))
_AUDIO_CACHE_MAX_BYTES = int(os.environ.get("MOMSVOICE_AUDIO_CACHE_MAX_BYTES", str(512 * 1024 * 1024)))


def _cache_key(text: str, voice_profile_id: str | None) -> str:
    """Generate a cache key from text + voice profile."""
    raw = f"{voice_profile_id or 'stock'}:{text}"
    return hashlib.md5(raw.encode()).hexdigest()


def _story_key(chunks: list[str], voice_profile_id: str | None) -> str:
    return _cache_key("\n".join(chunks), voice_profile_id)


def _prune_pregen_progress_locked() -> None:
    removable = [
        story_key for story_key, progress in _pregen_progress.items()
        if story_key not in _pregen_in_progress and not progress.get("in_progress")
    ]
    while len(_pregen_progress) > _PREGEN_PROGRESS_MAX_STORIES and removable:
        _pregen_progress.pop(removable.pop(0), None)


def cancel_pregeneration() -> None:
    """Ask all background pre-generation workers to stop after their current chunk."""
    with _pregen_lock:
        for event in _pregen_cancel_events.values():
            event.set()
        for story_key in list(_pregen_in_progress):
            progress = _pregen_progress.get(story_key)
            if progress is not None:
                progress["in_progress"] = False
                progress["cancelled"] = True
        _pregen_in_progress.clear()
        _prune_pregen_progress_locked()


def _prune_pregen_cache_locked() -> None:
    while len(_pregen_cache) > _PREGEN_CACHE_MAX_STORIES:
        _pregen_cache.popitem(last=False)


def get_pregeneration_status(
    chunks: list[str],
    voice_profile_id: str | None = None,
) -> dict[str, int | bool]:
    """Return how much of a story is actually cached or warmed."""
    total = len(chunks)
    if total == 0:
        return {"cached": 0, "total": 0, "in_progress": False, "complete": False}

    story_key = _story_key(chunks, voice_profile_id)
    with _pregen_lock:
        if story_key in _pregen_cache:
            _pregen_cache.move_to_end(story_key)
            return {
                "cached": len(_pregen_cache[story_key]),
                "total": total,
                "target": total,
                "in_progress": False,
                "complete": True,
                "errors": 0,
                "cancelled": False,
            }
        progress = _pregen_progress.get(story_key)
        if progress is not None:
            return dict(progress)

    target = min(total, _PREGEN_CHUNK_LIMIT)
    cached = sum(
        1 for chunk in chunks[:target]
        if _get_cached_audio(chunk, voice_profile_id) is not None
    )
    return {
        "cached": cached,
        "total": total,
        "target": target,
        "in_progress": False,
        "complete": cached == total,
        "errors": 0,
        "cancelled": False,
    }


def _get_cached_audio(chunk_text: str, voice_profile_id: str | None) -> np.ndarray | None:
    """Check if audio for this chunk is already cached on disk."""
    key = _cache_key(chunk_text, voice_profile_id)
    path = os.path.join(_CACHE_DIR, f"{key}.wav")
    if os.path.exists(path):
        try:
            wav, sr = sf.read(path, dtype='float32')
            return wav
        except Exception:
            pass
    return None


def _save_cached_audio(chunk_text: str, voice_profile_id: str | None, wav: np.ndarray, sr: int):
    """Save synthesized audio to disk cache."""
    key = _cache_key(chunk_text, voice_profile_id)
    path = os.path.join(_CACHE_DIR, f"{key}.wav")
    try:
        sf.write(path, wav, sr)
        _prune_audio_cache()
    except Exception:
        pass


def _prune_audio_cache():
    """Bound disk cache by deleting oldest cached audio files."""
    try:
        entries = [
            entry for entry in os.scandir(_CACHE_DIR)
            if entry.is_file() and entry.name.endswith(".wav")
        ]
        total_bytes = sum(entry.stat().st_size for entry in entries)
        if len(entries) <= _AUDIO_CACHE_MAX_FILES and total_bytes <= _AUDIO_CACHE_MAX_BYTES:
            return

        entries.sort(key=lambda entry: entry.stat().st_mtime)
        idx = 0
        while idx < len(entries):
            if len(entries) <= _AUDIO_CACHE_MAX_FILES and total_bytes <= _AUDIO_CACHE_MAX_BYTES:
                break
            entry = entries[idx]
            try:
                size = entry.stat().st_size
                os.remove(entry.path)
                total_bytes -= size
                entries.pop(idx)
            except OSError:
                idx += 1
                continue
    except OSError:
        return


def split_into_chunks(text: str) -> list[str]:
    """Split text into short chunks suitable for low-latency TTS streaming.
    
    Splits by sentence first, then further splits long sentences by clause
    boundaries (commas, semicolons, em-dashes) to keep each chunk under
    ~120 chars for faster TTS generation.
    """
    sentences = _SENTENCE_RE.split(text.strip())
    chunks = []
    for sent in sentences:
        sent = sent.strip()
        if not sent:
            continue
        if len(sent) <= _MAX_CHUNK_CHARS:
            chunks.append(sent)
        else:
            # Split long sentence into clauses
            clauses = _CLAUSE_RE.split(sent)
            current = ""
            for clause in clauses:
                clause = clause.strip()
                if not clause:
                    continue
                if current and len(current) + len(clause) + 1 > _MAX_CHUNK_CHARS:
                    chunks.append(current.strip())
                    current = clause
                else:
                    current = f"{current} {clause}".strip() if current else clause
            if current.strip():
                chunks.append(current.strip())
    return [c for c in chunks if c]


def pregenerate_story_audio(
    chunks: list[str],
    voice_profile_id: str | None = None,
    max_chunks: int = _PREGEN_CHUNK_LIMIT,
):
    """Pre-generate the first few story chunks in background.
    
    Call this on book selection to pre-warm initial playback. Non-blocking.
    """
    story_key = _story_key(chunks, voice_profile_id)
    target_chunks = chunks[:max(0, min(len(chunks), max_chunks))]
    target = len(target_chunks)
    if target == 0:
        return
    
    with _pregen_lock:
        if story_key in _pregen_in_progress or story_key in _pregen_cache:
            return  # Already running or done
        cancel_event = threading.Event()
        _pregen_in_progress.add(story_key)
        _pregen_cancel_events[story_key] = cancel_event
        _pregen_progress[story_key] = {
            "cached": 0,
            "total": len(chunks),
            "target": target,
            "in_progress": True,
            "complete": False,
            "errors": 0,
            "cancelled": False,
        }

    def _worker():
        results = []
        errors = 0
        sr = 24000
        for i, chunk in enumerate(target_chunks):
            if cancel_event.is_set():
                break
            # Check disk cache first
            cached = _get_cached_audio(chunk, voice_profile_id)
            if cached is not None:
                results.append((sr, cached))
                with _pregen_lock:
                    if _pregen_cancel_events.get(story_key) is cancel_event:
                        _pregen_progress[story_key]["cached"] = len(results)
                continue
            # Synthesize
            try:
                acquired = False
                while not cancel_event.is_set():
                    acquired = GPU_INFERENCE_LOCK.acquire(timeout=0.1)
                    if acquired:
                        break
                if not acquired or cancel_event.is_set():
                    break
                try:
                    wav, sample_rate = _synthesize_single(chunk, voice_profile_id)
                finally:
                    GPU_INFERENCE_LOCK.release()
                sr = sample_rate
                results.append((sr, wav))
                _save_cached_audio(chunk, voice_profile_id, wav, sr)
            except Exception as e:
                errors += 1
                logger.warning("Pre-gen failed on chunk %d: %s", i, e)
            with _pregen_lock:
                if _pregen_cancel_events.get(story_key) is cancel_event:
                    _pregen_progress[story_key]["cached"] = len(results)
                    _pregen_progress[story_key]["errors"] = errors
        
        with _pregen_lock:
            if _pregen_cancel_events.get(story_key) is not cancel_event:
                return
            cancelled = cancel_event.is_set()
            fully_cached = len(results) == len(chunks) and errors == 0 and not cancelled
            if fully_cached:
                _pregen_cache[story_key] = results
                _pregen_cache.move_to_end(story_key)
                _prune_pregen_cache_locked()
            _pregen_progress[story_key] = {
                "cached": len(results),
                "total": len(chunks),
                "target": target,
                "in_progress": False,
                "complete": fully_cached,
                "errors": errors,
                "cancelled": cancelled,
            }
            _pregen_in_progress.discard(story_key)
            _pregen_cancel_events.pop(story_key, None)
            _prune_pregen_progress_locked()
        logger.info(
            "Pre-generation finished: %d/%d initial chunks cached (errors=%d, cancelled=%s).",
            len(results), target, errors, cancelled,
        )

    threading.Thread(target=_worker, daemon=True).start()


def _synthesize_single(text: str, voice_profile_id: str | None) -> tuple[np.ndarray, int]:
    """Synthesize a single chunk, choosing backend based on profile."""
    if voice_profile_id:
        from voice_clone import synthesize_cloned, synthesize_custom_voice
        try:
            wav, sr = synthesize_cloned(text, voice_profile_id)
            return wav, sr
        except Exception:
            wav, sr = synthesize_custom_voice(text)
            return wav, sr
    else:
        from voice_clone import synthesize_custom_voice
        wav, sr = synthesize_custom_voice(text)
        return wav, sr


def generate_audio_stream(
    chunks: list[str],
    voice_profile_id: str | None = None,
    custom_voice_speaker: str = "vivian",
    add_transitions: bool = True,
):
    """
    Generator: synthesizes chunks with caching + pre-buffering + transitions.
    
    If pre-generated audio is available (from pregenerate_story_audio), serves
    instantly from cache. Otherwise synthesizes on-demand with background pre-buffer.
    
    Yields (sample_rate, wav_array, chunk_idx, total_chunks, error_msg).
    """
    n = len(chunks)
    sample_rate = 24000

    # Check if pre-generated cache is available
    story_key = _story_key(chunks, voice_profile_id)
    pregen_results = None
    with _pregen_lock:
        if story_key in _pregen_cache:
            pregen_results = _pregen_cache[story_key]

    if pregen_results and len(pregen_results) == n:
        # Serve from pre-generated cache — near-zero latency
        for i, (sr, wav) in enumerate(pregen_results):
            if wav is not None and len(wav) > 0:
                if add_transitions and i > 0:
                    # Prepend transition chime
                    wav = np.concatenate([_TRANSITION_CHIME, wav])
                yield sr, wav, i, n, None
        return

    # Fallback: on-demand synthesis with pre-buffering (4 chunks ahead)
    chunk_q: queue.Queue = queue.Queue(maxsize=4)

    if voice_profile_id:
        _start_qwen_worker(chunks, voice_profile_id, chunk_q)
    else:
        _start_custom_voice_worker(chunks, custom_voice_speaker, chunk_q)

    # Batch small audio segments into larger blocks for smoother playback
    audio_buffer = []
    buffer_samples = 0
    last_idx = 0
    _TARGET_SAMPLES = 24000 * 5  # ~5s blocks
    chunk_count = 0

    while True:
        item = chunk_q.get()
        if item is _SENTINEL:
            if audio_buffer:
                combined = np.concatenate(audio_buffer).astype(np.float32)
                yield sample_rate, combined, last_idx, n, None
            break
        i, wav, err = item
        if err:
            if audio_buffer:
                combined = np.concatenate(audio_buffer).astype(np.float32)
                yield sample_rate, combined, last_idx, n, None
                audio_buffer = []
                buffer_samples = 0
            yield sample_rate, np.zeros(0, dtype=np.float32), i, n, err
            break
        last_idx = i

        # Skip empty audio segments
        if wav is None or len(wav) == 0:
            continue
        
        # Add transition chime between chunks (not before first)
        if add_transitions and chunk_count > 0:
            audio_buffer.append(_TRANSITION_CHIME)
            buffer_samples += len(_TRANSITION_CHIME)
        
        audio_buffer.append(wav)
        buffer_samples += len(wav)
        chunk_count += 1

        # Cache this chunk for future replays
        _save_cached_audio(chunks[i], voice_profile_id, wav, sample_rate)

        # Yield when buffer reaches target size or this is the first chunk (fast start)
        if buffer_samples >= _TARGET_SAMPLES or (i == 0 and buffer_samples > 0):
            combined = np.concatenate(audio_buffer).astype(np.float32)
            yield sample_rate, combined, i, n, None
            audio_buffer = []
            buffer_samples = 0


def _start_qwen_worker(chunks, profile_id, chunk_q):
    """Background thread: synthesize chunks with Qwen3-TTS voice clone (Base 1.7B).
    Falls back to stock voice if cloned synthesis fails."""
    def _worker():
        from voice_clone import synthesize_cloned, synthesize_custom_voice
        use_fallback = False
        for i, stmt in enumerate(chunks):
            # Check cache first
            cached = _get_cached_audio(stmt, profile_id)
            if cached is not None:
                chunk_q.put((i, cached, None))
                continue
            try:
                if use_fallback:
                    wav, _sr = synthesize_custom_voice(stmt)
                else:
                    wav, _sr = synthesize_cloned(stmt, profile_id)
                chunk_q.put((i, wav, None))
            except Exception as exc:
                if not use_fallback:
                    logger.warning("Cloned voice failed (chunk %d): %s — falling back to stock voice", i, exc)
                    use_fallback = True
                    try:
                        wav, _sr = synthesize_custom_voice(stmt)
                        chunk_q.put((i, wav, None))
                    except Exception as exc2:
                        logger.exception("Stock voice also failed on chunk %d", i)
                        chunk_q.put((i, None, str(exc2)))
                        return
                else:
                    logger.exception("Stock voice synthesis failed on chunk %d", i)
                    chunk_q.put((i, None, str(exc)))
                    return
        chunk_q.put(_SENTINEL)

    threading.Thread(target=_worker, daemon=True).start()


def _start_custom_voice_worker(chunks, speaker, chunk_q):
    """Background thread: synthesize chunks with Qwen3-TTS CustomVoice (0.6B).
    Uses sub-segment streaming for lower latency — yields partial audio as generated."""
    def _worker():
        from voice_clone import synthesize_custom_voice_streaming
        import numpy as np
        for i, stmt in enumerate(chunks):
            # Check cache first
            cached = _get_cached_audio(stmt, None)
            if cached is not None:
                chunk_q.put((i, cached, None))
                continue
            try:
                segments = []
                for seg, _sr in synthesize_custom_voice_streaming(stmt, speaker=speaker):
                    segments.append(seg)
                if segments:
                    wav = np.concatenate(segments)
                    chunk_q.put((i, wav, None))
                else:
                    chunk_q.put((i, np.zeros(0, dtype=np.float32), None))
            except Exception as exc:
                logger.exception("CustomVoice synthesis failed on chunk %d", i)
                chunk_q.put((i, None, str(exc)))
                return
        chunk_q.put(_SENTINEL)

    threading.Thread(target=_worker, daemon=True).start()