Trixy7/plugins/tts_piper/main.py

# -*- coding: utf-8 -*-
"""
Piper TTS Plugin - Text-to-Speech mit Piper.

Piper ist ein schnelles, lokales TTS-System basierend auf VITS.

Unterstützt:
- Lokale Synthese (offline)
- Sehr schnell (Echtzeit auf CPU)
- Hochwertige Stimmen
- Mehrere Sprachen
- Auto-Download von Hugging Face
"""

import asyncio
import io
import struct
import time
import wave
from pathlib import Path
from typing import Any, AsyncIterator

from trixy_core.plugins import TrixyPlugin
from trixy_core.audio.tts import TTSProvider, TTSConfig, TTSResult, TTSState, Voice


# Piper Voice Katalog mit Hugging Face Pfaden
PIPER_VOICES = {
    # Deutsche Stimmen
    "de_DE-thorsten-high": {
        "name": "Thorsten (High)",
        "language": "de-DE",
        "gender": "male",
        "description": "Deutsche männliche Stimme, hohe Qualität",
        "sample_rate": 22050,
        "hf_path": "de/de_DE/thorsten/high/de_DE-thorsten-high.onnx",
    },
    "de_DE-thorsten-medium": {
        "name": "Thorsten (Medium)",
        "language": "de-DE",
        "gender": "male",
        "description": "Deutsche männliche Stimme, mittlere Qualität",
        "sample_rate": 22050,
        "hf_path": "de/de_DE/thorsten/medium/de_DE-thorsten-medium.onnx",
    },
    "de_DE-thorsten-low": {
        "name": "Thorsten (Low)",
        "language": "de-DE",
        "gender": "male",
        "description": "Deutsche männliche Stimme, niedrige Qualität (schnell)",
        "sample_rate": 22050,
        "hf_path": "de/de_DE/thorsten/low/de_DE-thorsten-low.onnx",
    },
    "de_DE-eva_k-x_low": {
        "name": "Eva",
        "language": "de-DE",
        "gender": "female",
        "description": "Deutsche weibliche Stimme",
        "sample_rate": 22050,
        "hf_path": "de/de_DE/eva_k/x_low/de_DE-eva_k-x_low.onnx",
    },
    "de_DE-kerstin-low": {
        "name": "Kerstin",
        "language": "de-DE",
        "gender": "female",
        "description": "Deutsche weibliche Stimme",
        "sample_rate": 16000,
        "hf_path": "de/de_DE/kerstin/low/de_DE-kerstin-low.onnx",
    },
    "de_DE-ramona-low": {
        "name": "Ramona",
        "language": "de-DE",
        "gender": "female",
        "description": "Deutsche weibliche Stimme",
        "sample_rate": 16000,
        "hf_path": "de/de_DE/ramona/low/de_DE-ramona-low.onnx",
    },
    "de_DE-karlsson-low": {
        "name": "Karlsson",
        "language": "de-DE",
        "gender": "male",
        "description": "Deutsche männliche Stimme",
        "sample_rate": 16000,
        "hf_path": "de/de_DE/karlsson/low/de_DE-karlsson-low.onnx",
    },
    "de_DE-pavoque-low": {
        "name": "Pavoque",
        "language": "de-DE",
        "gender": "male",
        "description": "Deutsche männliche Stimme",
        "sample_rate": 16000,
        "hf_path": "de/de_DE/pavoque/low/de_DE-pavoque-low.onnx",
    },
    "de_DE-mls-medium": {
        "name": "MLS",
        "language": "de-DE",
        "gender": "male",
        "description": "Deutsche männliche Stimme (Multi-Speaker)",
        "sample_rate": 22050,
        "hf_path": "de/de_DE/mls/medium/de_DE-mls-medium.onnx",
    },
    "de_DE-thorsten_emotional-medium": {
        "name": "Thorsten Emotional",
        "language": "de-DE",
        "gender": "male",
        "description": "Deutsche männliche Stimme mit Emotionen",
        "sample_rate": 22050,
        "hf_path": "de/de_DE/thorsten_emotional/medium/de_DE-thorsten_emotional-medium.onnx",
    },
    # Englische Stimmen
    "en_US-lessac-medium": {
        "name": "Lessac (Medium)",
        "language": "en-US",
        "gender": "female",
        "description": "Englische weibliche Stimme (US)",
        "sample_rate": 22050,
        "hf_path": "en/en_US/lessac/medium/en_US-lessac-medium.onnx",
    },
    "en_US-amy-medium": {
        "name": "Amy (Medium)",
        "language": "en-US",
        "gender": "female",
        "description": "Englische weibliche Stimme (US)",
        "sample_rate": 22050,
        "hf_path": "en/en_US/amy/medium/en_US-amy-medium.onnx",
    },
    "en_GB-alan-medium": {
        "name": "Alan (Medium)",
        "language": "en-GB",
        "gender": "male",
        "description": "Englische männliche Stimme (UK)",
        "sample_rate": 22050,
        "hf_path": "en/en_GB/alan/medium/en_GB-alan-medium.onnx",
    },
}

