"""Rime TTS provider — HTTP chunked PCM endpoint, pure aiohttp.

Both Arcana and Mist model families are supported. Requests Rime's
``audio/pcm`` accept header so chunks are raw PCM_S16LE suitable for direct
use in Patter's pipeline.
"""

from __future__ import annotations

import logging
import os
from enum import StrEnum
from typing import ClassVar, Any, AsyncIterator, Optional, Union

from getpatter.providers.base import TTSProvider

logger = logging.getLogger("getpatter.providers.rime_tts")

try:  # pragma: no cover - trivial import guard
    import aiohttp
except ImportError:  # pragma: no cover
    aiohttp = None  # type: ignore

RIME_BASE_URL = "https://users.rime.ai/v1/rime-tts"


class RimeModel(StrEnum):
    """Rime TTS model families. Arcana is higher-quality / latency-tolerant,
    Mist variants prioritise low latency."""

    ARCANA = "arcana"
    MIST = "mist"
    MIST_V2 = "mistv2"


class RimeAudioFormat(StrEnum):
    """Supported response Content-Type accept headers for Rime TTS."""

    PCM = "audio/pcm"
    MP3 = "audio/mp3"
    WAV = "audio/wav"
    MULAW = "audio/mulaw"


# Model-specific timeouts — Arcana can take up to ~80% of the audio duration
# it is synthesizing.
ARCANA_MODEL_TIMEOUT = 60 * 4
MIST_MODEL_TIMEOUT = 30


def _is_mist_model(model: str) -> bool:
    # Rime Mist-family model ids are ``mist``, ``mistv2``, etc. — always
    # prefixed with ``mist``. Use ``startswith`` so unrelated model ids that
    # happen to contain the substring don't accidentally match.
    return model.startswith(RimeModel.MIST.value)


def _timeout_for_model(model: str) -> int:
    if model == RimeModel.ARCANA.value:
        return ARCANA_MODEL_TIMEOUT
    return MIST_MODEL_TIMEOUT


class RimeTTS(TTSProvider):
    """Rime TTS over the HTTP chunked endpoint.

    Defaults to the ``arcana`` model with the ``astra`` voice. Output is
    PCM_S16LE at the configured ``sample_rate`` (default 16000 Hz).
    """

    #: Stable pricing/dashboard key — read by stream-handler/metrics.
    provider_key: ClassVar[str] = "rime"

    def __init__(
        self,
        api_key: Optional[str] = None,
        *,
        model: Union[RimeModel, str] = RimeModel.ARCANA,
        speaker: Optional[str] = None,
        lang: str = "eng",
        sample_rate: int = 16000,
        # Arcana-only options
        repetition_penalty: Optional[float] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        max_tokens: Optional[int] = None,
        # Mist-only options
        speed_alpha: Optional[float] = None,
        reduce_latency: Optional[bool] = None,
        pause_between_brackets: Optional[bool] = None,
        phonemize_between_brackets: Optional[bool] = None,
        base_url: str = RIME_BASE_URL,
        session: Optional["aiohttp.ClientSession"] = None,
    ) -> None:
        if aiohttp is None:
            raise ImportError(
                "aiohttp is required for RimeTTS. "
                "Install with: pip install getpatter[rime]"
            )

        resolved_key = api_key or os.environ.get("RIME_API_KEY")
        if not resolved_key:
            raise ValueError(
                "Rime API key is required, either as argument or set "
                "RIME_API_KEY environment variable"
            )

        if speaker is None:
            # Sensible defaults: "cove" (DefaultMistVoice) for Mist models,
            # "astra" for Arcana. Callers can override either.
            speaker = "cove" if _is_mist_model(model) else "astra"

        self.api_key = resolved_key
        self.model = model
        self.speaker = speaker
        self.lang = lang
        self.sample_rate = sample_rate
        self.repetition_penalty = repetition_penalty
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.speed_alpha = speed_alpha
        self.reduce_latency = reduce_latency
        self.pause_between_brackets = pause_between_brackets
        self.phonemize_between_brackets = phonemize_between_brackets
        self.base_url = base_url
        self._total_timeout = _timeout_for_model(model)
        self._owns_session = session is None
        self._session = session

    def __repr__(self) -> str:
        return (
            f"RimeTTS(model={self.model!r}, speaker={self.speaker!r}, "
            f"lang={self.lang!r}, sample_rate={self.sample_rate})"
        )

    def _ensure_session(self) -> "aiohttp.ClientSession":
        if self._session is None:
            self._session = aiohttp.ClientSession()
            self._owns_session = True
        return self._session

    def _build_payload(self, text: str) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "speaker": self.speaker,
            "text": text,
            "modelId": self.model,
        }

        if self.model == RimeModel.ARCANA.value:
            if self.repetition_penalty is not None:
                payload["repetition_penalty"] = self.repetition_penalty
            if self.temperature is not None:
                payload["temperature"] = self.temperature
            if self.top_p is not None:
                payload["top_p"] = self.top_p
            if self.max_tokens is not None:
                payload["max_tokens"] = self.max_tokens
            payload["lang"] = self.lang
            payload["samplingRate"] = self.sample_rate
        elif _is_mist_model(self.model):
            payload["lang"] = self.lang
            payload["samplingRate"] = self.sample_rate
            if self.speed_alpha is not None:
                payload["speedAlpha"] = self.speed_alpha
            if (
                self.model == RimeModel.MIST_V2.value
                and self.reduce_latency is not None
            ):
                payload["reduceLatency"] = self.reduce_latency
            if self.pause_between_brackets is not None:
                payload["pauseBetweenBrackets"] = self.pause_between_brackets
            if self.phonemize_between_brackets is not None:
                payload["phonemizeBetweenBrackets"] = self.phonemize_between_brackets

        return payload

    def _record_synthesis_cost(self, text: str) -> None:
        """Emit ``patter.cost.tts_chars`` for the synthesised text."""
        try:
            from getpatter.observability.attributes import record_patter_attrs

            record_patter_attrs(
                {
                    "patter.cost.tts_chars": len(text),
                    "patter.tts.provider": "rime",
                }
            )
        except Exception:  # pragma: no cover — defense in depth
            logger.debug("_record_synthesis_cost failed", exc_info=True)

    async def synthesize(self, text: str) -> AsyncIterator[bytes]:
        """Stream raw PCM_S16LE bytes for ``text`` over HTTP."""
        self._record_synthesis_cost(text)
        session = self._ensure_session()

        headers = {
            "accept": RimeAudioFormat.PCM.value,
            "Authorization": f"Bearer {self.api_key}",
            "content-type": "application/json",
        }

        async with session.post(
            self.base_url,
            headers=headers,
            json=self._build_payload(text),
            timeout=aiohttp.ClientTimeout(total=self._total_timeout),
        ) as resp:
            resp.raise_for_status()
            # Rime returns audio/pcm on success; any other content type is a
            # surface-level error (e.g. JSON error payload).
            content_type = resp.headers.get("Content-Type", "")
            if not content_type.startswith("audio"):
                body = await resp.text()
                raise RuntimeError(f"Rime returned non-audio response: {body[:500]}")

            async for chunk in resp.content.iter_chunked(4096):
                if chunk:
                    yield chunk

    async def close(self) -> None:
        """Close the underlying session (idempotent)."""
        if self._session is not None and self._owns_session:
            await self._session.close()
            self._session = None
