"""Hugging Face remote-code entry point for the ONNX Runtime TeraTTS release."""

from __future__ import annotations

from pathlib import Path
from typing import Iterator

import numpy as np
from transformers import PreTrainedModel

from .configuration_teratts import TeraTTSConfig
# Keep the bundled stress-runtime module in Transformers' remote-code module
# dependency set. ``teratts.py`` imports it at runtime after this entry point
# has been loaded.
from .teratts_ruaccent import RUAccent as _BundledRUAccent
from .teratts import (
    LoadedTTS,
    generate_speech,
    generate_speech_stream,
    load_model,
    normalize_text,
    write_wav,
)


class TeraTTSModel(PreTrainedModel):
    """A reusable ONNX Runtime speech generator loaded through ``AutoModel``."""

    config_class = TeraTTSConfig

    def __init__(self, config: TeraTTSConfig) -> None:
        super().__init__(config)
        self.runtime: LoadedTTS | None = None

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        *model_args: object,
        config: TeraTTSConfig | None = None,
        provider: str = "CPUExecutionProvider",
        threads: int | None = None,
        diffusion_model: str | None = None,
        russian_stress: bool = True,
        ruaccent_model_size: str = "turbo3.1",
        ruaccent_device: str = "CPU",
        ruaccent_mode: str = "full",
        **kwargs: object,
    ) -> "TeraTTSModel":
        """Download/load the release and initialize reusable ONNX sessions.

        ``provider``, ``threads``, and ``diffusion_model`` are TeraTTS-specific
        arguments. Standard Hub arguments such as ``revision``, ``token``,
        ``cache_dir``, and ``local_files_only`` are forwarded to the snapshot
        download when a Hub model ID is supplied.
        """
        if model_args:
            raise TypeError("TeraTTSModel.from_pretrained accepts no positional model arguments")
        # ``AutoModel`` forwards these framework-only values to a custom
        # class after it has already used them. This ONNX runtime has no
        # adapters or PyTorch state dict to load, so safely discard them.
        for key in (
            "trust_remote_code",
            "adapter_kwargs",
            "_from_auto",
            "use_safetensors",
            "weights_only",
            "low_cpu_mem_usage",
            "device_map",
            "torch_dtype",
            "dtype",
        ):
            kwargs.pop(key, None)
        source = Path(pretrained_model_name_or_path)
        if source.is_dir():
            release = source
        else:
            from huggingface_hub import snapshot_download

            download_keys = {"revision", "token", "cache_dir", "local_files_only", "force_download"}
            download_kwargs = {key: kwargs.pop(key) for key in list(kwargs) if key in download_keys}
            if kwargs:
                unexpected = ", ".join(sorted(kwargs))
                raise TypeError(f"unexpected TeraTTS loading arguments: {unexpected}")
            release = Path(
                snapshot_download(repo_id=str(pretrained_model_name_or_path), **download_kwargs)
            )
        if config is None:
            config = TeraTTSConfig.from_pretrained(release)
        instance = cls(config)
        instance.runtime = load_model(
            release,
            model=diffusion_model or config.default_diffusion_model,
            provider=provider,
            threads=threads,
            russian_stress=russian_stress,
            ruaccent_model_size=ruaccent_model_size,
            ruaccent_device=ruaccent_device,
            ruaccent_mode=ruaccent_mode,
        )
        return instance

    def _runtime(self) -> LoadedTTS:
        if self.runtime is None:
            raise RuntimeError("load this model with from_pretrained() before generating speech")
        return self.runtime

    def generate_speech(
        self,
        text: str,
        voice: str,
        *,
        duration_scale: float = 1.0,
        guidance: float = 3.0,
        seed: int = 1234,
    ) -> np.ndarray:
        return generate_speech(
            self._runtime(),
            text,
            voice,
            duration_scale=duration_scale,
            guidance=guidance,
            seed=seed,
        )

    def normalize_text(self, text: str) -> str:
        """Show the final tagged text passed to the TTS text encoder."""
        return normalize_text(self._runtime(), text)

    def generate_speech_stream(
        self,
        text: str,
        voice: str,
        *,
        duration_scale: float = 1.0,
        guidance: float = 3.0,
        seed: int = 1234,
        chunk_frames: int = 16,
    ) -> Iterator[np.ndarray]:
        yield from generate_speech_stream(
            self._runtime(),
            text,
            voice,
            duration_scale=duration_scale,
            guidance=guidance,
            seed=seed,
            chunk_frames=chunk_frames,
        )

    def save_wav(self, path: str | Path, waveform: np.ndarray) -> None:
        """Write a generated mono float32 waveform as a 44.1 kHz WAV file."""
        write_wav(Path(path), waveform)