from __future__ import annotations

import asyncio
from typing import List
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from xml.sax.saxutils import escape, quoteattr

import httpx
from livekit.agents import (
    APIConnectionError,
    APIStatusError,
    APITimeoutError,
    tts,
    utils,
)
from livekit.agents.types import APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS


AZURE_OUTPUT_FORMATS = {
    8000: "raw-8khz-16bit-mono-pcm",
    16000: "raw-16khz-16bit-mono-pcm",
    22050: "raw-22050hz-16bit-mono-pcm",
    24000: "raw-24khz-16bit-mono-pcm",
    44100: "raw-44100hz-16bit-mono-pcm",
    48000: "raw-48khz-16bit-mono-pcm",
}


def _append_query_param(url: str, key: str, value: str) -> str:
    if not value:
        return url

    parsed = urlsplit(url)
    query_pairs = parse_qsl(parsed.query, keep_blank_values=True)
    if any(existing_key == key for existing_key, _ in query_pairs):
        return url

    query_pairs.append((key, value))
    return urlunsplit(parsed._replace(query=urlencode(query_pairs)))


def _is_custom_domain_endpoint(url: str) -> bool:
    parsed = urlsplit((url or "").strip())
    host = (parsed.hostname or "").strip().lower()
    return host.endswith(".cognitiveservices.azure.com")


def _prepend_service_prefix(url: str, service: str) -> str:
    parsed = urlsplit((url or "").strip())
    path = (parsed.path or "").lstrip("/")
    service = (service or "").strip().strip("/")
    if not service or path.startswith(f"{service}/"):
        return url.rstrip("/")
    return urlunsplit(parsed._replace(path=f"/{service}/{path}")).rstrip("/")


