tts_engine.py

"""
LOCAL TTS ENGINE v3 — Zero API Cost, Sentence-Chunk Streaming
==============================================================
FIXES FROM v2:

  FIX A — AudioPlayer._play_loop exits immediately after stop()+restart_if_needed()
    ROOT CAUSE: stop() sets self._running = False. restart_if_needed() calls
    start() which spawns a new thread — BUT start() also correctly sets
    self._running = True before the thread starts. HOWEVER the new thread
    checks `while self._running` at top-of-loop. If _STOP_SPEAKING is still
    set when the new thread reads it (line 169), it breaks immediately.
    FIX: start() now always clears _STOP_SPEAKING before spawning the thread,
    AND drains the queue so no stale sentinel can kill the fresh loop.

  FIX B — wait_done() joins the WRONG thread after restart_if_needed()
    ROOT CAUSE: speak() calls restart_if_needed() which may create a new
    thread stored in self._thread. Then speak() calls wait_done() which
    joins self._thread — this is correct. BUT the old play_loop exits its
    sd.OutputStream context *after* wait_done() returns if it was running
    concurrently, causing the next stream.write() to fail with "stream closed".
    FIX: restart_if_needed() now waits (up to 50ms) for the old thread to
    finish before starting the new one, preventing overlapping sd streams.

  FIX C — on_done() fires even when stopped mid-speech (barge-in)
    ROOT CAUSE: _play_loop always calls on_done() at the end regardless of
    whether it was stopped or completed normally. When barge-in stops TTS,
    notify_tts_done() fires prematurely → recorder drains too early →
    microphone picks up remaining TTS echo as the user's command.
    FIX: on_done() is only called when _STOP_SPEAKING is NOT set (i.e.,
    clean completion). Barge-in does NOT trigger notify_tts_done().

  FIX D — speak() timeout is proportional to text length but ignores pre-warming
    Pre-warmed phrases synthesize in <5ms. The old timeout formula
    `len(text)*0.1 + 3.0` gave 3 seconds minimum wait for "Yes?" (4 chars),
    blocking the listen loop for no reason.
    FIX: timeout = max(0.5, min(len(text) * 0.06 + 1.5, 30.0))
    This gives 1.7s for short phrases and scales cleanly for long ones.

  FIX E — _is_speaking never set in VoiceService when wake confirmation removed
    This is handled in service.py (see that file's fix notes).
    tts_engine.py exposes notify_recording_start() so VoiceService can
    suppress on_done() during the post-wake drain window.
"""

import io
import logging
import queue
import re
import threading
import time
from typing import Callable, Iterator, List, Optional

import numpy as np

logger = logging.getLogger(__name__)

# Global stop flag — checked every 20ms by AudioPlayer
_STOP_SPEAKING = threading.Event()


def stop_speaking():
    _STOP_SPEAKING.set()


def clear_stop():
    _STOP_SPEAKING.clear()


# ── SENTENCE CHUNKER ──────────────────────────────────────────────────────

class SentenceChunker:
    """
    Converts a stream of LLM tokens into complete sentences for TTS.

    Hard boundaries: .  !  ?  followed by whitespace.
    Soft boundaries: :  —  after 80+ chars (prevents very long mono-sentences).
    """

    def __init__(self):
        self._buf = ""

    def feed(self, token: str) -> List[str]:
        self._buf += token
        return self._extract()

    def flush(self) -> List[str]:
        remaining = self._buf.strip()
        self._buf = ""
        return [remaining] if remaining else []

    def _extract(self) -> List[str]:
        sentences = []
        while True:
            m = re.search(r'[.!?]\s+(?=[A-Z])', self._buf)
            m = re.search(r'[.!?]\s', self._buf)
            if m:
                end      = m.end()
                sentence = self._buf[:end].strip()
                self._buf = self._buf[end:]
                if sentence:
                    sentences.append(sentence)
                continue
            if len(self._buf) > 80:
                m2 = re.search(r'[:\-—]\s', self._buf[40:])
                if m2:
                    end      = 40 + m2.end()
                    sentence = self._buf[:end].strip()
                    self._buf = self._buf[end:]
                    if sentence:
                        sentences.append(sentence)
                    continue
            break
        return sentences


