"""Deepgram STT client via SDKRouter server.

Real-time speech-to-text using Deepgram through SDKRouter server.
You only need SDKROUTER_API_KEY - the server handles Deepgram credentials.

Example:
    ```python
    from sdkrouter import AsyncSDKRouter

    sdk = AsyncSDKRouter(api_key="sk-...")

    async with sdk.audio.stt.stream_deepgram() as session:
        # Send audio chunks
        await session.send(audio_bytes)

        # Receive transcripts in real-time
        async for segment in session.transcripts():
            if segment.is_final:
                print(f">>> {segment.text}")
            else:
                print(f"... {segment.text}")
    ```
"""

from __future__ import annotations

import asyncio
import json
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from ....exceptions import SDKRouterError

if TYPE_CHECKING:
    from ...._config import SDKConfig

logger = logging.getLogger(__name__)


# =============================================================================
# Exceptions
# =============================================================================


class DeepgramConnectionError(SDKRouterError):
    """Failed to connect to Deepgram."""

    pass


class DeepgramStreamError(SDKRouterError):
    """Error during streaming."""

    def __init__(
        self,
        message: str,
        error_code: str = "STREAM_ERROR",
    ):
        super().__init__(message)
        self.error_code = error_code


# =============================================================================
# Types
# =============================================================================


class DeepgramConfig(BaseModel):
    """Configuration for Deepgram streaming STT."""

    model_config = ConfigDict(frozen=True)

    model: str = Field(default="nova-3", description="Deepgram model (nova-3, nova-2, etc.)")
    language: str = Field(default="en-US", description="Language code (en-US, es, etc.)")
    sample_rate: int = Field(default=16000, ge=8000, le=48000, description="Audio sample rate")
    channels: int = Field(default=1, ge=1, le=2, description="Audio channels")
    punctuate: bool = Field(default=True, description="Add punctuation")
    interim_results: bool = Field(default=True, description="Return interim results")
    diarize: bool = Field(default=False, description="Speaker diarization")
    smart_format: bool = Field(default=False, description="Smart formatting")
    # VAD parameters
    endpointing: int = Field(default=300, ge=0, le=5000, description="Silence detection threshold in ms")
    vad_events: bool = Field(default=True, description="Enable VAD events (SpeechStarted, UtteranceEnd)")
    utterance_end_ms: int | None = Field(default=None, ge=0, le=10000, description="Custom utterance boundary in ms")


class TranscriptWord(BaseModel):
    """Word-level transcript detail."""

    word: str
    start_time: float = 0.0
    end_time: float = 0.0
    confidence: float = 0.0
    speaker_id: int | None = None


class TranscriptSegment(BaseModel):
    """Transcript segment from Deepgram."""

    text: str
    is_final: bool = False
    confidence: float = 0.0
    start_time: float = 0.0
    end_time: float = 0.0
    words: list[TranscriptWord] = Field(default_factory=list)


class VADEvent(BaseModel):
    """Voice Activity Detection event from Deepgram."""

    type: Literal["speech_started", "utterance_end"]
    timestamp: float = 0.0
    last_word_end: float = 0.0


# =============================================================================
# Session
# =============================================================================


