"""
Deepgram STT Adapter

Streams 16kHz slin16 PCM audio to Deepgram's live transcription API
and fires on_transcript callbacks for both partial and final results.

Features:
- Nova-3 model family (2026 default — best Korean + English accuracy)
- Speaker diarization support
- Utterance end detection (endpointing)
- Keyterm prompting for domain-specific vocabulary

Deepgram live API format:
  Input:  raw 16-bit PCM, 16kHz, mono (matches DVGateway slin16)
  Output: JSON transcript events

API Endpoint: wss://api.deepgram.com/v1/listen
"""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass, field
from typing import Any, AsyncIterable, Callable
from urllib.parse import urlencode

import websockets
import websockets.client

from dvgateway.audio.codec import float32_to_slin16
from dvgateway.types import AudioChunk, SentimentResult, SttAdapter, TranscriptResult, TranscriptWord


@dataclass
class DeepgramAdapterOptions:
    api_key: str = ""
    language: str = "ko"
    model: str = "nova-3"
    diarize: bool = False
    endpointing_ms: int = 500
    interim_results: bool = True
    smart_format: bool = True
    keywords: list[str] = field(default_factory=list)
    punctuate: bool = True
    profanity_filter: bool = False
    utterance_end_ms: int = 1000
    sentiment: bool = False


