# 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 service for Pipecat."""

from __future__ import annotations

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

import aiohttp
from pipecat.frames.frames import (
    ErrorFrame,
    Frame,
    TTSAudioRawFrame,
    TTSStartedFrame,
    TTSStoppedFrame,
)
from pipecat.services.tts_service import TTSService

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

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


@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


class KugelAudioTTSService(TTSService):
    """KugelAudio Text-to-Speech service for Pipecat.

    Integrates KugelAudio's TTS API with Pipecat's pipeline framework,
    providing high-quality voice synthesis via WebSocket streaming.

    Example:
        ```python
        from kugelaudio.pipecat import KugelAudioTTSService

        tts = KugelAudioTTSService(
            api_key="your-api-key",
            model="kugel-1-turbo",
            sample_rate=24000,
        )

        # Use in a Pipecat pipeline
        pipeline = Pipeline([..., tts, ...])
        ```
    """

    def __init__(
        self,
        *,
        api_key: Optional[str] = None,
        model: TTSModels | str = DEFAULT_MODEL,
        voice_id: Optional[int] = None,
        sample_rate: int = DEFAULT_SAMPLE_RATE,
        cfg_scale: float = DEFAULT_CFG_SCALE,
        max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
        base_url: str = "https://api.kugelaudio.com",
        aiohttp_session: Optional[aiohttp.ClientSession] = None,
        **kwargs: Any,
    ) -> None:
        """Create a new KugelAudio TTS service for Pipecat.

        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.
            cfg_scale: CFG scale for generation quality. Defaults to 2.0.
            max_new_tokens: Maximum tokens to generate. Defaults to 2048.
            base_url: API base URL. Defaults to https://api.kugelaudio.com.
            aiohttp_session: Optional aiohttp session to reuse.
            **kwargs: Additional arguments passed to Pipecat TTSService.
        """
        super().__init__(sample_rate=sample_rate, **kwargs)

        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"
            )

        if sample_rate not in SUPPORTED_SAMPLE_RATES:
            raise ValueError(
                f"Unsupported sample rate: {sample_rate}. "
                f"Supported rates: {SUPPORTED_SAMPLE_RATES}"
            )

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

        self._session = aiohttp_session

    def can_generate_metrics(self) -> bool:
        """Check if this service can generate processing metrics.

        Returns:
            True, as KugelAudio service supports metrics generation.
        """
        return True

    def _ensure_session(self) -> aiohttp.ClientSession:
        """Ensure an aiohttp session exists, creating one if needed."""
        if not self._session:
            self._session = aiohttp.ClientSession()
        return self._session

    async def set_model(self, model: str) -> None:
        """Set the TTS model.

        Args:
            model: The model identifier to use.
        """
        self._opts.model = model
        await super().set_model(model)

    def set_voice(self, voice: str) -> None:
        """Set the voice ID.

        Args:
            voice: The voice identifier (integer as string).
        """
        self._opts.voice_id = int(voice) if voice else None
        super().set_voice(voice)

    async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
        """Synthesize text to speech via WebSocket streaming.

        Connects to the KugelAudio WebSocket TTS endpoint, sends the text,
        and yields audio frames as they arrive.

        Args:
            text: The text to synthesize into speech.

        Yields:
            Frame: TTSStartedFrame, TTSAudioRawFrame chunks, and TTSStoppedFrame.
        """
        logger.debug(f"Generating TTS [{text}]")

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

        request_data: dict[str, Any] = {
            "text": 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,
        }
        if self._opts.voice_id is not None:
            request_data["voice_id"] = self._opts.voice_id

        try:
            await self.start_ttfb_metrics()
            await self.start_tts_usage_metrics(text)

            session = self._ensure_session()
            async with session.ws_connect(
                ws_url,
                timeout=aiohttp.ClientTimeout(total=60, sock_connect=30),
            ) as ws:
                await ws.send_str(json.dumps(request_data))

                yield TTSStartedFrame()

                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        data = json.loads(msg.data)

                        if data.get("error"):
                            logger.error(f"KugelAudio error: {data['error']}")
                            yield ErrorFrame(
                                error=f"KugelAudio API error: {data['error']}"
                            )
                            return

                        if data.get("audio"):
                            audio_bytes = base64.b64decode(data["audio"])
                            await self.stop_ttfb_metrics()
                            yield TTSAudioRawFrame(
                                audio=audio_bytes,
                                sample_rate=self._opts.sample_rate,
                                num_channels=1,
                            )

                        if data.get("final"):
                            break

                    elif msg.type in (
                        aiohttp.WSMsgType.CLOSED,
                        aiohttp.WSMsgType.ERROR,
                    ):
                        logger.warning(f"WebSocket closed unexpectedly: {msg.type}")
                        break

        except aiohttp.ClientError as e:
            logger.error(f"Connection error: {e}")
            yield ErrorFrame(error=f"KugelAudio connection error: {e}")
        except Exception as e:
            logger.error(f"Unexpected error: {e}")
            yield ErrorFrame(error=f"KugelAudio error: {e}")
        finally:
            await self.stop_ttfb_metrics()
            yield TTSStoppedFrame()
            logger.debug(f"Finished TTS [{text}]")
