"""
Custom handler pour Hugging Face Inference Endpoints — assistant vocal
WAGADU, synthèse wolof (bilalfaye/speecht5_tts-wolof).

À déployer : ce fichier doit être uploadé sous le nom EXACT "handler.py"
à la racine d'une copie du modèle sur votre compte Hugging Face (le
modèle original n'appartenant pas à Wagadu, on ne peut pas y ajouter de
fichier directement — voir les instructions de déploiement fournies par
Claude).

Reproduit exactement la logique déjà validée dans
backend/apps/communication/voice_assistant.py (même modèle, même
embedding de voix — speaker index 7306, "slt", voix féminine neutre) —
seule différence : tourne sur GPU (endpoint dédié) au lieu du CPU du
worker Render, ce qui doit réduire le temps de synthèse de 49-188s à
quelques secondes (décodeur autoregressif, largement accéléré par un
GPU par rapport à un seul cœur CPU).
"""

import base64
import io
from typing import Any, Dict


class EndpointHandler:
    def __init__(self, path: str = ""):
        import torch
        from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor

        self._device = "cuda" if torch.cuda.is_available() else "cpu"

        self.processor = SpeechT5Processor.from_pretrained(path)
        self.model = SpeechT5ForTextToSpeech.from_pretrained(path).to(self._device)
        self.model.eval()
        self.vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan").to(self._device)
        self.speaker_embedding = self._load_speaker_embedding().to(self._device)

    def _load_speaker_embedding(self):
        """Même méthode que _load_speaker_embedding() dans
        voice_assistant.py : contourne datasets.load_dataset (cassé pour
        ce dataset — format "script" non supporté par datasets>=4) en
        téléchargeant directement l'archive et en extrayant le fichier
        voulu, indexé exactement comme le script de chargement d'origine."""
        import zipfile

        import numpy as np
        import torch
        from huggingface_hub import hf_hub_download

        zip_path = hf_hub_download("Matthijs/cmu-arctic-xvectors", "spkrec-xvect.zip", repo_type="dataset")
        with zipfile.ZipFile(zip_path) as zf:
            names = sorted(n for n in zf.namelist() if n.endswith(".npy"))
            with zf.open(names[7306]) as f:
                return torch.tensor(np.load(f)).unsqueeze(0)

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        import numpy as np
        import scipy.io.wavfile
        import torch

        text = (data.get("inputs") or "").strip()
        if not text:
            return {"error": "champ 'inputs' requis (texte wolof à synthétiser)"}

        inputs = self.processor(
            text=text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.model.config.max_text_positions,
        )
        inputs = {k: v.to(self._device) for k, v in inputs.items()}

        # Réglages recommandés par l'auteur du modèle (fiche Hugging Face,
        # section "Model Loading and Speech Generation Code") : generate()
        # avec recherche en faisceau plutôt que generate_speech() (glouton,
        # utilisé côté Render CPU où la vitesse prime) — meilleure qualité,
        # mais nettement plus lent. Se justifie uniquement ici, sur GPU, où
        # on a la marge de calcul pour se le permettre.
        max_length = int(self.model.config.max_length * 1.2)
        with torch.no_grad():
            waveform = self.model.generate(
                inputs["input_ids"],
                speaker_embeddings=self.speaker_embedding,
                vocoder=self.vocoder,
                max_length=max_length,
                min_length=max_length // 3,
                num_beams=7,
                temperature=0.6,
                no_repeat_ngram_size=3,
                repetition_penalty=1.5,
                eos_token_id=None,
                use_cache=True,
            )
        waveform = waveform.detach().cpu().numpy()

        # PCM 16 bits — même format que voice_assistant.py côté Render,
        # pour que le reste du pipeline (Celery, polling, lecture client)
        # n'ait besoin d'aucun changement.
        pcm = np.clip(waveform, -1.0, 1.0)
        pcm = (pcm * 32767).astype(np.int16)

        buf = io.BytesIO()
        scipy.io.wavfile.write(buf, rate=16000, data=pcm)

        return {"audio_base64": base64.b64encode(buf.getvalue()).decode("ascii")}