from __future__ import annotations

import asyncio
import base64
import json
import logging
import os
import time
import weakref
from dataclasses import dataclass, replace
from typing import Any

import aiohttp

from livekit.agents import (
    APIConnectionError,
    APIConnectOptions,
    APIStatusError,
    APITimeoutError,
    tokenize,
    tts,
    utils,
)
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, NotGivenOr
from livekit.agents.utils import is_given

from .log import logger
from .models import (
    TTSDefaultVoiceId,
    TTSDefaultVoiceStyle,
    TTSEncoding,
    TTSLocales,
    TTSModels,
    TTSStyles,
)

API_AUTH_HEADER = "api-key"


@dataclass
class _TTSOptions:
    api_key: str
    locale: TTSLocales | str | None
    model: TTSModels | str
    voice: str
    style: str | None
    speed: int | None
    pitch: int | None
    sample_rate: int
    encoding: TTSEncoding
    base_url: str
    min_buffer_size: int
    max_buffer_delay_in_ms: int
    verbose: bool

    def get_http_url(self, path: str) -> str:
        return f"{self.base_url}{path}"

    def get_ws_url(self, path: str) -> str:
        return f"{self.base_url.replace('http', 'ws', 1)}{path}"


class TTS(tts.TTS):
    def __init__(
        self,
        *,
        api_key: str | None = None,
        model: TTSModels | str = "FALCON",
        locale: TTSLocales | str | None = None,
        voice: str = TTSDefaultVoiceId,
        style: TTSStyles | str | None = None,
        speed: int | None = None,
        pitch: int | None = None,
        sample_rate: int = 24000,
        encoding: TTSEncoding = "pcm",
        base_url: str = "https://global.api.murf.ai",
        http_session: aiohttp.ClientSession | None = None,
        tokenizer: NotGivenOr[tokenize.SentenceTokenizer] = NOT_GIVEN,
        text_pacing: tts.SentenceStreamPacer | bool = False,
        min_buffer_size: int = 3,
        max_buffer_delay_in_ms: int = 0,
        streaming: bool = True,
        verbose: bool = False,
    ) -> None:
        """
        Create a new instance of Murf AI TTS.

        See https://murf.ai/api/docs/api-reference/text-to-speech/stream-input for more details on the the Murf AI API.

        Args:
            api_key (str | None, optional): The Murf AI API key. If not provided, it will be read from the MURF_API_KEY environment variable.
            model (TTSModels | str, optional): The Murf AI TTS model to use. Defaults to "FALCON".
            locale (str | None, optional): The locale for synthesis (e.g., "en-US", "en-UK"). If not provided, will be inferred from voice.
            voice (str, optional): The voice ID from Murf AI's voice library (e.g., "en-US-matthew"). Defaults to TTSDefaultVoiceId.
            style (TTSStyles | str | None, optional): The voice style to apply (e.g., "Conversation"). Can be None for default style.
            speed (int | None, optional): The speech speed control. Higher values = faster speech. None for default speed.
            pitch (int | None, optional): The speech pitch control. Higher values = higher pitch. None for default pitch.
            sample_rate (int, optional): The audio sample rate in Hz. Defaults to 24000.
            encoding (str, optional): The audio encoding format. Defaults to "pcm".
            http_session (aiohttp.ClientSession | None, optional): An existing aiohttp ClientSession to use. If not provided, a new session will be created.
            base_url (str, optional): The base URL for the Murf AI API. Defaults to "https://global.api.murf.ai".
            tokenizer (tokenize.SentenceTokenizer, optional): The tokenizer to use. Defaults to tokenize.basic.SentenceTokenizer(min_sentence_len=min_buffer_size).
            text_pacing (tts.SentenceStreamPacer | bool, optional): Stream pacer for the TTS. Set to True to use the default pacer, False to disable.
            min_buffer_size (int, optional):Minimum characters to buffer before sending text to audio when no sentence boundary is detected. Higher values improve quality; lower values reduce TTFB. Defaults to 3.
            max_buffer_delay_in_ms (int, optional): Maximum wait time before sending buffered text if min_buffer_size isn’t reached. Defaults to 0.
            verbose (bool, optional): Enable detailed Murf logging. When True, logs TTFB, latency metrics, buffer configuration, and other diagnostic information. Also enabled when logger level is DEBUG. Defaults to False.
            streaming (bool, optional): If True, uses WebSocket streaming for real-time audio. If False, uses HTTP requests. Defaults to True.
        """  # noqa: E501

        self._streaming = streaming

        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=streaming),
            sample_rate=sample_rate,
            num_channels=1,
        )

        murf_api_key = api_key or os.environ.get("MURF_API_KEY")
        if not murf_api_key:
            raise ValueError("MURF_API_KEY must be set")

        self._opts = _TTSOptions(
            api_key=murf_api_key,
            model=model,
            locale=locale,
            voice=voice,
            style=style or TTSDefaultVoiceStyle,
            speed=speed,
            pitch=pitch,
            sample_rate=sample_rate,
            encoding=encoding,
            base_url=base_url,
            min_buffer_size=min_buffer_size,
            max_buffer_delay_in_ms=max_buffer_delay_in_ms,
            verbose=verbose,
        )
        self._session = http_session
        self._pool = utils.ConnectionPool[aiohttp.ClientWebSocketResponse](
            connect_cb=self._connect_ws,
            close_cb=self._close_ws,
            max_session_duration=300,
            mark_refreshed_on_get=True,
        )
        self._streams = weakref.WeakSet[SynthesizeStream]()
        self._sentence_tokenizer = (
            tokenizer if is_given(tokenizer) else tokenize.blingfire.SentenceTokenizer(min_sentence_len=min_buffer_size)
        )
        self._stream_pacer: tts.SentenceStreamPacer | None = None
        if text_pacing is True:
            self._stream_pacer = tts.SentenceStreamPacer()
        elif isinstance(text_pacing, tts.SentenceStreamPacer):
            self._stream_pacer = text_pacing

    @property
    def model(self) -> str:
        return self._opts.model

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

    def _is_verbose(self) -> bool:
        return self._opts.verbose or logger.isEnabledFor(logging.DEBUG)

    async def _connect_ws(self, timeout: float) -> aiohttp.ClientWebSocketResponse:
        session = self._ensure_session()
        url = self._opts.get_ws_url(
            f"/v1/speech/stream-input?api-key={self._opts.api_key}&sample_rate={self._opts.sample_rate}&format={self._opts.encoding}&model={self._opts.model}"
        )
        return await asyncio.wait_for(session.ws_connect(url), timeout)

    async def _close_ws(self, ws: aiohttp.ClientWebSocketResponse) -> None:
        await ws.close()

    def _ensure_session(self) -> aiohttp.ClientSession:
        if not self._session:
            self._session = utils.http_context.http_session()

        return self._session

    def prewarm(self) -> None:
        self._pool.prewarm()

    def update_options(
        self,
        *,
        locale: NotGivenOr[str] = NOT_GIVEN,
        voice: NotGivenOr[str] = NOT_GIVEN,
        style: NotGivenOr[str | None] = NOT_GIVEN,
        speed: NotGivenOr[int | None] = NOT_GIVEN,
        pitch: NotGivenOr[int | None] = NOT_GIVEN,
    ) -> None:
        """
        Update the Text-to-Speech (TTS) configuration options.

        This method allows updating the TTS settings, including model, locale, voice, style,
        speed and pitch. If any parameter is not provided, the existing value will be retained.

        Args:
            locale (str, optional): The locale for synthesis (e.g., "en-US", "en-UK").
            voice (str, optional): The voice ID from Murf AI's voice library. (e.g. "en-US-matthew")
            style (str | None, optional): The voice style to apply (e.g., "Conversation", "Promo").
            speed (int | None, optional): Controls the speech speed. Positive values increase speed, negative values decrease it. Valid range: -50 to 50.
            pitch (int | None, optional): Controls the speech pitch. Positive values raise pitch, negative values lower it. Valid range: -50 to 50.
        """
        if is_given(locale):
            self._opts.locale = locale
        if is_given(voice):
            self._opts.voice = voice
        if is_given(style):
            self._opts.style = style
        if is_given(speed):
            self._opts.speed = speed
        if is_given(pitch):
            self._opts.pitch = pitch

    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 stream(
        self, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
    ) -> SynthesizeStream:
        stream = SynthesizeStream(tts=self, conn_options=conn_options)
        self._streams.add(stream)
        return stream

    async def aclose(self) -> None:
        for stream in list(self._streams):
            await stream.aclose()

        self._streams.clear()
        await self._pool.aclose()