# ── AUDIO PLAYER ──────────────────────────────────────────────────────────

class AudioPlayer:
    """
    Plays float32 audio arrays via sounddevice in a background thread.
    Checks _STOP_SPEAKING every 20ms for instant barge-in.
    Calls on_done() ONLY on clean completion (not on barge-in stop).

    FIX A: start() clears _STOP_SPEAKING + drains queue before spawning thread.
    FIX B: restart_if_needed() waits for old thread before starting new one.
    FIX C: on_done() suppressed on barge-in stop.
    """

    def __init__(
        self,
        sample_rate: int = 22050,
        on_done: Optional[Callable] = None,
    ):
        self._sr      = sample_rate
        self._on_done = on_done
        self._q: queue.Queue = queue.Queue(maxsize=64)
        self._thread: Optional[threading.Thread] = None
        self._running = False
        self._lock    = threading.Lock()  # guards start/restart

    def start(self):
        """Start the audio player thread. Safe to call after stop()."""
        with self._lock:
            # FIX A: Clear stop flag and drain queue before new thread touches them
            _STOP_SPEAKING.clear()
            self._drain_queue()
            self._running = True
            self._thread  = threading.Thread(
                target=self._play_loop, daemon=True, name="Jarvis-AudioPlayer"
            )
            self._thread.start()

    def restart_if_needed(self):
        """
        Re-start the player thread if it has stopped (e.g. after barge-in).
        FIX B: waits up to 60ms for the old thread to exit before starting
        a fresh one, preventing two sd.OutputStreams opening simultaneously.
        """
        with self._lock:
            if self._thread and self._thread.is_alive():
                return  # Already running — nothing to do
            # Old thread is dead or never started — wait briefly then restart
            if self._thread:
                self._thread.join(timeout=0.06)  # up to 60ms drain
            # FIX A: clear flags/queue for the fresh thread
            _STOP_SPEAKING.clear()
            self._drain_queue()
            self._running = True
            self._thread  = threading.Thread(
                target=self._play_loop, daemon=True, name="Jarvis-AudioPlayer"
            )
            self._thread.start()

    def enqueue(self, audio: np.ndarray):
        try:
            self._q.put_nowait(audio)
        except queue.Full:
            pass

    def enqueue_sentinel(self):
        self._q.put(None)

    def stop(self):
        """
        Instant barge-in stop. Does NOT fire on_done (FIX C).
        Does NOT join the thread — returns immediately.
        """
        _STOP_SPEAKING.set()
        self._running = False
        self._drain_queue()

    def wait_done(self, timeout: float = 30.0):
        """Block until the current play_loop thread exits."""
        # Capture thread reference under lock to avoid TOCTOU with restart_if_needed
        with self._lock:
            t = self._thread
        if t and t.is_alive():
            t.join(timeout=timeout)

    def _drain_queue(self):
        """Discard all pending audio chunks and sentinels."""
        while not self._q.empty():
            try:
                self._q.get_nowait()
            except queue.Empty:
                break

    def _play_loop(self):
        try:
            import sounddevice as sd
        except ImportError:
            logger.error("[TTS] sounddevice not installed — no audio output")
            self._running = False
            return

        BLOCK_SIZE  = int(self._sr * 0.02)   # 20ms blocks for barge-in granularity
        clean_exit  = False                    # FIX C: track whether we finished naturally

        try:
            with sd.OutputStream(
                samplerate=self._sr, channels=1, dtype="float32", blocksize=BLOCK_SIZE
            ) as stream:
                while self._running:
                    if _STOP_SPEAKING.is_set():
                        break
                    try:
                        chunk = self._q.get(timeout=0.05)
                    except queue.Empty:
                        continue

                    if chunk is None:
                        # Sentinel — utterance finished naturally
                        clean_exit = True
                        break

                    # Play in 20ms sub-blocks for responsive barge-in
                    for i in range(0, len(chunk), BLOCK_SIZE):
                        if _STOP_SPEAKING.is_set():
                            break
                        block = chunk[i:i + BLOCK_SIZE].astype(np.float32)
                        if len(block) < BLOCK_SIZE:
                            block = np.pad(block, (0, BLOCK_SIZE - len(block)))
                        stream.write(block.reshape(-1, 1))

        except Exception as e:
            logger.warning(f"[TTS] AudioPlayer stream error: {e}")

        finally:
            self._running = False
            # FIX C: Only fire on_done if we finished cleanly (not barge-in)
            if clean_exit and self._on_done:
                try:
                    self._on_done()
                except Exception as ex:
                    logger.warning(f"[TTS] on_done callback error: {ex}")


