"""Синтез речи моделями Silero (https://github.com/snakers4/silero-models)."""

from __future__ import annotations

import logging
import os
import threading
import urllib.request
import warnings
from pathlib import Path
from typing import Any

import numpy as np

from agr_assistent.config import TTSConfig
from agr_assistent.tts.base import TTSError

log = logging.getLogger(__name__)

_MODEL_URL = "https://models.silero.ai/models/tts/{language}/{model}.pt"
_MAX_TORCH_THREADS = 4
_WARM_UP_TEXT = "Готова к работе."


class SileroTTS:
    def __init__(self, config: TTSConfig, models_dir: Path) -> None:
        self._config = config
        self._model_path = models_dir / "silero" / f"{config.model}.pt"
        self._model: Any = None
        self._lock = threading.Lock()

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

    def load(self) -> None:
        with self._lock:
            if self._model is not None:
                return
            if not self._model_path.exists():
                self._download()

            # torch импортируется долго, поэтому только при первой загрузке модели
            import torch

            torch.set_num_threads(min(_MAX_TORCH_THREADS, os.cpu_count() or 1))
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    importer = torch.package.PackageImporter(str(self._model_path))
                    model = importer.load_pickle("tts_models", "model")
                model.to(torch.device(self._config.device))
            except Exception as exc:
                raise TTSError(f"Не удалось загрузить модель {self._model_path}: {exc}") from exc

            if self._config.speaker not in model.speakers:
                raise TTSError(
                    f"Голос '{self._config.speaker}' отсутствует в модели {self._config.model} "
                    f"(есть: {', '.join(model.speakers)})"
                )
            # Первый синтез заметно медленнее последующих — прогреваем заранее
            self._apply_tts(model, _WARM_UP_TEXT)
            self._model = model
            log.info("Модель Silero %s загружена", self._config.model)

    def synthesize(self, text: str) -> np.ndarray:
        self.load()
        return self._apply_tts(self._model, text)

    def _apply_tts(self, model: Any, text: str) -> np.ndarray:
        import torch

        with torch.inference_mode(), warnings.catch_warnings():
            warnings.simplefilter("ignore")
            audio = model.apply_tts(
                text=text,
                speaker=self._config.speaker,
                sample_rate=self._config.sample_rate,
            )
        return audio.cpu().numpy().astype(np.float32, copy=False)

    def _download(self) -> None:
        language = self._config.model.rsplit("_", 1)[-1]
        url = _MODEL_URL.format(language=language, model=self._config.model)
        partial_path = self._model_path.with_suffix(".part")
        log.info("Скачиваю модель Silero: %s", url)
        try:
            self._model_path.parent.mkdir(parents=True, exist_ok=True)
            urllib.request.urlretrieve(url, partial_path)
            partial_path.replace(self._model_path)
        except OSError as exc:
            partial_path.unlink(missing_ok=True)
            raise TTSError(f"Не удалось скачать модель Silero {url}: {exc}") from exc