class DeepgramAdapter(SttAdapter):
    """Deepgram STT adapter implementing the SttAdapter interface."""

    def __init__(
        self,
        api_key: str = "",
        language: str = "ko",
        model: str = "nova-3",
        diarize: bool = False,
        endpointing_ms: int = 500,
        interim_results: bool = True,
        smart_format: bool = True,
        keywords: list[str] | None = None,
        punctuate: bool = True,
        profanity_filter: bool = False,
        utterance_end_ms: int = 1000,
        sentiment: bool = False,
    ) -> None:
        self._opts = DeepgramAdapterOptions(
            api_key=api_key,
            language=language,
            model=model,
            diarize=diarize,
            endpointing_ms=endpointing_ms,
            interim_results=interim_results,
            smart_format=smart_format,
            keywords=keywords or [],
            punctuate=punctuate,
            profanity_filter=profanity_filter,
            utterance_end_ms=utterance_end_ms,
            sentiment=sentiment,
        )
        self._transcript_handler: Callable[[TranscriptResult], None] | None = None
        self._ws: websockets.client.WebSocketClientProtocol | None = None  # type: ignore[type-arg]
        self._stopped = False

    def on_transcript(self, handler: Callable[[TranscriptResult], None]) -> None:
        self._transcript_handler = handler

    async def start_stream(self, linked_id: str, audio_stream: AsyncIterable[AudioChunk]) -> None:
        self._stopped = False
        self._ws = await self._connect_deepgram()

        # Pipe audio in background
        asyncio.ensure_future(self._pipe_audio(linked_id, audio_stream))

        # Keep alive until stream ends or stop() is called
        assert self._ws is not None
        try:
            async for message in self._ws:
                if self._stopped:
                    break
                try:
                    msg = json.loads(message if isinstance(message, str) else message.decode("utf-8"))
                    self._handle_transcript(msg)
                except (json.JSONDecodeError, UnicodeDecodeError):
                    pass
        except websockets.exceptions.ConnectionClosed:
            pass

    async def stop(self) -> None:
        self._stopped = True
        if self._ws:
            try:
                await self._ws.send(json.dumps({"type": "CloseStream"}))
                await asyncio.sleep(0.2)
                await self._ws.close(1000, "stop")
            except Exception:
                pass
            self._ws = None

    async def _connect_deepgram(self) -> websockets.client.WebSocketClientProtocol:  # type: ignore[type-arg]
        params: dict[str, str] = {
            "model": self._opts.model,
            "language": self._opts.language,
            "encoding": "linear16",
            "sample_rate": "16000",
            "channels": "1",
            "interim_results": str(self._opts.interim_results).lower(),
            "smart_format": str(self._opts.smart_format).lower(),
            "endpointing": str(self._opts.endpointing_ms),
            "punctuate": str(self._opts.punctuate).lower(),
            "utterance_end_ms": str(self._opts.utterance_end_ms),
        }
        if self._opts.diarize:
            params["diarize"] = "true"
        if self._opts.profanity_filter:
            params["profanity_filter"] = "true"
        if self._opts.sentiment:
            params["sentiment"] = "true"

        url = f"wss://api.deepgram.com/v1/listen?{urlencode(params)}"

        # Append keywords
        for kw in self._opts.keywords:
            url += f"&keyterm={kw}"

        ws = await websockets.client.connect(
            url,
            extra_headers={"Authorization": f"Token {self._opts.api_key}"},
            ping_interval=30,
            ping_timeout=60,
        )
        return ws

    async def _pipe_audio(self, linked_id: str, audio_stream: AsyncIterable[AudioChunk]) -> None:
        async for chunk in audio_stream:
            if self._stopped or not self._ws:
                break
            pcm = float32_to_slin16(chunk.samples)
            try:
                await self._ws.send(pcm)
            except Exception:
                break

    def _handle_transcript(self, msg: dict[str, Any]) -> None:
        if msg.get("type") != "Results":
            # Log non-Results messages to check if sentiment comes separately
            import sys
            print(f"[DeepgramDebug] non-Results msg type={msg.get('type')}, keys={list(msg.keys())}", file=sys.stderr)
            return

        alt = None
        channel = msg.get("channel", {})
        alternatives = channel.get("alternatives", [])
        if alternatives:
            alt = alternatives[0]

        if not alt or not alt.get("transcript"):
            return

        is_final = msg.get("speech_final", False) or msg.get("is_final", False)

        words: list[TranscriptWord] | None = None
        raw_words = alt.get("words", [])
        if raw_words:
            words = [
                TranscriptWord(
                    word=w.get("punctuated_word", w.get("word", "")),
                    start_ms=round(w.get("start", 0) * 1000),
                    end_ms=round(w.get("end", 0) * 1000),
                    confidence=w.get("confidence"),
                )
                for w in raw_words
            ]

        # Debug: dump sentiment-related fields from raw response
        import sys
        if self._opts.sentiment and is_final:
            word0 = raw_words[0] if raw_words else {}
            print(
                f"[DeepgramDebug] sentiment check:"
                f" msg.sentiments={msg.get('sentiments')}"
                f" word0.keys={list(word0.keys())}"
                f" word0.sentiment={word0.get('sentiment')}"
                f" alt.keys={list(alt.keys())}",
                file=sys.stderr,
            )

        speaker: str | None = None
        if self._opts.diarize and raw_words and raw_words[0].get("speaker") is not None:
            speaker = f"speaker_{raw_words[0]['speaker']}"

        # Extract sentiment if enabled
        #   Streaming API: sentiment is on each word in alternatives[].words[]
        #   Pre-recorded API: sentiment is in top-level sentiments.segments/average
        sentiment_result: SentimentResult | None = None
        if self._opts.sentiment:
            # 1) Try streaming format: word-level sentiment → aggregate
            word_sentiments = [
                w for w in raw_words if w.get("sentiment") is not None
            ]
            if word_sentiments:
                avg_score = sum(w.get("sentiment_score", 0.0) for w in word_sentiments) / len(word_sentiments)
                if avg_score >= 0.333333:
                    label = "positive"
                elif avg_score <= -0.333333:
                    label = "negative"
                else:
                    label = "neutral"
                # Normalise Deepgram score (-1…1) → SDK score (0…1)
                normalised = (avg_score + 1.0) / 2.0
                sentiment_result = SentimentResult(
                    sentiment=label,
                    sentiment_score=normalised,
                )
            else:
                # 2) Fallback: pre-recorded / batch format
                sentiments = msg.get("sentiments", {})
                segments = sentiments.get("segments", [])
                seg = segments[0] if segments else sentiments.get("average")
                if seg:
                    raw_score = seg.get("sentiment_score", 0.0)
                    sentiment_result = SentimentResult(
                        sentiment=seg.get("sentiment", "neutral"),
                        sentiment_score=(raw_score + 1.0) / 2.0,
                    )

        import time

        result = TranscriptResult(
            linked_id="",
            text=alt["transcript"],
            is_final=is_final,
            confidence=alt.get("confidence"),
            language=self._opts.language,
            timestamp_ms=time.time() * 1000,
            words=words,
            speaker=speaker,
            sentiment=sentiment_result,
        )

        if self._transcript_handler:
            self._transcript_handler(result)