class ChunkedStream(tts.ChunkedStream):
    """Synthesize chunked text using the http streaming output endpoint"""

    def __init__(self, *, tts: TTS, input_text: str, conn_options: APIConnectOptions) -> None:
        super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
        self._tts: TTS = tts
        self._opts = replace(tts._opts)

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        request_id = utils.shortuuid()
        request_start_time = time.perf_counter()

        if self._tts._is_verbose():
            logger.info(
                "[Murf TTS] HTTP request started - request_id=%s, voice=%s, style=%s, locale=%s, "
                "model=%s, endpoint=%s",
                request_id,
                self._opts.voice,
                self._opts.style,
                self._opts.locale,
                self._opts.model,
                self._opts.base_url,
            )

        try:
            async with self._tts._ensure_session().post(
                self._opts.get_http_url("/v1/speech/stream"),
                headers={API_AUTH_HEADER: self._opts.api_key},
                json={
                    "text": self._input_text,
                    "model": self._opts.model,
                    "multiNativeLocale": self._opts.locale,
                    "voice_id": self._opts.voice,
                    "style": self._opts.style,
                    "rate": self._opts.speed,
                    "pitch": self._opts.pitch,
                    "format": self._opts.encoding,
                    "sample_rate": self._opts.sample_rate,
                },
                timeout=aiohttp.ClientTimeout(total=30, sock_connect=self._conn_options.timeout),
            ) as resp:
                resp.raise_for_status()

                output_emitter.initialize(
                    request_id=request_id,
                    sample_rate=self._opts.sample_rate,
                    num_channels=1,
                    mime_type="audio/pcm",
                )

                first_chunk = True
                async for data, _ in resp.content.iter_chunks():
                    if first_chunk:
                        first_chunk = False
                        if self._tts._is_verbose():
                            ttfb_ms = (time.perf_counter() - request_start_time) * 1000.0
                            logger.info("[Murf TTS] HTTP TTFB: %.2f ms, request_id=%s", ttfb_ms, request_id)
                    output_emitter.push(data)

                output_emitter.flush()
        except asyncio.TimeoutError:
            logger.error("[Murf TTS] HTTP request timed out - request_id=%s", request_id)
            raise APITimeoutError() from None
        except aiohttp.ClientResponseError as e:
            logger.error("[Murf TTS] HTTP error %d: %s - request_id=%s", e.status, e.message, request_id)
            raise APIStatusError(
                message=e.message, status_code=e.status, request_id=None, body=None
            ) from None
        except Exception as e:
            logger.error("[Murf TTS] HTTP request failed: %s - request_id=%s", str(e), request_id)
            raise APIConnectionError() from e