HUGGINGFACE_REPO = "rhasspy/piper-voices"


def _resample_wav(wav_data: bytes, source_rate: int, target_rate: int) -> bytes:
    """Resampled WAV-Daten von source_rate auf target_rate.

    Verwendet lineare Interpolation (schnell, gute Qualitaet fuer Sprache).
    """
    # PCM aus WAV extrahieren (44 Bytes Header)
    pcm = wav_data[44:]
    # int16 → float
    n_samples = len(pcm) // 2
    samples = struct.unpack(f"<{n_samples}h", pcm)

    # Resample-Faktor
    ratio = target_rate / source_rate
    new_length = int(n_samples * ratio)

    # Lineare Interpolation
    resampled = []
    for i in range(new_length):
        src_pos = i / ratio
        idx = int(src_pos)
        frac = src_pos - idx
        if idx + 1 < n_samples:
            val = samples[idx] * (1.0 - frac) + samples[idx + 1] * frac
        elif idx < n_samples:
            val = samples[idx]
        else:
            val = 0
        resampled.append(int(max(-32768, min(32767, val))))

    # Neues WAV bauen
    out = io.BytesIO()
    with wave.open(out, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(target_rate)
        wf.writeframes(struct.pack(f"<{len(resampled)}h", *resampled))

    return out.getvalue()


class PiperTTSProvider(TTSProvider):
    """Piper-basierter TTS-Provider."""

    def __init__(
        self,
        config: TTSConfig | None = None,
        models_dir: Path | None = None,
        voice_id: str = "de_DE-thorsten-medium",
        auto_download: bool = True,
        event_manager: Any = None,
        shared_models_dir: Path | None = None,
    ):
        super().__init__(config)
        self._models_dir = models_dir            # plugins/tts_piper/models/
        self._shared_models_dir = shared_models_dir  # ./models/piper/ (persistent)
        self._voice_id = voice_id
        self._auto_download = auto_download
        self._event_manager = event_manager
        self._piper = None
        self._model_path: Path | None = None

    def _search_dirs(self) -> list[Path]:
        """Gibt alle Verzeichnisse zurueck in denen nach Modellen gesucht wird."""
        dirs: list[Path] = []
        if self._shared_models_dir and self._shared_models_dir.exists():
            dirs.append(self._shared_models_dir)
        if self._models_dir:
            dirs.append(self._models_dir)
        return dirs

    def _find_model_file(self, voice_id: str) -> Path | None:
        """Sucht eine .onnx-Datei fuer die gegebene voice_id in allen Verzeichnissen."""
        filename = f"{voice_id}.onnx"
        for d in self._search_dirs():
            candidate = d / filename
            if candidate.exists():
                return candidate
        return None

    @property
    def name(self) -> str:
        return "piper"

    @property
    def supported_languages(self) -> list[str]:
        return ["de-DE", "en-US", "en-GB", "fr-FR", "es-ES", "it-IT", "nl-NL", "pl-PL"]

    @property
    def supports_streaming(self) -> bool:
        return True

    async def initialize(self) -> None:
        """Lädt das Piper-Modell."""
        try:
            from piper import PiperVoice

            self._state = TTSState.SYNTHESIZING

            # Lokale Modelle in allen Such-Verzeichnissen automatisch registrieren
            self._discover_local_models()

            # Modell-Pfad bestimmen: zuerst in allen Such-Verzeichnissen
            # (./models/piper/ hat Vorrang vor plugins/tts_piper/models/)
            existing = self._find_model_file(self._voice_id)
            if existing:
                self._model_path = existing
            else:
                # Nicht lokal vorhanden — Download-Pfad im Plugin-Ordner
                self._model_path = self._models_dir / f"{self._voice_id}.onnx"

            # Auto-Download wenn nicht vorhanden
            if not self._model_path.exists():
                if self._auto_download:
                    await self._download_model()
                else:
                    raise RuntimeError(
                        f"Piper-Modell nicht gefunden: {self._model_path}. "
                        "Auto-Download ist deaktiviert."
                    )

            # Piper laden (in Thread)
            loop = asyncio.get_event_loop()
            self._piper = await loop.run_in_executor(
                None,
                lambda: PiperVoice.load(str(self._model_path))
            )

            # Aktuelle Stimme setzen
            if self._voice_id in PIPER_VOICES:
                voice_info = PIPER_VOICES[self._voice_id]
                self._current_voice = Voice(
                    id=self._voice_id,
                    name=voice_info["name"],
                    language=voice_info["language"],
                    gender=voice_info["gender"],
                    description=voice_info["description"],
                    sample_rate=voice_info["sample_rate"],
                )

            self._model_loaded = True
            self._state = TTSState.READY

        except ImportError:
            raise RuntimeError(
                "Piper nicht installiert. "
                "Installieren mit: pip install piper-tts"
            )
        except Exception as e:
            self._state = TTSState.ERROR
            raise RuntimeError(f"Fehler beim Laden des Piper-Modells: {e}")

    async def _download_model(self) -> None:
        """Lädt das Modell von Hugging Face herunter."""
        from trixy_core.utils.debug import pinfo
        from trixy_core.utils.download import download_from_huggingface

        if self._voice_id not in PIPER_VOICES:
            raise RuntimeError(
                f"Unbekannte Stimme: {self._voice_id}. "
                f"Verfügbar: {', '.join(PIPER_VOICES.keys())}"
            )

        voice_info = PIPER_VOICES[self._voice_id]
        hf_path = voice_info["hf_path"]
        hf_json_path = hf_path + ".json"

        self._models_dir.mkdir(parents=True, exist_ok=True)

        # ONNX-Modell herunterladen
        pinfo(f"Lade Piper-Modell: {self._voice_id}...")

        model_dest = self._models_dir / f"{self._voice_id}.onnx"
        success = await download_from_huggingface(
            repo_id=HUGGINGFACE_REPO,
            filename=hf_path,
            dest_path=model_dest,
            event_manager=self._event_manager,
            download_id=f"piper-{self._voice_id}",
        )

        if not success:
            raise RuntimeError(f"Download des Modells fehlgeschlagen: {self._voice_id}")

        # JSON-Config herunterladen
        json_dest = self._models_dir / f"{self._voice_id}.onnx.json"
        await download_from_huggingface(
            repo_id=HUGGINGFACE_REPO,
            filename=hf_json_path,
            dest_path=json_dest,
            event_manager=self._event_manager,
            download_id=f"piper-{self._voice_id}-json",
        )

        pinfo(f"Piper-Modell heruntergeladen: {self._voice_id}")

    def _discover_local_models(self) -> None:
        """Erkennt lokale Modelle in allen Such-Verzeichnissen.

        Scannt sowohl plugins/tts_piper/models/ als auch (falls gesetzt)
        ./models/piper/. Liest .onnx.json fuer Metadaten.
        Bereits in PIPER_VOICES eingetragene Stimmen werden nicht ueberschrieben.
        """
        from trixy_core.utils.debug import pinfo
        import json as _json

        discovered = 0
        seen_voices: set[str] = set()

        for models_dir in self._search_dirs():
            for onnx_file in sorted(models_dir.glob("*.onnx")):
                voice_id = onnx_file.stem  # z.B. "de_DE-Feli-medium"
                if voice_id in PIPER_VOICES or voice_id in seen_voices:
                    continue  # Bereits bekannt

                json_file = onnx_file.with_suffix(".onnx.json")
                if not json_file.exists():
                    continue  # Ohne Config nicht nutzbar

                try:
                    with open(json_file, "r", encoding="utf-8") as f:
                        meta = _json.load(f)
                except Exception:
                    continue

                audio = meta.get("audio", {})
                language = meta.get("language", {})
                dataset = meta.get("dataset", "")

                # Stimmenname aus voice_id extrahieren: "de_DE-Feli-medium" → "Feli"
                parts = voice_id.split("-")
                if len(parts) >= 2:
                    name = parts[1].replace("_", " ").title()
                else:
                    name = voice_id

                quality = parts[-1] if len(parts) >= 3 else "unknown"

                lang_code = language.get("code", "")
                if lang_code:
                    lang_formatted = lang_code.replace("_", "-")
                else:
                    lang_formatted = parts[0].replace("_", "-") if parts else ""

                sample_rate = audio.get("sample_rate", 22050)

                PIPER_VOICES[voice_id] = {
                    "name": f"{name} ({quality.title()})",
                    "language": lang_formatted,
                    "gender": "unknown",
                    "description": f"Lokales Modell: {dataset or name}",
                    "sample_rate": sample_rate,
                    "hf_path": "",
                }
                seen_voices.add(voice_id)
                discovered += 1

        if discovered > 0:
            pinfo(f"Piper TTS: {discovered} lokale Modelle erkannt")

    async def shutdown(self) -> None:
        """Gibt Ressourcen frei."""
        self._piper = None
        self._model_loaded = False
        self._state = TTSState.UNINITIALIZED

    async def get_voices(self, language: str | None = None) -> list[Voice]:
        """Gibt verfügbare Stimmen zurück."""
        voices = []
        for voice_id, info in PIPER_VOICES.items():
            voices.append(Voice(
                id=voice_id,
                name=info["name"],
                language=info["language"],
                gender=info["gender"],
                description=info["description"],
                sample_rate=info["sample_rate"],
            ))

        if language:
            lang_prefix = language.replace("-", "_").split("_")[0]
            voices = [v for v in voices if v.language.replace("-", "_").startswith(lang_prefix)]

        return voices

    async def synthesize(
        self,
        text: str,
        voice_id: str | None = None,
        language: str | None = None,
    ) -> TTSResult:
        """Synthetisiert Text mit Piper."""
        if not self._piper:
            raise RuntimeError("Piper nicht geladen")

        self._state = TTSState.SYNTHESIZING
        start_time = time.time()

        try:
            # Audio-Buffer
            audio_buffer = io.BytesIO()

            # WAV-Writer
            with wave.open(audio_buffer, "wb") as wav_file:
                wav_file.setnchannels(1)
                wav_file.setsampwidth(2)  # 16-bit
                wav_file.setframerate(self._piper.config.sample_rate)

                # Synthese in Thread
                loop = asyncio.get_event_loop()

                def do_synthesize():
                    for audio_chunk in self._piper.synthesize(text):
                        wav_file.writeframes(audio_chunk.audio_int16_bytes)

                await loop.run_in_executor(None, do_synthesize)

            processing_time = (time.time() - start_time) * 1000
            audio_data = audio_buffer.getvalue()

            self._state = TTSState.READY

            source_rate = self._piper.config.sample_rate
            target_rate = 22050  # Standard-Ausgaberate fuer TTS-Socket

            # Resampling wenn Modell-Rate != Ausgabe-Rate
            if source_rate != target_rate:
                audio_data = _resample_wav(audio_data, source_rate, target_rate)

            # Dauer berechnen
            sample_rate = target_rate
            # WAV hat 44 bytes Header
            pcm_data = audio_data[44:]
            num_samples = len(pcm_data) // 2  # 16-bit = 2 bytes
            duration = num_samples / sample_rate

            return TTSResult(
                audio_data=audio_data,
                sample_rate=sample_rate,
                channels=1,
                text=text,
                duration_seconds=duration,
                processing_time_ms=processing_time,
                voice_id=voice_id or self._voice_id,
                voice_name=self._current_voice.name if self._current_voice else "",
                language=language or self._config.language,
                provider="piper",
                model=self._voice_id,
            )

        except Exception as e:
            self._state = TTSState.ERROR
            raise RuntimeError(f"Piper-Synthese fehlgeschlagen: {e}")

    async def synthesize_stream(
        self,
        text: str,
        voice_id: str | None = None,
        language: str | None = None,
        chunk_size: int = 4096,
    ) -> AsyncIterator[bytes]:
        """Synthetisiert Text als Stream."""
        if not self._piper:
            raise RuntimeError("Piper nicht geladen")

        self._state = TTSState.SYNTHESIZING

        try:
            loop = asyncio.get_event_loop()

            # Generator für Chunks
            def generate_chunks():
                for audio_chunk in self._piper.synthesize(text):
                    audio_bytes = audio_chunk.audio_int16_bytes
                    # In Chunks aufteilen
                    for i in range(0, len(audio_bytes), chunk_size):
                        yield audio_bytes[i:i + chunk_size]

            # Chunks yielden
            chunks = await loop.run_in_executor(None, lambda: list(generate_chunks()))

            for chunk in chunks:
                yield chunk

            self._state = TTSState.READY

        except Exception as e:
            self._state = TTSState.ERROR
            raise RuntimeError(f"Piper-Stream fehlgeschlagen: {e}")


class PiperTTSPlugin(TrixyPlugin):
    """Piper TTS Plugin für Trixy."""

    def __init__(self, application, plugin_path, config: dict | None = None):
        super().__init__(application, plugin_path, config)
        self._provider: PiperTTSProvider | None = None

    async def on_load(self) -> None:
        """Plugin wird geladen."""
        from trixy_core.utils.debug import pinfo
        pinfo("Piper TTS Plugin: Lade...")

        # Konfiguration
        voice_id = self.config.get("voice", "de_DE-thorsten-medium")
        language = self.config.get("language", "de-DE")
        auto_download = self.config.get("auto_download", True)

        # Models-Verzeichnis im Plugin-Ordner (fuer Standard-Downloads)
        models_dir = self.plugin_path / "models"
        models_dir.mkdir(parents=True, exist_ok=True)

        # Shared Models-Verzeichnis: projekt-weit fuer Custom-Voices
        # ./models/piper/ — wird nicht vom Installer ueberschrieben
        # plugin_path ist z.B. /home/pi/trixy/source/plugins/tts_piper
        # parents[1] → /home/pi/trixy/source
        shared_models_dir = self.plugin_path.parents[1] / "models" / "piper"

        tts_config = TTSConfig(
            language=language,
            voice_id=voice_id,
        )

        # Provider erstellen (mit EventManager für Download-Events)
        self._provider = PiperTTSProvider(
            tts_config,
            models_dir=models_dir,
            shared_models_dir=shared_models_dir,
            voice_id=voice_id,
            auto_download=auto_download,
            event_manager=self.application.events,
        )

        # Modell laden (und ggf. downloaden)
        await self._provider.initialize()

        # Extension registrieren
        if hasattr(self.application, "extension_points"):
            ext_point = self.application.extension_points.get("conversation.tts")
            if ext_point:
                ext_point.register(self._provider)
                pinfo("Piper TTS Plugin: Extension registriert")

        # Event-Handler registrieren
        self._register_event_handlers()

        pinfo(f"Piper TTS Plugin: Geladen (Stimme: {voice_id})")

    def _register_event_handlers(self) -> None:
        """Registriert Event-Handler."""
        em = self.application.events

        @em.on("tts_request")
        async def on_tts_request(event_name: str, data: dict) -> None:
            """Verarbeitet TTS-Anfragen."""
            if not self._provider or not self._provider.is_ready:
                return

            text = data.get("text")
            satellite_id = data.get("satellite_id")
            session_id = data.get("session_id")
            voice_id = data.get("voice_id")
            request_id = data.get("request_id", "")

            if not text:
                return

            from trixy_core.utils.debug import pinfo
            pinfo(f"Piper TTS: Synthetisiere '{text[:50]}...'")

            try:
                result = await self._provider.synthesize(text, voice_id=voice_id)

                await em.emit("tts_completed", {
                    "request_id": request_id,
                    "audio_data": result.audio_data.hex(),
                    "sample_rate": result.sample_rate,
                    "duration_seconds": result.duration_seconds,
                    "provider": "piper",
                    "text": text,
                    "satellite_id": satellite_id,
                    "session_id": session_id,
                    "processing_time_ms": result.processing_time_ms,
                })

                pinfo(f"Piper TTS: Fertig ({result.duration_seconds:.1f}s Audio, "
                      f"{result.processing_time_ms:.0f}ms)")

            except Exception as e:
                from trixy_core.utils.debug import perror
                perror(f"Piper TTS Fehler: {e}")

                await em.emit("tts_error", {
                    "error": str(e),
                    "provider": "piper",
                    "text": text,
                    "satellite_id": satellite_id,
                    "session_id": session_id,
                })

    async def on_unload(self) -> None:
        """Plugin wird entladen."""
        if self._provider:
            await self._provider.shutdown()
            self._provider = None

        from trixy_core.utils.debug import pinfo
        pinfo("Piper TTS Plugin: Entladen")

    @property
    def provider(self) -> PiperTTSProvider | None:
        """TTS-Provider."""
        return self._provider


# Plugin-Export
Plugin = PiperTTSPlugin