src/services/speech.py

"""語音服務 — STT (faster-whisper) + TTS (Kokoro)（透過 ModelRegistry）

【這個模組在做什麼？】
  STT (Speech-to-Text): 把使用者錄音轉成文字
  TTS (Text-to-Speech): 把 AI 回答轉成語音播放

【模型選型】
  STT: faster-whisper V3 Turbo（809M params, ~2-3GB RAM）
    - Whisper 的 CTranslate2 加速版，比原版快 6 倍
    - 支援 99+ 語言含中文
  TTS: Kokoro-82M（82M params, ~500MB RAM）
    - 超低延遲 <0.3s，36x 即時
    - 自然度接近商用 TTS

【為什麼不用雲端 TTS/STT？】
  跟整個 Phase C 的理念一致：零外部 API 依賴。
  語音在使用者裝置本地處理（隱私），ML service 只做轉換。

【安裝依賴】
  uv add faster-whisper kokoro soundfile

【注意】
  模型會在第一次呼叫時下載（~3GB），之後快取。
  faster-whisper 需要 ctranslate2，CPU 推論足夠（~2-5 秒/段語音）。

【ModelRegistry 整合】
  STT → stt capability → whisper-turbo model（provider: specialized, loader: faster_whisper）
  TTS → tts capability → kokoro-tts model（provider: specialized, loader: kokoro）
  _get_whisper_model() / _get_tts_model() 透過 registry.get() 取得已載入的模型。
  若 registry 未初始化，fallback 到舊版 lazy loading。
"""

from __future__ import annotations

import io
import logging
import os
import tempfile
from pathlib import Path

from pydantic import BaseModel

logger = logging.getLogger(__name__)

# ============================================================================
# STT — faster-whisper
# ============================================================================

# 保留舊版全域變數作為 fallback（registry 未啟動時）
_whisper_model = None
WHISPER_MODEL_SIZE = os.environ.get("WHISPER_MODEL_SIZE", "turbo")


class TranscribeResult(BaseModel):
    text: str
    language: str
    confidence: float


async def _get_whisper_model():
    """
    透過 ModelRegistry 取得 faster-whisper 模型。

    【Registry 優先，Lazy loading fallback】
    1. 嘗試從 registry 取得 stt capability 的已載入 WhisperModel
    2. 若 registry 未初始化或載入失敗，fallback 到舊版全域 lazy loading
    """
    global _whisper_model

    # 優先嘗試從 ModelRegistry 取得
    try:
        from src.registry import get_registry

        reg = get_registry()
        loaded = await reg.get("stt")
        if loaded.model is not None:
            # 更新全域快取，讓同步 transcribe() 也能直接使用 registry 載入的模型
            _whisper_model = loaded.model
            return loaded.model
    except Exception:
        # Registry 未初始化或 stt 模型載入失敗，走 fallback
        pass

    # Fallback：舊版 lazy loading
    if _whisper_model is not None:
        return _whisper_model

    try:
        from faster_whisper import WhisperModel

        logger.info(
            "Loading faster-whisper model (%s) [fallback]...", WHISPER_MODEL_SIZE
        )
        _whisper_model = WhisperModel(
            WHISPER_MODEL_SIZE,
            device="cpu",  # CPU 推論（MPS 不支援 ctranslate2）
            compute_type="int8",  # 量化加速（CPU 上 int8 最快）
        )
        logger.info("faster-whisper ready [fallback]")
        return _whisper_model
    except ImportError:
        logger.warning("faster-whisper not installed. Run: uv add faster-whisper")
        return None
    except Exception as e:
        logger.error("Failed to load faster-whisper: %s", e)
        return None


def _load_whisper():
    """
    同步版 whisper 載入（供 transcribe() 在非 async 呼叫環境使用）。
    備注：主要入口 transcribe() 會在 asyncio.to_thread 中呼叫，
    若在那之前 registry 已載入模型，此函式為 no-op。
    """
    global _whisper_model
    if _whisper_model is not None:
        return

    try:
        from faster_whisper import WhisperModel

        logger.info("Loading faster-whisper model (%s)...", WHISPER_MODEL_SIZE)
        _whisper_model = WhisperModel(
            WHISPER_MODEL_SIZE,
            device="cpu",
            compute_type="int8",
        )
        logger.info("faster-whisper ready")
    except ImportError:
        logger.warning("faster-whisper not installed. Run: uv add faster-whisper")
    except Exception as e:
        logger.error("Failed to load faster-whisper: %s", e)