# ── TTS BACKENDS ──────────────────────────────────────────────────────────

class KokoroBackend:
    """
    Kokoro TTS — 82M ONNX model, near ElevenLabs quality, 100% local.
    GPU-accelerated when CUDA is available.
    """

    def __init__(self, voice: str = "af_heart", speed: float = 1.1):
        self._voice       = voice
        self._speed       = speed
        self._model       = None
        self._sample_rate = 24000
        self._load()

    def _load(self):
        try:
            from kokoro_onnx import Kokoro
            providers = self._get_ort_providers()
            # ── THE FIX: Pass only the two file paths directly ──
            self._model = Kokoro("kokoro-v1.0.onnx", "voices.bin", providers=providers)
            logger.info(
                f"[TTS] Kokoro loaded | voice={self._voice} | "
                f"speed={self._speed}x | providers={providers}"
            )
        except ImportError:
            logger.warning("[TTS] kokoro-onnx not installed")
        except TypeError:
            try:
                from kokoro_onnx import Kokoro
                # ── THE FIX: Pass only the two file paths directly ──
                self._model = Kokoro("kokoro-v1.0.onnx", "voices.bin")
                logger.info("[TTS] Kokoro loaded (no GPU provider)")
            except Exception as e:
                logger.warning(f"[TTS] Kokoro load failed: {e}")
        except Exception as e:
            logger.warning(f"[TTS] Kokoro unavailable: {e}")

    def _get_ort_providers(self) -> List[str]:
        providers = []
        try:
            import torch
            if torch.cuda.is_available():
                providers.append("CUDAExecutionProvider")
        except ImportError:
            pass
        providers.append("CPUExecutionProvider")
        return providers

    @property
    def available(self) -> bool:
        return self._model is not None

    @property
    def sample_rate(self) -> int:
        return self._sample_rate

    def synthesize(self, text: str) -> np.ndarray:
        if not self._model:
            return np.zeros(self._sample_rate // 2, dtype=np.float32)
        try:
            samples, sr = self._model.create(text, voice=self._voice, speed=self._speed, lang="en-us")
            audio = samples.astype(np.float32)
            if audio.max() > 1.0:
                audio /= 32768.0
            return audio
        except Exception as e:
            logger.error(f"[TTS] Kokoro synthesis error: {e}")
            return np.zeros(self._sample_rate // 2, dtype=np.float32)


class PiperBackend:
    """
    Piper TTS — local, CPU-only, <50ms/sentence, zero cloud dependency.
    """

    def __init__(self):
        self._piper       = None
        self._sample_rate = 22050
        self._load()

    def _load(self):
        import wave as wave_mod
        model_paths = [
            ("models/en_US-lessac-medium.onnx", "models/en_US-lessac-medium.onnx.json"),
            ("en_US-lessac-medium.onnx",         "en_US-lessac-medium.onnx.json"),
        ]
        for model, config in model_paths:
            try:
                import piper
                self._piper = piper.PiperVoice.load(model, config_path=config, use_cuda=False)
                logger.info(f"[TTS] Piper loaded | model={model}")
                self._sample_rate = self._piper.config.sample_rate
                return
            except FileNotFoundError:
                continue
            except ImportError:
                logger.warning("[TTS] piper-tts not installed")
                return
            except Exception as e:
                logger.warning(f"[TTS] Piper load failed ({model}): {e}")

    @property
    def available(self) -> bool:
        return self._piper is not None

    @property
    def sample_rate(self) -> int:
        return self._sample_rate

    def synthesize(self, text: str) -> np.ndarray:
        import wave as wave_mod
        buf = io.BytesIO()
        with wave_mod.open(buf, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(self._sample_rate)
            self._piper.synthesize(text, wf)
        buf.seek(0)
        raw = np.frombuffer(buf.read(), dtype=np.int16)
        return raw.astype(np.float32) / 32768.0


class EdgeTTSBackend:
    """
    Microsoft edge-tts — cloud fallback, always available if installed.
    Uses a dedicated per-call event loop in a fresh thread to avoid
    'event loop already running' errors.
    """

    def __init__(self, voice: str = "en-US-ChristopherNeural", rate: str = "+18%"):
        self._voice       = voice
        self._rate        = rate
        self._sample_rate = 24000
        logger.info(f"[TTS] Edge TTS ready | voice={voice}")

    @property
    def available(self) -> bool:
        try:
            import edge_tts  # noqa: F401
            return True
        except ImportError:
            return False

    @property
    def sample_rate(self) -> int:
        return self._sample_rate

    def synthesize(self, text: str) -> np.ndarray:
        result_holder: List = [None]
        error_holder:  List = [None]

        def _worker():
            import asyncio
            loop = asyncio.new_event_loop()
            try:
                result_holder[0] = loop.run_until_complete(self._async_synth(text))
            except Exception as e:
                error_holder[0] = e
            finally:
                loop.close()

        t = threading.Thread(target=_worker, daemon=True)
        t.start()
        t.join(timeout=15.0)

        if error_holder[0]:
            logger.error(f"[TTS] EdgeTTS error: {error_holder[0]}")
            return np.zeros(self._sample_rate // 2, dtype=np.float32)

        return result_holder[0] if result_holder[0] is not None else \
               np.zeros(self._sample_rate // 2, dtype=np.float32)

    async def _async_synth(self, text: str) -> np.ndarray:
        import edge_tts
        import soundfile as sf

        comm = edge_tts.Communicate(text, self._voice, rate=self._rate)
        buf  = io.BytesIO()
        async for chunk in comm.stream():
            if chunk["type"] == "audio":
                buf.write(chunk["data"])
        buf.seek(0)
        try:
            data, _ = sf.read(buf, dtype="float32")
            return data if data.ndim == 1 else data[:, 0]
        except Exception:
            return np.zeros(self._sample_rate // 2, dtype=np.float32)


# ── MAIN TTS ENGINE ───────────────────────────────────────────────────────

class LocalTTSEngine:
    """
    Unified local TTS engine with sentence-chunk streaming.

    Auto-selects best available backend:
      Kokoro (GPU if available) → Piper → edge-tts

    on_done: callable invoked ONLY after each utterance finishes playing
             cleanly (NOT on barge-in stop). Wire to
             AudioRecorder.notify_tts_done() for accurate drain timing.

    FIX A/B/C/D applied — see module docstring.
    """

    _PREWARM_PHRASES = [
        "Yes?", "On it.", "Opening.", "Closing.", "Playing.",
        "Searching.", "One moment.", "Done.", "Listening.",
        "Still working.", "Almost there.", "That didn't work.",
        "Let me look into that.", "Stopping.",
    ]

    def __init__(
        self,
        voice:    str               = "af_heart",
        speed:    float             = 1.1,
        on_done:  Optional[Callable] = None,
    ):
        self._on_done  = on_done
        self._backend  = self._select_backend(voice, speed)
        self._player   = AudioPlayer(
            sample_rate=self._backend.sample_rate,
            on_done=on_done,
        )
        self._player.start()
        self._cache: dict      = {}
        self._cache_lock       = threading.Lock()

        # Pre-warm common phrases in background
        threading.Thread(
            target=self._prewarm, daemon=True, name="TTS-Prewarm"
        ).start()

    # ── PUBLIC API ─────────────────────────────────────────────────────────

    def speak(self, text: str, priority: bool = False):
        """
        Synthesize and play text. Blocks until playback finishes.
        priority=True: interrupts current speech first.

        FIX D: Timeout formula corrected so short pre-warmed phrases
        don't block the listen loop for 3+ seconds.
        """
        if not text or not text.strip():
            return
        text = text.strip()

        if priority:
            # Signal stop, wait briefly, then clear for new speech
            _STOP_SPEAKING.set()
            time.sleep(0.04)
            _STOP_SPEAKING.clear()

        # FIX A+B: Restart player cleanly (waits for old thread, clears flags)
        self._player.restart_if_needed()

        audio = self._get_audio(text)
        self._player.enqueue(audio)
        self._player.enqueue_sentinel()

        # FIX D: Corrected timeout — proportional to text, but fast for short phrases
        timeout = max(0.5, min(len(text) * 0.06 + 1.5, 30.0))
        self._player.wait_done(timeout=timeout)

    def speak_streaming(
        self,
        token_iterator: Iterator[str],
        on_sentence: Optional[Callable[[str], None]] = None,
    ):
        """
        Non-blocking streaming TTS.
        Speaks sentence 1 while LLM generates sentence 2, etc.
        """
        clear_stop()
        chunker = SentenceChunker()

        def _stream():
            self._player.restart_if_needed()
            try:
                for token in token_iterator:
                    if _STOP_SPEAKING.is_set():
                        break
                    for sentence in chunker.feed(token):
                        if _STOP_SPEAKING.is_set():
                            break
                        self._enqueue_sentence(sentence, on_sentence)
                if not _STOP_SPEAKING.is_set():
                    for sentence in chunker.flush():
                        self._enqueue_sentence(sentence, on_sentence)
            finally:
                self._player.enqueue_sentinel()

        threading.Thread(target=_stream, daemon=True, name="TTS-Stream").start()

    def stop(self):
        """Instant barge-in stop. Does NOT fire on_done."""
        stop_speaking()
        self._player.stop()

    def cleanup(self):
        self._player.stop()

    # ── INTERNAL ──────────────────────────────────────────────────────────

    def _select_backend(self, voice: str, speed: float):
        for BackendClass, kwargs in [
            (KokoroBackend,  {"voice": voice, "speed": speed}),
            (PiperBackend,   {}),
            (EdgeTTSBackend, {}),
        ]:
            try:
                b = BackendClass(**kwargs)
                if b.available:
                    return b
            except Exception as e:
                logger.debug(f"[TTS] Backend {BackendClass.__name__} failed: {e}")
        logger.error("[TTS] All backends failed — TTS will be silent")
        return EdgeTTSBackend()

    def _get_audio(self, text: str) -> np.ndarray:
        with self._cache_lock:
            if text in self._cache:
                return self._cache[text]
        try:
            import re
            
            def normalize_numbers(t: str) -> str:
                def convert(match):
                    raw = match.group()
                    num = int(raw.replace(",", ""))
                    try:
                        from num2words import num2words
                        return num2words(num, lang="en_IN")
                    except Exception:
                        return raw

                t = re.sub(r'\b\d[\d,]*\b', convert, t)
                t = t.replace("₹", " rupees ")
                t = t.replace("$", " dollars ")
                return t

            clean_text = normalize_numbers(text)
            audio = self._backend.synthesize(clean_text)
            
        except Exception as e:
            logger.error(f"[TTS] Synthesis error for '{text[:30]}': {e}")
            audio = np.zeros(int(self._backend.sample_rate * 0.3), dtype=np.float32)
            
        with self._cache_lock:
            if len(self._cache) >= 300:
                oldest = next(iter(self._cache))
                del self._cache[oldest]
            self._cache[text] = audio
        return audio

    def _enqueue_sentence(self, sentence: str, callback: Optional[Callable]):
        audio = self._get_audio(sentence)
        try:
            from ui_bridge import ui_bridge
            ui_bridge.broadcast("speaking", sentence[:80])
        except Exception:
            pass
            
        self._player.enqueue(audio)
        if callback:
            try:
                callback(sentence)
            except Exception:
                pass

    def _prewarm(self):
        """Cache common phrases before user says anything."""
        warmed = 0
        for phrase in self._PREWARM_PHRASES:
            if _STOP_SPEAKING.is_set():
                break
            try:
                audio = self._backend.synthesize(phrase)
                with self._cache_lock:
                    self._cache[phrase] = audio
                warmed += 1
            except Exception:
                pass
        logger.info(f"[TTS] Pre-warmed {warmed}/{len(self._PREWARM_PHRASES)} phrases")