# Copyright 2024 KugelAudio
#
# 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.

"""KugelAudio TTS plugin for LiveKit Agents.

Performance optimizations:
- Persistent WebSocket connection reuse for /ws/tts (non-streaming).
  The server supports multiple requests over a single WebSocket, so we
  keep one connection alive across synthesize() calls to avoid repeated
  TCP+TLS+WS handshake overhead (~35-290ms per call).
- prewarm() method to pre-establish the WebSocket connection and
  optionally warm the server-side voice cache before the first real
  synthesis request.
"""

from __future__ import annotations

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

import aiohttp
from livekit.agents import (
    APIConnectionError,
    APIConnectOptions,
    APIStatusError,
    APITimeoutError,
    tts,
    utils,
)
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, NotGivenOr
from livekit.agents.utils import is_given
from livekit.agents.voice.io import TimedString

from .models import (
    DEFAULT_MODEL,
    DEFAULT_SAMPLE_RATE,
    SUPPORTED_SAMPLE_RATES,
    TTSModels,
)

logger = logging.getLogger("kugelaudio.livekit")

# Timeout in seconds before an idle persistent connection is closed.
_WS_INACTIVITY_TIMEOUT = 300


def _word_timestamps_to_timed(
    timestamps: list[dict[str, Any]],
) -> list[TimedString]:
    """Convert server word_timestamps payload to LiveKit TimedString list.

    The server sends timestamps with ``start_ms`` / ``end_ms`` (integers),
    while LiveKit expects ``start_time`` / ``end_time`` in seconds (floats).
    """
    return [
        TimedString(
            ts["word"],
            start_time=ts["start_ms"] / 1000.0,
            end_time=ts["end_ms"] / 1000.0,
        )
        for ts in timestamps
    ]


@dataclass
class _TTSOptions:
    model: TTSModels | str
    voice_id: int | None
    sample_rate: int
    cfg_scale: float
    max_new_tokens: int
    api_key: str
    base_url: str
    word_timestamps: bool = True
    normalize: bool = True


class _PersistentConnection:
    """Manages a persistent WebSocket connection to the /ws/tts endpoint.

    The server-side /ws/tts endpoint supports multiple sequential requests
    on a single WebSocket connection (it runs a ``while True`` receive loop).
    By keeping the connection alive, we avoid per-request TCP+TLS+WS
    handshake overhead (~35-290ms depending on network conditions).
    """

    def __init__(
        self,
        opts: _TTSOptions,
        session: aiohttp.ClientSession,
    ) -> None:
        self._opts = opts
        self._session = session
        self._ws: Optional[aiohttp.ClientWebSocketResponse] = None
        self._lock = asyncio.Lock()
        self._closed = False

    @property
    def connected(self) -> bool:
        return self._ws is not None and not self._ws.closed

    async def ensure_connected(self, timeout: float = 10.0) -> None:
        """Ensure the WebSocket connection is established.

        This is idempotent: if already connected, returns immediately.
        Thread-safe via asyncio.Lock.
        """
        if self.connected:
            return

        async with self._lock:
            # Double-check after acquiring the lock
            if self.connected:
                return
            if self._closed:
                raise APIConnectionError(
                    "Connection has been closed, create a new TTS instance"
                )
            await self._connect(timeout)

    async def _connect(self, timeout: float) -> None:
        """Establish the WebSocket connection."""
        ws_url = self._opts.base_url.replace("https://", "wss://").replace(
            "http://", "ws://"
        )
        ws_url = (
            f"{ws_url}/ws/tts" f"?api_key={self._opts.api_key}&model={self._opts.model}"
        )

        try:
            self._ws = await self._session.ws_connect(
                ws_url,
                timeout=aiohttp.ClientTimeout(total=None, sock_connect=timeout),
            )
            logger.debug("Persistent WS connection established to /ws/tts")
        except asyncio.TimeoutError:
            raise APITimeoutError() from None
        except aiohttp.ClientError as e:
            raise APIConnectionError(
                f"Failed to establish persistent WS connection: {e}"
            ) from e

    async def send_request(self, request_data: dict[str, Any]) -> None:
        """Send a synthesis request over the persistent connection."""
        if not self.connected:
            raise APIConnectionError("Not connected")
        assert self._ws is not None
        await self._ws.send_str(json.dumps(request_data))

    async def receive(self) -> aiohttp.WSMessage:
        """Receive a message from the persistent connection."""
        if not self.connected:
            raise APIConnectionError("Not connected")
        assert self._ws is not None
        return await self._ws.receive()

    async def aclose(self) -> None:
        """Close the persistent connection."""
        self._closed = True
        if self._ws and not self._ws.closed:
            await self._ws.close()
            logger.debug("Persistent WS connection closed")
        self._ws = None


