"""Bark TTS implementation using HuggingFace transformers."""

from __future__ import annotations

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


class BarkTTS(TTSModel):
    META = {
        "display_name": "Bark",
        "description": "Suno's generative text-to-audio model. Expressive speech, runs locally.",
        "kind": "local",
        "pip_dependencies": ["transformers>=4.46.0,<5.0.0", "torch"],
        "hf_models": ["suno/bark-small"],
        "estimated_download_mb": 1800,
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "suno/bark-small",
            },
        ],
    }

    def __init__(self, model_id: str = "suno/bark-small") -> None:
        self.model_id = model_id
        self._processor = None
        self._model = None

    def load(self) -> None:
        from transformers import AutoProcessor, BarkModel

        self._processor = AutoProcessor.from_pretrained(self.model_id)
        self._model = BarkModel.from_pretrained(self.model_id)

    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, return_tensors="pt")
        with torch.no_grad():
            audio_array = self._model.generate(**inputs)
        audio = audio_array.cpu().numpy().squeeze()
        return audio, 24000

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

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