class AzureRESTTTS(tts.TTS):
    def __init__(
        self,
        *,
        voice: str,
        language: str | None = None,
        sample_rate: int = 16000,
        speech_key: str | None = None,
        speech_region: str | None = None,
        speech_endpoint: str | None = None,
        deployment_id: str | None = None,
        speech_auth_token: str | None = None,
        user_agent: str = "tia-azure-tts/1.0",
        timeout_s: float = 30.0,
    ) -> None:
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=False, aligned_transcript=False),
            sample_rate=sample_rate,
            num_channels=1,
        )
        if sample_rate not in AZURE_OUTPUT_FORMATS:
            raise ValueError(
                f"Unsupported sample rate {sample_rate}. Supported: {sorted(AZURE_OUTPUT_FORMATS)}"
            )
        if not (voice or "").strip():
            raise ValueError("voice is required")
        if not ((speech_key or "").strip() or (speech_auth_token or "").strip()):
            raise ValueError("speech_key or speech_auth_token is required")
        if not ((speech_region or "").strip() or (speech_endpoint or "").strip()):
            raise ValueError("speech_region or speech_endpoint is required")

        self._voice = voice.strip()
        self._language = (language or "").strip() or None
        self._speech_key = (speech_key or "").strip() or None
        self._speech_region = (speech_region or "").strip() or None
        self._speech_endpoint = (speech_endpoint or "").strip().rstrip("/") or None
        self._deployment_id = (deployment_id or "").strip() or None
        self._speech_auth_token = (speech_auth_token or "").strip() or None
        self._user_agent = (user_agent or "tia-azure-tts/1.0").strip()
        self._timeout_s = float(timeout_s)

    @property
    def model(self) -> str:
        return self._deployment_id or self._voice

    @property
    def provider(self) -> str:
        return "azure"

    def synthesize(
        self,
        text: str,
        *,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> "ChunkedStream":
        return ChunkedStream(tts=self, input_text=text, conn_options=conn_options)

    def _base_endpoint(self) -> str:
        if self._speech_endpoint:
            return self._speech_endpoint
        assert self._speech_region
        service = "voice" if self._deployment_id else "tts"
        return f"https://{self._speech_region}.{service}.speech.microsoft.com/cognitiveservices/v1"

    def _request_endpoints(self) -> List[str]:
        base = _append_query_param(self._base_endpoint(), "deploymentId", self._deployment_id or "")
        endpoints = [base]

        parsed = urlsplit(base)
        path = (parsed.path or "").rstrip("/")
        if _is_custom_domain_endpoint(base) and path == "/cognitiveservices/v1":
            service = "voice" if self._deployment_id else "tts"
            endpoints.append(
                _append_query_param(
                    _prepend_service_prefix(base, service),
                    "deploymentId",
                    self._deployment_id or "",
                )
            )
        return endpoints

    def _build_ssml(self, text: str) -> bytes:
        language = self._language or "pt-BR"
        escaped_text = escape((text or "").strip())
        return (
            f"<speak version='1.0' "
            f"xmlns='http://www.w3.org/2001/10/synthesis' "
            f"xmlns:mstts='http://www.w3.org/2001/mstts' "
            f"xml:lang={quoteattr(language)}>"
            f"<voice name={quoteattr(self._voice)}>{escaped_text}</voice>"
            f"</speak>"
        ).encode("utf-8")

    def _headers(self) -> dict[str, str]:
        headers = {
            "Content-Type": "application/ssml+xml",
            "X-Microsoft-OutputFormat": AZURE_OUTPUT_FORMATS[self.sample_rate],
            "User-Agent": self._user_agent,
        }
        if self._speech_auth_token:
            headers["Authorization"] = f"Bearer {self._speech_auth_token}"
        elif self._speech_key:
            headers["Ocp-Apim-Subscription-Key"] = self._speech_key
        return headers

    def synthesize_pcm(self, text: str) -> bytes:
        text = (text or "").strip()
        if not text:
            return b""

        headers = self._headers()
        body = self._build_ssml(text)
        last_status_error: httpx.HTTPStatusError | None = None
        endpoints = self._request_endpoints()

        with httpx.Client(timeout=self._timeout_s, follow_redirects=True) as client:
            for index, endpoint in enumerate(endpoints):
                try:
                    response = client.post(endpoint, headers=headers, content=body)
                    response.raise_for_status()
                    if not response.content:
                        raise RuntimeError("Azure TTS returned empty audio.")
                    self._speech_endpoint = endpoint
                    return response.content
                except httpx.HTTPStatusError as exc:
                    should_try_next = (
                        exc.response is not None
                        and exc.response.status_code == 404
                        and index < len(endpoints) - 1
                    )
                    if should_try_next:
                        last_status_error = exc
                        continue
                    raise

        if last_status_error is not None:
            raise last_status_error
        raise RuntimeError("Azure TTS returned no response.")

    async def aclose(self) -> None:
        return None


class ChunkedStream(tts.ChunkedStream):
    def __init__(
        self,
        *,
        tts: AzureRESTTTS,
        input_text: str,
        conn_options: APIConnectOptions,
    ) -> None:
        super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
        self._tts: AzureRESTTTS = tts

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        output_emitter.initialize(
            request_id=utils.shortuuid(),
            sample_rate=self._tts.sample_rate,
            num_channels=self._tts.num_channels,
            mime_type="audio/pcm",
        )

        try:
            audio = await asyncio.to_thread(self._tts.synthesize_pcm, self._input_text)
        except httpx.TimeoutException as exc:
            raise APITimeoutError() from exc
        except httpx.HTTPStatusError as exc:
            status_code = exc.response.status_code if exc.response is not None else -1
            request_id = exc.response.headers.get("X-RequestId") if exc.response is not None else None
            body = exc.response.text if exc.response is not None else None
            message = "Azure TTS request failed."
            if body:
                message = f"{message} {body}"
            raise APIStatusError(
                message=message,
                status_code=status_code,
                request_id=request_id,
                body=body,
            ) from exc
        except httpx.RequestError as exc:
            raise APIConnectionError(f"Could not connect to Azure TTS: {exc}") from exc
        except RuntimeError as exc:
            raise APIConnectionError(str(exc), retryable=False) from exc

        output_emitter.push(audio)
        output_emitter.flush()
