app/services/tts.py

import json
import os
import re
import time
from pathlib import Path

import numpy as np

from app.config import Config
from app.logging_config import get_logger

logger = get_logger('tts')

TTS = None
STYLE_CLS = None


def _ensure_supertonic():
    global TTS, STYLE_CLS
    if TTS is None:
        try:
            from supertonic import TTS as _TTS
            from supertonic.core import Style as _Style
            TTS = _TTS
            STYLE_CLS = _Style
        except ImportError as exc:
            raise ImportError('supertonic not found. Install with: pip install supertonic') from exc


class TTSService:
    def __init__(self):
        self._tts = None
        self._model_loaded = False
        self._voices_dir: str | None = None
        self._voice_presets_dir: str | None = None
        self._preset_voices: dict[str, str] = {}
        self._style_cache: dict[str, object] = {}
        # Pre-parse text into segments for simulated streaming
        self._sentence_splitter = re.compile(r'(?<=[.!?])\s+')

    @property
    def is_loaded(self) -> bool:
        return self._model_loaded and self._tts is not None

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

    @property
    def device(self) -> str:
        return 'cpu'

    def load_model(self) -> None:
        _ensure_supertonic()

        logger.info('Loading SuperTonic3 model (auto-download)...')
        t0 = time.time()

        try:
            self._tts = TTS(auto_download=Config.AUTO_DOWNLOAD)
            self._model_loaded = True
            load_time = time.time() - t0
            logger.info(f'SuperTonic3 model loaded in {load_time:.2f}s')
        except Exception as e:
            logger.error(f'Failed to load SuperTonic3 model: {e}')
            raise

        self.discover_presets()

    def set_voices_dir(self, voices_dir: str | None) -> None:
        if voices_dir and os.path.isdir(voices_dir):
            self._voices_dir = voices_dir
            logger.info(f'Voices directory set to: {voices_dir}')
        elif voices_dir:
            logger.warning(f'Voices directory not found: {voices_dir}')
            self._voices_dir = None
        else:
            self._voices_dir = None

    def discover_presets(self) -> None:
        self._preset_voices = {}
        presets_dir = Config.VOICE_PRESETS_DIR
        if not os.path.isdir(presets_dir):
            logger.info(f'Voice presets directory not found: {presets_dir}')
            return
        self._voice_presets_dir = presets_dir
        for f in sorted(Path(presets_dir).iterdir()):
            if f.suffix.lower() in Config.VOICE_PRESET_EXTENSIONS and f.stem.isidentifier():
                self._preset_voices[f.stem] = str(f)
        if self._preset_voices:
            logger.info(
                f'Discovered {len(self._preset_voices)} voice presets: '
                f'{", ".join(self._preset_voices.keys())}'
            )

    def get_voice_style(self, voice_id: str):
        if not self.is_loaded:
            raise RuntimeError('Model not loaded. Call load_model() first.')

        # Check if it's a cached preset voice
        if voice_id in self._style_cache:
            return self._style_cache[voice_id]

        # Check if it's a known preset voice (load from JSON)
        if voice_id in self._preset_voices:
            return self._load_preset_style(voice_id)

        voice_name = self._resolve_voice(voice_id)
        t0 = time.time()
        style = self._tts.get_voice_style(voice_name=voice_name)
        logger.debug(f'Voice style loaded in {time.time() - t0:.2f}s: {voice_name}')
        return style

    def _load_preset_style(self, voice_id: str):
        filepath = self._preset_voices.get(voice_id)
        if not filepath:
            raise ValueError(f'Preset voice not found: {voice_id}')

        t0 = time.time()
        with open(filepath) as f:
            data = json.load(f)

        style_ttl = np.array(data['style_ttl']['data'], dtype=data['style_ttl']['type']).reshape(
            data['style_ttl']['dims']
        )
        style_dp = np.array(data['style_dp']['data'], dtype=data['style_dp']['type']).reshape(
            data['style_dp']['dims']
        )

        style = STYLE_CLS(style_ttl_onnx=style_ttl, style_dp_onnx=style_dp)
        self._style_cache[voice_id] = style
        logger.info(f'Loaded preset voice {voice_id} in {time.time() - t0:.2f}s')
        return style

    def _resolve_voice(self, voice_id: str) -> str:
        voice_lower = voice_id.lower()
        if voice_lower in [v.lower() for v in Config.BUILTIN_VOICES]:
            return next(v for v in Config.BUILTIN_VOICES if v.lower() == voice_lower)

        if voice_id.startswith(('http://', 'https://')):
            raise ValueError(
                f'URL scheme not allowed for security reasons: {voice_id[:50]}. '
                "Use 'hf://' for HuggingFace models or a local file path."
            )

        if self._voices_dir:
            p = Path(self._voices_dir) / voice_id
            if p.exists():
                return str(p)

        p = Path(voice_id)
        if p.exists():
            return str(p)

        return voice_id

    def validate_voice(self, voice_id: str) -> tuple[bool, str]:
        if voice_id.startswith(('http://', 'https://')):
            return (
                False,
                'HTTP/HTTPS URLs are not allowed for security reasons. Use hf:// for HuggingFace models.',
            )
        # Check built-in voices
        voice_lower = voice_id.lower()
        if voice_lower in [v.lower() for v in Config.BUILTIN_VOICES]:
            return True, f'Built-in voice: {voice_id}'
        # Check preset voices (JSON style embeddings)
        if voice_id in self._preset_voices:
            return True, f'Preset voice: {voice_id}'
        if self._voices_dir:
            p = Path(self._voices_dir) / voice_id
            if p.exists():
                return True, f'Local voice: {p}'
        p = Path(voice_id)
        if p.exists():
            return True, f'Local voice: {p}'
        try:
            self._tts.get_voice_style(voice_name=voice_id)
            return True, f'Voice: {voice_id}'
        except Exception:
            return False, f'Voice not found: {voice_id}'

    def generate_audio(self, voice_style, text: str, lang: str = 'en'):
        if not self.is_loaded:
            raise RuntimeError('Model not loaded')

        t0 = time.time()
        audio, duration = self._tts.synthesize(text=text, voice_style=voice_style, lang=lang)
        gen_time = time.time() - t0
        dur = float(duration.item() if hasattr(duration, 'item') else duration)
        logger.info(f'Generated {len(text)} chars in {gen_time:.2f}s (audio: {dur:.2f}s)')
        return audio if audio.ndim == 1 else audio[0]

    def generate_audio_stream(self, voice_style, text: str, lang: str = 'en'):
        if not self.is_loaded:
            raise RuntimeError('Model not loaded')

        import numpy as np

        sentences = self._sentence_splitter.split(text)
        if not sentences:
            sentences = [text]

        logger.info(f'Starting streaming generation for {len(text)} chars ({len(sentences)} segments)')

        for i, segment in enumerate(sentences):
            if not segment.strip():
                continue
            t0 = time.time()
            chunk, _ = self._tts.synthesize(text=segment, voice_style=voice_style, lang=lang)
            gen_time = time.time() - t0
            logger.debug(f'Stream chunk {i + 1}/{len(sentences)} in {gen_time:.2f}s')
            yield (chunk[0] if chunk.ndim > 1 else chunk).astype(np.float32)

    def list_voices(self) -> list[dict]:
        voices: list[dict] = []
        for name in Config.BUILTIN_VOICES:
            voices.append({'id': name, 'name': name, 'type': 'builtin'})
        for name in sorted(self._preset_voices):
            voices.append({'id': name, 'name': name, 'type': 'preset'})
        if self._voices_dir:
            for f in sorted(Path(self._voices_dir).iterdir()):
                if f.suffix.lower() in Config.VOICE_EXTENSIONS:
                    stem = f.stem
                    clean_name = stem.replace('_', ' ').replace('-', ' ').title()
                    voices.append({'id': stem, 'name': clean_name, 'type': 'custom'})
        return voices


_tts_service: TTSService | None = None


def get_tts_service() -> TTSService:
    global _tts_service
    if _tts_service is None:
        _tts_service = TTSService()
    return _tts_service