class DeepgramSession:
    """WebSocket session for Deepgram streaming STT.

    Handles the WebSocket protocol for streaming audio and receiving transcripts.
    """

    __slots__ = (
        "_ws",
        "_config",
        "_connected",
        "_transcript_queue",
        "_vad_event_queue",
        "_receive_task",
        "_session_id",
        "_total_duration",
    )

    def __init__(self, ws: Any, config: DeepgramConfig) -> None:
        self._ws = ws
        self._config = config
        self._connected = False
        self._transcript_queue: asyncio.Queue[TranscriptSegment | None] = asyncio.Queue()
        self._vad_event_queue: asyncio.Queue[VADEvent | None] = asyncio.Queue()
        self._receive_task: asyncio.Task[None] | None = None
        self._session_id: str | None = None
        self._total_duration: float = 0.0

    @property
    def is_connected(self) -> bool:
        """Check if session is connected."""
        return self._connected

    @property
    def session_id(self) -> str | None:
        """Get session ID from server."""
        return self._session_id

    @property
    def total_duration(self) -> float:
        """Get total audio duration processed."""
        return self._total_duration

    @property
    def config(self) -> DeepgramConfig:
        """Get session configuration."""
        return self._config

    async def _start(self) -> None:
        """Initialize the session - send config and wait for ready."""
        # Send config message
        config_msg = {
            "type": "config",
            "sample_rate": self._config.sample_rate,
            "channels": self._config.channels,
        }
        await self._ws.send(json.dumps(config_msg))

        # Wait for ready
        msg = await self._ws.recv()
        data = json.loads(msg)

        if data.get("type") == "error":
            raise DeepgramConnectionError(f"Server error: {data.get('message')}")

        if data.get("type") == "ready":
            self._session_id = data.get("session_id")
            self._connected = True
            # Start receiving transcripts in background
            self._receive_task = asyncio.create_task(self._receive_loop())
        else:
            raise DeepgramConnectionError(f"Unexpected message: {data}")

    async def send(self, audio_chunk: bytes) -> None:
        """Send audio chunk to Deepgram.

        Args:
            audio_chunk: Raw audio bytes (PCM 16-bit recommended)

        Raises:
            DeepgramStreamError: If session not connected
        """
        if not self._connected:
            raise DeepgramStreamError("Session not connected", "NOT_CONNECTED")
        await self._ws.send(audio_chunk)

    async def finalize(self) -> None:
        """Signal end of audio stream."""
        if self._connected:
            try:
                await self._ws.send(json.dumps({"type": "stop"}))
            except Exception:
                pass

    async def transcripts(self) -> AsyncIterator[TranscriptSegment]:
        """Iterate over transcript segments as they arrive.

        Yields interim and final results in real-time.

        Yields:
            TranscriptSegment with text, confidence, timing info
        """
        while self._connected:
            try:
                segment = await asyncio.wait_for(
                    self._transcript_queue.get(),
                    timeout=1.0,
                )
                if segment is None:
                    break
                yield segment
            except asyncio.TimeoutError:
                continue

    async def vad_events(self) -> AsyncIterator[VADEvent]:
        """Iterate over VAD events as they arrive.

        Yields VAD events: speech_started, utterance_end.
        Only available if vad_events=True in config.

        Yields:
            VADEvent with type and timing info
        """
        while self._connected:
            try:
                event = await asyncio.wait_for(
                    self._vad_event_queue.get(),
                    timeout=1.0,
                )
                if event is None:
                    break
                yield event
            except asyncio.TimeoutError:
                continue

    async def _receive_loop(self) -> None:
        """Background task to receive messages from server."""
        try:
            async for message in self._ws:
                if isinstance(message, bytes):
                    continue

                data = json.loads(message)
                msg_type = data.get("type")

                if msg_type == "transcript":
                    # Parse words
                    words = [
                        TranscriptWord(
                            word=w.get("word", ""),
                            start_time=w.get("start", 0.0),
                            end_time=w.get("end", 0.0),
                            confidence=w.get("confidence", 0.0),
                        )
                        for w in data.get("words", [])
                    ]

                    segment = TranscriptSegment(
                        text=data.get("text", ""),
                        is_final=data.get("is_final", False),
                        confidence=data.get("confidence", 0.0),
                        start_time=data.get("start_time", 0.0),
                        end_time=data.get("end_time", 0.0),
                        words=words,
                    )
                    await self._transcript_queue.put(segment)

                elif msg_type == "speech_started":
                    event = VADEvent(
                        type="speech_started",
                        timestamp=data.get("timestamp", 0.0),
                    )
                    await self._vad_event_queue.put(event)

                elif msg_type == "utterance_end":
                    event = VADEvent(
                        type="utterance_end",
                        last_word_end=data.get("last_word_end", 0.0),
                    )
                    await self._vad_event_queue.put(event)

                elif msg_type == "done":
                    self._total_duration = data.get("duration", 0.0)
                    logger.debug(f"Stream done, duration: {self._total_duration}s")
                    break

                elif msg_type == "error":
                    error_msg = data.get("message", "Unknown error")
                    logger.error(f"Server error: {error_msg}")
                    raise DeepgramStreamError(error_msg)

        except asyncio.CancelledError:
            pass
        except DeepgramStreamError:
            raise
        except Exception as e:
            logger.error(f"Receive error: {e}")
        finally:
            self._connected = False
            await self._transcript_queue.put(None)
            await self._vad_event_queue.put(None)

    async def close(self) -> None:
        """Close the session."""
        self._connected = False

        if self._receive_task:
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass

        try:
            await self._ws.close()
        except Exception:
            pass


# =============================================================================
# Client
# =============================================================================


