# Copyright 2023 LiveKit, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import base64
import json
import os
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

SUPPORTED_SAMPLE_RATE = 48000


@dataclass
class _TTSOptions:
    voice: str | None
    voice_id: str | None
    pronunciation_id: str | None
    word_tokenizer: tokenize.WordTokenizer
    json_config: dict[str, Any] | None = None


class TTS(tts.TTS):
    def __init__(
        self,
        *,
        api_key: str | None = None,
        model_endpoint: str | None = "wss://api.gradium.ai/api/speech/tts",
        model_name: str = "default",
        voice: str | None = None,
        voice_id: str | None = "YTpq7expH9539ERJ",
        pronunciation_id: str | None = None,
        json_config: dict[str, Any] | None = None,
        http_session: aiohttp.ClientSession | None = None,
        word_tokenizer: tokenize.WordTokenizer | None = None,
    ) -> None:
        """
        Initialize the Gradium TTS.

        Args:
            api_key (str): Gradium API key, or `GRADIUM_API_KEY` env var.
            model_endpoint (str): Gradium model endpoint, or `GRADIUM_MODEL_ENDPOINT` env var.
            model_name (str): Model name.
            voice (str): Speaker voice.
            voice_id (str): Speaker voice ID.
            pronunciation_id (str): Optional pronunciation ID for controlling TTS pronunciation.
            word_tokenizer (tokenize.WordTokenizer): Tokenizer for processing text. Defaults to basic WordTokenizer.
        """
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=True),
            sample_rate=SUPPORTED_SAMPLE_RATE,
            num_channels=1,
        )

        api_key = api_key or os.environ.get("GRADIUM_API_KEY")

        if not api_key:
            raise ValueError(
                "Gradium API key is required. "
                "Pass one in via the `api_key` parameter, "
                "or set it as the `GRADIUM_API_KEY` environment variable"
            )

        model_endpoint = model_endpoint or os.environ.get("GRADIUM_MODEL_ENDPOINT")

        if not model_endpoint:
            raise ValueError(
                "The model endpoint is required, you can find it in the Gradium dashboard"
            )

        self._api_key = api_key
        self._model_endpoint = model_endpoint
        self._model_name = model_name

        if not word_tokenizer:
            word_tokenizer = tokenize.basic.WordTokenizer(ignore_punctuation=False)
        self._opts = _TTSOptions(
            voice=voice,
            voice_id=voice_id,
            pronunciation_id=pronunciation_id,
            word_tokenizer=word_tokenizer,
            json_config=json_config,
        )
        self._session = http_session

    @property
    def model(self) -> str:
        return "unknown"

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

    async def _connect_ws(self, timeout: float) -> aiohttp.ClientWebSocketResponse:
        return await asyncio.wait_for(
            self._ensure_session().ws_connect(
                self._model_endpoint,
                headers={"x-api-key": self._api_key, "x-api-source": "livekit"},
            ),
            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 update_options(
        self,
        *,
        voice: NotGivenOr[str] = NOT_GIVEN,
        json_config: NotGivenOr[dict[str, Any]] = NOT_GIVEN,
    ) -> None:
        if is_given(voice):
            self._opts.voice = voice
        if is_given(json_config):
            self._opts.json_config = json_config

    def stream(
        self, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
    ) -> SynthesizeStream:
        return SynthesizeStream(tts=self, conn_options=conn_options)

    def synthesize(
        self, text: str, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
    ) -> ChunkedStream:
        return ChunkedStream(
            tts=self,
            api_key=self._api_key,
            input_text=text,
            model_endpoint=self._model_endpoint,
            model_name=self._model_name,
            conn_options=conn_options,
        )


class ChunkedStream(tts.ChunkedStream):
    def __init__(
        self,
        *,
        tts: TTS,
        api_key: str,
        model_endpoint: str,
        model_name: str,
        input_text: str,
        conn_options: APIConnectOptions,
    ) -> None:
        super().__init__(
            tts=tts,
            input_text=input_text,
            conn_options=conn_options,
        )

        self._tts: TTS = tts
        self._api_key = api_key
        self._model_endpoint = model_endpoint
        self._model_name = model_name
        self._opts = replace(tts._opts)

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        # TODO(laurent): once we support the POST requests, we should use it here rather than the websocket API.
        try:
            async with self._tts._ensure_session().ws_connect(
                self._tts._model_endpoint,
                headers={"x-api-key": self._tts._api_key},
                timeout=aiohttp.ClientWSTimeout(ws_receive=self._conn_options.timeout, ws_close=10),
            ) as ws:
                output_emitter.initialize(
                    request_id=utils.shortuuid(),
                    sample_rate=SUPPORTED_SAMPLE_RATE,
                    num_channels=1,
                    mime_type="audio/pcm",
                )

                setup_msg: dict[str, Any] = {
                    "type": "setup",
                    "model_name": self._model_name,
                    "output_format": "pcm",
                }
                if self._opts.voice is not None:
                    setup_msg["voice"] = self._opts.voice
                if self._opts.voice_id is not None:
                    setup_msg["voice_id"] = self._opts.voice_id
                if self._opts.pronunciation_id is not None:
                    setup_msg["pronunciation_id"] = self._opts.pronunciation_id
                if self._opts.json_config is not None:
                    setup_msg["json_config"] = json.dumps(self._opts.json_config)
                await ws.send_str(json.dumps(setup_msg))
                text_msg = {"type": "text", "text": self._input_text}
                await ws.send_str(json.dumps(text_msg))
                flush_msg = {"type": "end_of_stream"}
                await ws.send_str(json.dumps(flush_msg))

                while True:
                    msg = await ws.receive()
                    if msg.type in (
                        aiohttp.WSMsgType.CLOSE,
                        aiohttp.WSMsgType.CLOSED,
                        aiohttp.WSMsgType.CLOSING,
                    ):
                        # TODO(laurent): once we support returning eos in the api, we should enable this back.
                        # raise APIStatusError("Gradium websocket connection closed unexpectedly")
                        break

                    if msg.type == aiohttp.WSMsgType.TEXT:
                        msg_data = json.loads(msg.data)
                        type_ = msg_data.get("type")
                        if type_ == "text":
                            # TODO(laurent): handle text timestamps?
                            pass
                        elif type_ == "ready":
                            pass
                        elif type_ == "audio":
                            audio = base64.b64decode(msg_data["audio"])
                            output_emitter.push(audio)
                        elif type_ == "end_of_stream":
                            break
                        else:
                            logger.warning(f"unknown message type: {type_}")
                output_emitter.flush()

        except asyncio.TimeoutError:
            raise APITimeoutError() from None
        except aiohttp.ClientResponseError as e:
            raise APIStatusError(
                message=e.message, status_code=e.status, request_id=None, body=None
            ) from None
        except Exception as e:
            raise APIConnectionError() from e


class SynthesizeStream(tts.SynthesizeStream):
    """Stream-based text-to-speech synthesis using Gradium WebSocket API."""

    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:
        segments_ch = utils.aio.Chan[tokenize.WordStream]()
        request_id = utils.shortuuid()
        output_emitter.initialize(
            request_id=request_id,
            sample_rate=SUPPORTED_SAMPLE_RATE,
            num_channels=1,
            mime_type="audio/pcm",
            stream=True,
        )

        async def _tokenize_input() -> None:
            # Converts incoming text into WordStreams and sends them into _segments_ch
            word_stream = None
            async for input in self._input_ch:
                if isinstance(input, str):
                    if word_stream is None:
                        word_stream = self._opts.word_tokenizer.stream()
                        segments_ch.send_nowait(word_stream)
                    word_stream.push_text(input)
                elif isinstance(input, self._FlushSentinel):
                    if word_stream:
                        word_stream.end_input()
                    word_stream = None

            segments_ch.close()

        async def _run_segments() -> None:
            async for word_stream in segments_ch:
                await self._run_ws(word_stream, output_emitter)

        tasks = [
            asyncio.create_task(_tokenize_input()),
            asyncio.create_task(_run_segments()),
        ]
        try:
            await asyncio.gather(*tasks)
        except asyncio.TimeoutError:
            raise APITimeoutError() from None
        except aiohttp.ClientResponseError as e:
            raise APIStatusError(
                message=e.message, status_code=e.status, request_id=request_id, body=None
            ) from None
        except Exception as e:
            raise APIConnectionError() from e
        finally:
            await utils.aio.gracefully_cancel(*tasks)

    async def _run_ws(
        self, word_stream: tokenize.WordStream, output_emitter: tts.AudioEmitter
    ) -> None:
        segment_id = utils.shortuuid()
        output_emitter.start_segment(segment_id=segment_id)

        async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
            setup_msg: dict[str, Any] = {
                "type": "setup",
                "model_name": self._tts._model_name,
                "output_format": "pcm",
            }
            if self._opts.voice is not None:
                setup_msg["voice"] = self._opts.voice
            if self._opts.voice_id is not None:
                setup_msg["voice_id"] = self._opts.voice_id
            if self._opts.pronunciation_id is not None:
                setup_msg["pronunciation_id"] = self._opts.pronunciation_id
            if self._opts.json_config is not None:
                setup_msg["json_config"] = json.dumps(self._opts.json_config)

            await ws.send_str(json.dumps(setup_msg))

            async for word in word_stream:
                text_msg = {"type": "text", "text": f"{word.token} "}
                self._mark_started()
                await ws.send_str(json.dumps(text_msg))

            flush_msg = {"type": "end_of_stream"}
            await ws.send_str(json.dumps(flush_msg))

        async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
            while True:
                msg = await ws.receive()
                if msg.type in (
                    aiohttp.WSMsgType.CLOSE,
                    aiohttp.WSMsgType.CLOSED,
                    aiohttp.WSMsgType.CLOSING,
                ):
                    # TODO(laurent): once we support returning eos in the api, we should enable this back.
                    # raise APIStatusError("Gradium websocket connection closed unexpectedly")
                    output_emitter.end_segment()
                    break

                if msg.type == aiohttp.WSMsgType.TEXT:
                    msg_data = json.loads(msg.data)
                    type_ = msg_data.get("type")
                    if type_ == "text":
                        # TODO(laurent): handle text timestamps?
                        pass
                    elif type_ == "ready":
                        pass
                    elif type_ == "audio":
                        audio = base64.b64decode(msg_data["audio"])
                        output_emitter.push(audio)
                    elif type_ == "end_of_stream":
                        output_emitter.end_segment()
                        break
                    else:
                        logger.warning(f"unknown message type: {type_}")

        async with self._tts._ensure_session().ws_connect(
            self._tts._model_endpoint,
            headers={"x-api-key": self._tts._api_key},
            timeout=aiohttp.ClientWSTimeout(ws_receive=self._conn_options.timeout, ws_close=10),
        ) as ws:
            tasks = [
                asyncio.create_task(send_task(ws)),
                asyncio.create_task(recv_task(ws)),
            ]

            try:
                await asyncio.gather(*tasks)
            finally:
                await utils.aio.gracefully_cancel(*tasks)