from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
from collections.abc import AsyncGenerator, AsyncIterator
from datetime import datetime, timezone
from email.utils import format_datetime
from typing import Any
from urllib.parse import urlencode, urlparse

from pipecat.frames.frames import ErrorFrame, Frame
from pipecat.services.tts_service import TTSService
from websockets.asyncio.client import connect


DEFAULT_XFYUN_TTS_URL = "wss://tts-api.xfyun.cn/v2/tts"


class XfyunTTSService(TTSService):
    """iFlytek/Xfyun online TTS service for Pipecat.

    Xfyun's API is not OpenAI-compatible. It uses a signed WebSocket URL,
    receives one JSON request per synthesis, and streams text WebSocket
    messages containing base64-encoded audio chunks. This service requests
    raw PCM so the chunks can become Pipecat audio frames without MP3 decode.
    """

    def __init__(
        self,
        *,
        app_id: str,
        api_key: str,
        api_secret: str,
        voice: str,
        url: str | None = None,
        sample_rate: int = 16000,
        source_sample_rate: int = 16000,
        encoding: str = "raw",
        text_encoding: str = "UTF8",
        speed: int = 50,
        volume: int = 50,
        pitch: int = 50,
        timeout: float = 30.0,
        **kwargs,
    ) -> None:
        super().__init__(sample_rate=sample_rate, **kwargs)
        self._app_id = app_id or os.environ.get("XFYUN_APP_ID", "")
        self._api_key = api_key or os.environ.get("XFYUN_API_KEY", "")
        self._api_secret = api_secret or os.environ.get("XFYUN_API_SECRET", "")
        self._voice = voice
        self._url = url or DEFAULT_XFYUN_TTS_URL
        self._source_sample_rate = source_sample_rate
        self._encoding = encoding
        self._text_encoding = text_encoding
        self._speed = speed
        self._volume = volume
        self._pitch = pitch
        self._timeout = timeout

    async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame, None]:
        if not text:
            return

        if not self._app_id or not self._api_key or not self._api_secret:
            yield ErrorFrame(error="Xfyun TTS requires app_id, api_key, and api_secret")
            return

        if len(text.encode("utf-8")) >= 8000:
            yield ErrorFrame(error="Xfyun TTS text must be less than 8000 UTF-8 bytes")
            return

        if self._encoding != "raw":
            yield ErrorFrame(error="Xfyun TTS is configured for PCM output; set aue/encoding to raw")
            return

        try:
            await self.start_tts_usage_metrics(text)

            first_frame = True
            async for frame in self._stream_audio_frames_from_iterator(
                self._iter_audio_chunks(text),
                in_sample_rate=self._source_sample_rate,
                context_id=context_id,
            ):
                if first_frame:
                    await self.stop_ttfb_metrics()
                    first_frame = False
                yield frame

            if first_frame:
                yield ErrorFrame(error="Xfyun TTS request finished without audio data")
        except Exception as exc:
            yield ErrorFrame(error=f"Xfyun TTS request failed: {exc}")

    async def _iter_audio_chunks(self, text: str) -> AsyncIterator[bytes]:
        request = self._build_request_frame(text)
        auth_url = _build_auth_url(self._url, self._api_key, self._api_secret)

        async with connect(auth_url, max_size=None, open_timeout=self._timeout) as websocket:
            await websocket.send(json.dumps(request, ensure_ascii=False))

            async for message in websocket:
                payload = json.loads(message)
                code = payload.get("code", -1)
                if code != 0:
                    message = payload.get("message", "unknown error")
                    sid = payload.get("sid")
                    raise RuntimeError(f"code={code}, sid={sid}, message={message}")

                data = payload.get("data")
                if not isinstance(data, dict):
                    continue

                audio_b64 = data.get("audio")
                if audio_b64:
                    yield base64.b64decode(audio_b64)

                if data.get("status") == 2:
                    break

    def _build_request_frame(self, text: str) -> dict[str, Any]:
        business: dict[str, Any] = {
            "aue": self._encoding,
            "auf": f"audio/L16;rate={self._source_sample_rate}",
            "vcn": self._voice,
            "speed": self._speed,
            "volume": self._volume,
            "pitch": self._pitch,
            "tte": self._text_encoding,
        }

        return {
            "common": {"app_id": self._app_id},
            "business": business,
            "data": {
                "status": 2,
                "text": base64.b64encode(text.encode("utf-8")).decode("utf-8"),
            },
        }


def _build_auth_url(url: str, api_key: str, api_secret: str) -> str:
    parsed = urlparse(url)
    host = parsed.netloc
    path = parsed.path or "/v2/tts"
    date = format_datetime(datetime.now(timezone.utc), usegmt=True)
    request_line = f"GET {path} HTTP/1.1"
    signature_origin = f"host: {host}\ndate: {date}\n{request_line}"
    signature_sha = hmac.new(
        api_secret.encode("utf-8"),
        signature_origin.encode("utf-8"),
        digestmod=hashlib.sha256,
    ).digest()
    signature = base64.b64encode(signature_sha).decode("utf-8")
    authorization_origin = (
        f'api_key="{api_key}", algorithm="hmac-sha256", '
        f'headers="host date request-line", signature="{signature}"'
    )
    authorization = base64.b64encode(authorization_origin.encode("utf-8")).decode("utf-8")
    query = urlencode({"authorization": authorization, "date": date, "host": host})
    return f"{url}?{query}"