class DeepgramClient:
    """Deepgram STT client via SDKRouter server.

    Connects to FastAPI server's WebSocket endpoint which handles
    Deepgram authentication. You only need SDKROUTER_API_KEY.

    Example:
        ```python
        from sdkrouter import AsyncSDKRouter

        sdk = AsyncSDKRouter(api_key="sk-...")

        # Via SDK
        async with sdk.audio.stt.stream_deepgram() as session:
            await session.send(audio_chunk)
            async for segment in session.transcripts():
                print(segment.text)
        ```
    """

    def __init__(self, sdk_config: "SDKConfig") -> None:
        self._sdk_config = sdk_config
        # Build WebSocket URL from audio_url
        audio_url = sdk_config.audio_url.rstrip("/")
        # Convert http(s) to ws(s)
        if audio_url.startswith("https://"):
            self._ws_base = audio_url.replace("https://", "wss://")
        else:
            self._ws_base = audio_url.replace("http://", "ws://")

    @asynccontextmanager
    async def stream(
        self,
        config: DeepgramConfig | None = None,
    ) -> AsyncIterator[DeepgramSession]:
        """Open streaming session to Deepgram.

        Args:
            config: Deepgram configuration (model, language, sample_rate, etc.)

        Yields:
            DeepgramSession for sending audio and receiving transcripts

        Raises:
            DeepgramConnectionError: If connection fails
            ImportError: If websockets not installed

        Example:
            ```python
            config = DeepgramConfig(model="nova-3", language="en-US")

            async with client.stream(config) as session:
                # Send audio
                await session.send(audio_bytes)

                # Get transcripts
                async for segment in session.transcripts():
                    print(segment.text)
            ```
        """
        try:
            from websockets.asyncio.client import connect as ws_connect
        except ImportError:
            raise ImportError(
                "websockets package required. Install with: pip install websockets"
            )

        config = config or DeepgramConfig()

        # Build WebSocket URL with query params
        params = [
            f"token={self._sdk_config.api_key}",
            f"model={config.model}",
            f"language={config.language}",
            f"punctuate={str(config.punctuate).lower()}",
            f"interim_results={str(config.interim_results).lower()}",
            f"diarize={str(config.diarize).lower()}",
            # VAD parameters
            f"endpointing={config.endpointing}",
            f"vad_events={str(config.vad_events).lower()}",
        ]
        if config.utterance_end_ms is not None:
            params.append(f"utterance_end_ms={config.utterance_end_ms}")
        url = f"{self._ws_base}/v1/audio/transcriptions/deepgram?{'&'.join(params)}"

        logger.debug(f"Connecting to Deepgram: {url.split('?')[0]}")

        try:
            ws = await ws_connect(url)
        except Exception as e:
            raise DeepgramConnectionError(f"Failed to connect: {e}") from e

        session = DeepgramSession(ws, config)

        try:
            await session._start()
            yield session
        finally:
            await session.finalize()
            await session.close()

    async def transcribe_file(
        self,
        audio_data: bytes,
        *,
        config: DeepgramConfig | None = None,
        chunk_size: int = 4096,
    ) -> list[TranscriptSegment]:
        """Transcribe audio file using streaming.

        Convenience method to transcribe complete audio data.

        Args:
            audio_data: Raw audio bytes (PCM).
            config: Streaming configuration.
            chunk_size: Size of chunks to send.

        Returns:
            List of final TranscriptSegment objects.
        """
        segments: list[TranscriptSegment] = []

        async with self.stream(config) as session:
            # Send audio in chunks
            for i in range(0, len(audio_data), chunk_size):
                chunk = audio_data[i:i + chunk_size]
                await session.send(chunk)
                await asyncio.sleep(0.02)

            # Collect transcripts
            async for segment in session.transcripts():
                if segment.is_final:
                    segments.append(segment)

        return segments


# =============================================================================
# Factory
# =============================================================================


def create_deepgram_client(sdk_config: "SDKConfig") -> DeepgramClient:
    """Create a Deepgram client from SDK config.

    Args:
        sdk_config: SDKRouter configuration

    Returns:
        DeepgramClient instance
    """
    return DeepgramClient(sdk_config)


__all__ = [
    # Client
    "DeepgramClient",
    "DeepgramSession",
    # Config & Types
    "DeepgramConfig",
    "TranscriptSegment",
    "TranscriptWord",
    "VADEvent",
    # Exceptions
    "DeepgramConnectionError",
    "DeepgramStreamError",
    # Factory
    "create_deepgram_client",
]