class SynthesizeStream(tts.SynthesizeStream):
    def __init__(self, *, tts: TTS, conn_options: APIConnectOptions):
        super().__init__(tts=tts, conn_options=conn_options)
        self._tts: TTS = tts
        self._opts = replace(tts._opts)

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

        input_sent_event = asyncio.Event()
        first_chunk_sent_time: float | None = None

        sent_tokenizer_stream = self._tts._sentence_tokenizer.stream()
        if self._tts._stream_pacer:
            sent_tokenizer_stream = self._tts._stream_pacer.wrap(
                sent_stream=sent_tokenizer_stream,
                audio_emitter=output_emitter,
            )

        async def _sentence_stream_task(ws: aiohttp.ClientWebSocketResponse) -> None:
            nonlocal first_chunk_sent_time
            context_id = utils.shortuuid()
            base_pkt = _to_murf_websocket_pkt(self._opts)
            first_sent = True
            async for ev in sent_tokenizer_stream:
                token_pkt = base_pkt.copy()
                token_pkt["context_id"] = context_id
                token_pkt["text"] = ev.token + " "
                self._mark_started()
                await ws.send_str(json.dumps(token_pkt))
                if first_sent:
                    first_sent = False
                    first_chunk_sent_time = time.perf_counter()
                    if self._tts._is_verbose():
                        logger.info(
                            "[Murf TTS] Stream started - context_id=%s, voice=%s, style=%s, locale=%s, "
                            "min_buffer_size=%d, max_buffer_delay_in_ms=%d, endpoint=%s",
                            context_id,
                            self._opts.voice,
                            self._opts.style,
                            self._opts.locale,
                            self._opts.min_buffer_size,
                            self._opts.max_buffer_delay_in_ms,
                            self._opts.base_url,
                        )
                input_sent_event.set()

            end_pkt = base_pkt.copy()
            end_pkt["context_id"] = context_id
            end_pkt["end"] = True
            await ws.send_str(json.dumps(end_pkt))
            input_sent_event.set()

        async def _input_task() -> None:
            async for data in self._input_ch:
                if isinstance(data, self._FlushSentinel):
                    sent_tokenizer_stream.flush()
                    continue

                sent_tokenizer_stream.push_text(data)

            sent_tokenizer_stream.end_input()

        async def _recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
            nonlocal first_chunk_sent_time
            current_segment_id: str | None = None
            first_audio_received = True
            await input_sent_event.wait()
            while True:
                msg = await ws.receive()
                if msg.type in (
                    aiohttp.WSMsgType.CLOSED,
                    aiohttp.WSMsgType.CLOSE,
                    aiohttp.WSMsgType.CLOSING,
                ):
                    logger.error("[Murf TTS] Connection closed unexpectedly")
                    raise APIStatusError(
                        "Murf AI connection closed unexpectedly", request_id=request_id
                    )

                if msg.type != aiohttp.WSMsgType.TEXT:
                    logger.warning("[Murf TTS] Unexpected message type %s", msg.type)
                    continue

                data = json.loads(msg.data)
                segment_id = data.get("context_id")
                if current_segment_id is None:
                    current_segment_id = segment_id
                    output_emitter.start_segment(segment_id=current_segment_id)
                if data.get("audio"):
                    if first_audio_received:
                        first_audio_received = False
                        if first_chunk_sent_time is not None:
                            ttfb_ms = (time.perf_counter() - first_chunk_sent_time) * 1000.0
                            if self._tts._is_verbose():
                                logger.info("[Murf TTS] Murf TTFB (first sentence to first audio): %.2f ms", ttfb_ms)
                    b64data = base64.b64decode(data["audio"])
                    output_emitter.push(b64data)
                elif data.get("final"):
                    if sent_tokenizer_stream.closed:
                        # close only if the input stream is closed
                        output_emitter.end_input()
                        break
                else:
                    logger.warning("[Murf TTS] Unexpected message %s", data)

        try:
            async with self._tts._pool.connection(timeout=self._conn_options.timeout) as ws:
                tasks = [
                    asyncio.create_task(_input_task()),
                    asyncio.create_task(_sentence_stream_task(ws)),
                    asyncio.create_task(_recv_task(ws)),
                ]

                try:
                    await asyncio.gather(*tasks)
                finally:
                    input_sent_event.set()
                    await sent_tokenizer_stream.aclose()
                    await utils.aio.gracefully_cancel(*tasks)
        except asyncio.TimeoutError:
            raise APITimeoutError() from None
        except aiohttp.ClientResponseError as e:
            logger.error("[Murf TTS] WebSocket error %d: %s", e.status, e.message)
            raise APIStatusError(
                message=e.message, status_code=e.status, request_id=None, body=None
            ) from None
        except Exception as e:
            logger.error("[Murf TTS] WebSocket connection failed: %s", str(e))
            raise APIConnectionError() from e


def _to_murf_websocket_pkt(opts: _TTSOptions) -> dict[str, Any]:
    voice_config: dict[str, Any] = {}

    if opts.voice:
        voice_config["voice_id"] = opts.voice

    if opts.style:
        voice_config["style"] = opts.style

    if opts.speed:
        voice_config["rate"] = opts.speed

    if opts.pitch:
        voice_config["pitch"] = opts.pitch

    if opts.locale:
        voice_config["multi_native_locale"] = opts.locale

    return {
        "voice_config": voice_config,
        "min_buffer_size": opts.min_buffer_size,
        "max_buffer_delay_in_ms": opts.max_buffer_delay_in_ms,
    }