class TTS(tts.TTS):
    """KugelAudio Text-to-Speech plugin for LiveKit Agents.

    This plugin integrates KugelAudio's TTS API with LiveKit's agent framework,
    providing high-quality voice synthesis with streaming support.

    Example:
        ```python
        from kugelaudio.livekit import TTS

        tts = TTS(api_key="your-api-key")

        # Use with VoicePipelineAgent
        agent = VoicePipelineAgent(
            tts=tts,
            ...
        )
        ```

    Or via the livekit.plugins namespace (after registering):
        ```python
        from kugelaudio.livekit import register_plugin
        register_plugin()

        from livekit.plugins import kugelaudio
        tts = kugelaudio.TTS()
        ```
    """

    def __init__(
        self,
        *,
        api_key: str | None = None,
        model: TTSModels | str = DEFAULT_MODEL,
        voice_id: int | None = None,
        sample_rate: int = DEFAULT_SAMPLE_RATE,
        cfg_scale: float = 2.0,
        max_new_tokens: int = 2048,
        normalize: bool = True,
        base_url: str = "https://api.kugelaudio.com",
        http_session: aiohttp.ClientSession | None = None,
    ) -> None:
        """Create a new KugelAudio TTS instance.

        Args:
            api_key: KugelAudio API key. If not provided, reads from
                KUGELAUDIO_API_KEY environment variable.
            model: TTS model to use. "kugel-1-turbo" (fast, 1.5B) or
                "kugel-1" (premium, 7B).
            voice_id: Voice ID to use. If None, uses server default.
            sample_rate: Output sample rate in Hz. Supported rates: 24000 (native),
                22050, 16000, 8000. Lower rates use server-side resampling with
                minimal latency impact (~0.1ms per 100ms of audio).
            cfg_scale: CFG scale for generation quality. Defaults to 2.0.
            max_new_tokens: Maximum tokens to generate. Defaults to 2048.
            normalize: Apply loudness normalization to the output audio. Defaults
                to True.
            base_url: API base URL. Defaults to https://api.kugelaudio.com.
            http_session: Optional aiohttp session to reuse.
        """
        super().__init__(
            capabilities=tts.TTSCapabilities(
                streaming=True,
                aligned_transcript=True,
            ),
            sample_rate=sample_rate,
            num_channels=1,
        )

        kugelaudio_api_key = api_key or os.environ.get("KUGELAUDIO_API_KEY")
        if not kugelaudio_api_key:
            raise ValueError(
                "KUGELAUDIO_API_KEY must be set or api_key must be provided"
            )

        self._opts = _TTSOptions(
            model=model,
            voice_id=voice_id,
            sample_rate=sample_rate,
            cfg_scale=cfg_scale,
            max_new_tokens=max_new_tokens,
            normalize=normalize,
            api_key=kugelaudio_api_key,
            base_url=base_url.rstrip("/"),
        )

        self._session = http_session
        self._persistent_conn: Optional[_PersistentConnection] = None
        self._conn_lock = asyncio.Lock()

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

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

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

    async def _ensure_persistent_conn(
        self, timeout: float = 10.0
    ) -> _PersistentConnection:
        """Get or create the persistent WebSocket connection.

        The connection is reused across multiple synthesize() calls.
        If the connection is dropped, a new one is created transparently.
        """
        async with self._conn_lock:
            if self._persistent_conn and self._persistent_conn.connected:
                return self._persistent_conn

            session = self._ensure_session()
            conn = _PersistentConnection(self._opts, session)
            await conn.ensure_connected(timeout)
            self._persistent_conn = conn
            return conn

    def prewarm(self) -> None:
        """Pre-warm the WebSocket connection to reduce TTFA on first synthesis.

        Call this from LiveKit's ``before_tts_cb`` or during agent setup
        to eagerly establish the WebSocket connection and authenticate,
        eliminating ~35-290ms of handshake latency from the first request.

        This is a synchronous method (as required by the LiveKit TTS interface)
        that schedules the async connection establishment in the background.

        Example:
            ```python
            tts = TTS(api_key="your-api-key")
            tts.prewarm()  # establishes WS connection in background
            ```
        """
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            logger.warning("prewarm() called outside event loop, skipping")
            return

        async def _do_prewarm() -> None:
            try:
                await self._ensure_persistent_conn(timeout=10.0)
                logger.info("KugelAudio TTS connection pre-warmed")
            except Exception:
                logger.warning(
                    "Failed to prewarm KugelAudio connection, "
                    "will retry on first synthesis",
                    exc_info=True,
                )

        loop.create_task(_do_prewarm())

    def update_options(
        self,
        *,
        model: NotGivenOr[TTSModels | str] = NOT_GIVEN,
        voice_id: NotGivenOr[int | None] = NOT_GIVEN,
        cfg_scale: NotGivenOr[float] = NOT_GIVEN,
        max_new_tokens: NotGivenOr[int] = NOT_GIVEN,
        normalize: NotGivenOr[bool] = NOT_GIVEN,
    ) -> None:
        """Update TTS options dynamically.

        Args:
            model: TTS model to use.
            voice_id: Voice ID to use.
            cfg_scale: CFG scale for generation.
            max_new_tokens: Maximum tokens to generate.
            normalize: Apply loudness normalization to the output audio.
        """
        if is_given(model):
            self._opts.model = model
        if is_given(voice_id):
            self._opts.voice_id = voice_id
        if is_given(cfg_scale):
            self._opts.cfg_scale = cfg_scale
        if is_given(max_new_tokens):
            self._opts.max_new_tokens = max_new_tokens
        if is_given(normalize):
            self._opts.normalize = normalize

    def synthesize(
        self,
        text: str,
        *,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> ChunkedStream:
        """Synthesize text to speech (non-streaming).

        Args:
            text: Text to synthesize.
            conn_options: Connection options.

        Returns:
            ChunkedStream that yields audio frames.
        """
        return ChunkedStream(tts=self, input_text=text, conn_options=conn_options)

    def stream(
        self,
        *,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> "SynthesizeStream":
        """Create a streaming TTS session.

        Args:
            conn_options: Connection options.

        Returns:
            SynthesizeStream for streaming text input.
        """
        return SynthesizeStream(tts=self, conn_options=conn_options)

    async def aclose(self) -> None:
        """Close the TTS instance and release resources.

        Closes the persistent WebSocket connection if one is active.
        """
        if self._persistent_conn:
            await self._persistent_conn.aclose()
            self._persistent_conn = None


class ChunkedStream(tts.ChunkedStream):
    """Non-streaming TTS synthesis using WebSocket."""

    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
        self._opts = replace(tts._opts)

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        """Run the synthesis and emit audio frames.

        Uses the persistent WebSocket connection from the TTS instance
        to avoid per-request handshake overhead. Falls back to a fresh
        connection if the persistent one is unavailable.
        """
        request_id = utils.shortuuid()

        output_emitter.initialize(
            request_id=request_id,
            sample_rate=self._opts.sample_rate,
            num_channels=1,
            mime_type="audio/pcm",
            # Use a small frame size to minimise buffering delay.
            # The default (200ms) adds ~200ms to perceived TTFA because
            # AudioByteStream waits until a full frame is accumulated
            # before emitting the first SynthesizedAudio event.
            frame_size_ms=20,
        )

        request_data = {
            "text": self._input_text,
            "model_id": self._opts.model,
            "cfg_scale": self._opts.cfg_scale,
            "max_new_tokens": self._opts.max_new_tokens,
            "sample_rate": self._opts.sample_rate,
            "word_timestamps": self._opts.word_timestamps,
            "normalize": self._opts.normalize,
        }
        if self._opts.voice_id is not None:
            request_data["voice_id"] = self._opts.voice_id

        try:
            conn = await self._tts._ensure_persistent_conn(
                timeout=self._conn_options.timeout
            )
            await conn.send_request(request_data)

            while True:
                msg = await conn.receive()

                if msg.type == aiohttp.WSMsgType.TEXT:
                    data = json.loads(msg.data)

                    if data.get("error"):
                        raise APIStatusError(
                            message=data["error"],
                            status_code=400,
                            request_id=request_id,
                            body=None,
                        )

                    if data.get("audio"):
                        audio_bytes = base64.b64decode(data["audio"])
                        output_emitter.push(audio_bytes)

                    if data.get("word_timestamps"):
                        timed_strings = _word_timestamps_to_timed(
                            data["word_timestamps"]
                        )
                        output_emitter.push_timed_transcript(timed_strings)

                    if data.get("final"):
                        break

                elif msg.type in (
                    aiohttp.WSMsgType.CLOSED,
                    aiohttp.WSMsgType.ERROR,
                ):
                    # Connection dropped — invalidate so next call reconnects
                    self._tts._persistent_conn = None
                    raise APIConnectionError("WebSocket connection closed unexpectedly")

            output_emitter.flush()

        except asyncio.TimeoutError:
            raise APITimeoutError() from None
        except aiohttp.ClientError as e:
            # Invalidate broken connection so it is re-created next call
            self._tts._persistent_conn = None
            raise APIConnectionError() from e


class SynthesizeStream(tts.SynthesizeStream):
    """Streaming TTS synthesis using the /ws/tts/stream WebSocket endpoint.

    This uses KugelAudio's streaming endpoint which supports:
    - Continuous text input (token-by-token from LLM)
    - Server-side sentence boundary detection and buffering
    - Flush commands to force generation of buffered text
    - Close commands for clean session shutdown

    Protocol:
        Client -> Server:
            First message includes config + text:
                {"text": "...", "voice_id": 123, "sample_rate": 24000, ...}
            Subsequent messages:
                {"text": "..."}              - Add text to buffer
                {"text": "...", "flush": true} - Add text and force flush
                {"flush": true}              - Flush current buffer
                {"close": true}              - End session

        Server -> Client:
            {"generation_started": true, "chunk_id": 0, "text": "..."}
            {"audio": "base64...", "chunk_id": 0, "idx": 0, ...}
            {"chunk_complete": true, "chunk_id": 0, ...}
            {"session_closed": true, "total_audio_seconds": 5.4, ...}
    """

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

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        """Run the streaming synthesis using /ws/tts/stream."""
        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,
        )

        ws_url = self._opts.base_url.replace("https://", "wss://").replace(
            "http://", "ws://"
        )
        ws_url = (
            f"{ws_url}/ws/tts/stream"
            f"?api_key={self._opts.api_key}&model={self._opts.model}"
        )

        try:
            session = self._tts._ensure_session()
            async with session.ws_connect(
                ws_url,
                timeout=aiohttp.ClientTimeout(total=None, sock_connect=30),
            ) as ws:
                segment_id = utils.shortuuid()
                output_emitter.start_segment(segment_id=segment_id)

                async def send_task() -> None:
                    """Stream text tokens to WebSocket, flush on sentinel."""
                    config_sent = False

                    async for data in self._input_ch:
                        if isinstance(data, self._FlushSentinel):
                            await ws.send_str(json.dumps({"flush": True}))
                            continue

                        if not data:
                            continue
                        msg: dict[str, Any] = {"text": data}

                        # Attach config on first message
                        if not config_sent:
                            msg.update(self._build_config())
                            config_sent = True

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

                    # Input channel closed -> flush remaining + close session
                    await ws.send_str(json.dumps({"flush": True}))
                    await ws.send_str(json.dumps({"close": True}))

                async def recv_task() -> None:
                    """Receive audio chunks from WebSocket."""
                    async for msg in ws:
                        if msg.type == aiohttp.WSMsgType.TEXT:
                            data = json.loads(msg.data)

                            if data.get("error"):
                                logger.error("KugelAudio error: %s", data["error"])
                                raise APIStatusError(
                                    message=data["error"],
                                    status_code=400,
                                    request_id=request_id,
                                    body=None,
                                )

                            if data.get("generation_started"):
                                # Server started generating for a text chunk
                                continue

                            if data.get("audio"):
                                audio_bytes = base64.b64decode(data["audio"])
                                self._mark_started()
                                output_emitter.push(audio_bytes)

                            if data.get("word_timestamps"):
                                timed_strings = _word_timestamps_to_timed(
                                    data["word_timestamps"]
                                )
                                output_emitter.push_timed_transcript(timed_strings)

                            if data.get("chunk_complete"):
                                # One text chunk is done, flush to emit segment
                                output_emitter.flush()

                            if data.get("session_closed"):
                                # Server confirmed session end
                                break

                        elif msg.type in (
                            aiohttp.WSMsgType.CLOSED,
                            aiohttp.WSMsgType.ERROR,
                        ):
                            break

                send = asyncio.create_task(send_task())
                recv = asyncio.create_task(recv_task())

                try:
                    await asyncio.gather(send, recv)
                finally:
                    await utils.aio.gracefully_cancel(send, recv)
                    output_emitter.flush()
                    output_emitter.end_segment()

        except asyncio.TimeoutError:
            raise APITimeoutError() from None
        except aiohttp.ClientError as e:
            raise APIConnectionError() from e

    def _build_config(self) -> dict[str, Any]:
        """Build config dict to send with the first message."""
        config: dict[str, Any] = {
            "sample_rate": self._opts.sample_rate,
            "cfg_scale": self._opts.cfg_scale,
            "max_new_tokens": self._opts.max_new_tokens,
            "model_id": self._opts.model,
            "word_timestamps": self._opts.word_timestamps,
            "normalize": self._opts.normalize,
        }
        if self._opts.voice_id is not None:
            config["voice_id"] = self._opts.voice_id
        return config