def transcribe(audio_bytes: bytes, filename: str = "audio.webm") -> TranscribeResult:
    # 【學習重點：STT 的完整流程】
    # 瀏覽器錄音（MediaRecorder API）→ WebM/OGG bytes → 這裡
    # → faster-whisper 做 beam search 解碼 → 文字 + 語言偵測
    """
    語音轉文字。

    【流程】
    1. 嘗試從 _whisper_model 全域取得已載入模型
       （registry 會預先載入，或由同步 _load_whisper 延遲載入）
    2. 把音檔暫存到磁碟（faster-whisper 需要檔案路徑）
    3. 執行轉錄
    4. 回傳文字 + 偵測到的語言 + 信心分數

    注意：此函式設計為同步執行於 asyncio.to_thread 中。
    模型應由 _get_whisper_model() 或 startup event 預先載入。
    """
    # 嘗試使用全域模型（registry 啟動時預載，或 fallback lazy load）
    model = _whisper_model
    if model is None:
        _load_whisper()
        model = _whisper_model

    if model is None:
        return TranscribeResult(
            text="[STT not available]", language="unknown", confidence=0.0
        )

    # 暫存音檔
    suffix = Path(filename).suffix or ".webm"
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
        f.write(audio_bytes)
        tmp_path = f.name

    try:
        segments, info = model.transcribe(
            tmp_path,
            beam_size=5,
            language=None,  # 自動偵測語言
        )

        # 合併所有 segments 的文字
        text = " ".join(seg.text.strip() for seg in segments)

        return TranscribeResult(
            text=text,
            language=info.language,
            confidence=round(info.language_probability, 3),
        )
    finally:
        os.unlink(tmp_path)


# ============================================================================
# TTS — Kokoro
# ============================================================================

# 保留舊版全域變數作為 fallback
_tts_model = None


class SynthesizeResult(BaseModel):
    audio_base64: str
    sample_rate: int
    duration_seconds: float


async def _get_tts_model():
    """
    透過 ModelRegistry 取得 Kokoro TTS pipeline。

    【Registry 優先，Lazy loading fallback】
    1. 嘗試從 registry 取得 tts capability 的已載入 KPipeline
    2. 若 registry 未初始化或載入失敗，fallback 到舊版全域 lazy loading
    """
    global _tts_model

    # 優先嘗試從 ModelRegistry 取得
    try:
        from src.registry import get_registry

        reg = get_registry()
        loaded = await reg.get("tts")
        if loaded.model is not None:
            # 更新全域快取，讓同步 synthesize() 也能直接使用 registry 載入的模型
            _tts_model = loaded.model
            return loaded.model
    except Exception:
        # Registry 未初始化或 tts 模型載入失敗，走 fallback
        pass

    # Fallback：舊版 lazy loading
    if _tts_model is not None:
        return _tts_model

    try:
        import kokoro

        logger.info("Loading Kokoro TTS model [fallback]...")
        _tts_model = kokoro.KPipeline(lang_code="a")
        logger.info("Kokoro TTS ready [fallback]")
        return _tts_model
    except ImportError:
        logger.warning("kokoro not installed. Run: uv add kokoro soundfile")
        return None
    except Exception as e:
        logger.error("Failed to load Kokoro TTS: %s", e)
        return None


def _load_tts():
    """
    同步版 Kokoro 載入（備用，供非 async 環境使用）。

    【Kokoro-82M 特點】
    - 只有 82M 參數（~500MB），載入極快
    - 延遲 <0.3s，即時度 36x
    - 支援英文，中文需要額外配置
    """
    global _tts_model
    if _tts_model is not None:
        return

    try:
        import kokoro

        logger.info("Loading Kokoro TTS model...")
        _tts_model = kokoro.KPipeline(lang_code="a")  # 'a' = American English
        logger.info("Kokoro TTS ready")
    except ImportError:
        logger.warning("kokoro not installed. Run: uv add kokoro soundfile")
    except Exception as e:
        logger.error("Failed to load Kokoro TTS: %s", e)


def synthesize(text: str, voice: str = "af_heart") -> SynthesizeResult | None:
    # 【學習重點：TTS 的完整流程】
    # 文字 → Kokoro 生成 audio tensor（numpy array）
    # → soundfile 寫成 WAV bytes → base64 編碼 → JSON 回傳
    # → 前端用 Audio API 播放 base64 audio
    """
    文字轉語音。

    【流程】
    1. 嘗試從 _tts_model 全域取得已載入模型
       （registry 預先載入，或由 _load_tts fallback 延遲載入）
    2. 用 Kokoro 生成音頻 tensor
    3. 轉成 WAV bytes → base64 編碼
    4. 回傳 base64 audio + metadata

    【voice 選項】
    - af_heart: 女聲（預設，自然）
    - am_adam: 男聲

    注意：此函式設計為同步執行於 asyncio.to_thread 中。
    """
    model = _tts_model
    if model is None:
        _load_tts()
        model = _tts_model

    if model is None:
        return None

    import base64
    import soundfile as sf
    import numpy as np

    try:
        # 生成音頻
        generator = model(text, voice=voice)
        audio_segments = []
        sample_rate = 24000  # Kokoro 預設 24kHz

        for _gs, _ps, audio in generator:
            if audio is not None:
                audio_segments.append(audio)

        if not audio_segments:
            return None

        # 合併所有 segments
        full_audio = (
            np.concatenate(audio_segments)
            if len(audio_segments) > 1
            else audio_segments[0]
        )

        # 轉成 WAV bytes
        buf = io.BytesIO()
        sf.write(buf, full_audio, sample_rate, format="WAV")
        wav_bytes = buf.getvalue()

        duration = len(full_audio) / sample_rate

        return SynthesizeResult(
            audio_base64=base64.b64encode(wav_bytes).decode("ascii"),
            sample_rate=sample_rate,
            duration_seconds=round(duration, 2),
        )
    except Exception as e:
        logger.error("TTS synthesis failed: %s", e)
        return None