"""SpeechT5 TTS implementation using HuggingFace transformers."""

from __future__ import annotations

import numpy as np
import torch
from roohai.base import TTSModel


class SpeechT5TTS(TTSModel):
    META = {
        "display_name": "SpeechT5",
        "description": "Microsoft's unified speech-text model for text-to-speech. Runs locally.",
        "kind": "local",
        "pip_dependencies": ["transformers>=4.46.0,<5.0.0", "torch", "sentencepiece"],
        "hf_models": ["microsoft/speecht5_tts", "microsoft/speecht5_hifigan"],
        "estimated_download_mb": 600,
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "microsoft/speecht5_tts",
            },
            {
                "name": "vocoder",
                "label": "Vocoder",
                "type": "text",
                "required": False,
                "default": "microsoft/speecht5_hifigan",
            },
        ],
    }

    def __init__(
        self,
        model_id: str = "microsoft/speecht5_tts",
        vocoder: str = "microsoft/speecht5_hifigan",
    ) -> None:
        self.model_id = model_id
        self.vocoder_id = vocoder
        self._processor = None
        self._model = None
        self._vocoder = None
        self._speaker_embedding = None

    def load(self) -> None:
        from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan

        self._processor = SpeechT5Processor.from_pretrained(self.model_id)
        self._model = SpeechT5ForTextToSpeech.from_pretrained(self.model_id)
        self._vocoder = SpeechT5HifiGan.from_pretrained(self.vocoder_id)
        # Use fixed random speaker embedding instead of broken cmu-arctic-xvectors dataset
        self._speaker_embedding = torch.randn(1, 512)

    def synthesize(self, text: str) -> tuple[np.ndarray, int]:
        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")
        inputs = self._processor(text=text, return_tensors="pt")
        with torch.no_grad():
            speech = self._model.generate_speech(
                inputs["input_ids"],
                self._speaker_embedding,
                vocoder=self._vocoder,
            )
        return speech.numpy(), 16000

    def unload(self) -> None:
        if self._model is not None:
            from roohai.base import _release_torch_memory
            del self._model
            del self._processor
            del self._vocoder
            del self._speaker_embedding
            self._model = None
            self._processor = None
            self._vocoder = None
            self._speaker_embedding = None
            _release_torch_memory()

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