"""Deepgram STT implementation with V1 streaming (interim + utterance_end) and REST fallback."""

from __future__ import annotations

import asyncio
import io
import logging
from typing import Callable

import numpy as np
import soundfile as sf

from roohai.base import STTModel

logger = logging.getLogger(__name__)


class DeepgramSTT(STTModel):
    META = {
        "display_name": "Deepgram",
        "description": "Cloud-based streaming speech recognition. Fast, accurate, requires API key.",
        "kind": "cloud",
        "pip_dependencies": ["deepgram-sdk>=6.0"],
        "hf_models": [],
        "estimated_download_mb": 0,
        "streaming": True,
        "config_fields": [
            {
                "name": "api_key",
                "label": "API Key",
                "type": "secret",
                "required": True,
                "env_var": "DEEPGRAM_API_KEY",
                "placeholder": "dg-xxx",
            },
            {
                "name": "model",
                "label": "Model",
                "type": "select",
                "required": False,
                "default": "nova-3",
                "options": ["nova-3", "nova-2"],
            },
            {
                "name": "language",
                "label": "Language",
                "type": "text",
                "required": False,
                "default": "en",
            },
        ],
    }

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "nova-3",
        language: str = "en",
        **kwargs,
    ) -> None:
        self.model = model
        self.language = language
        self._api_key = api_key
        self._client = None  # AsyncDeepgramClient
        self._ws = None  # active V1 streaming connection
        self._ws_ctx = None  # context manager for the connection
        self._listen_task: asyncio.Task | None = None
        self._on_transcript: Callable | None = None
        self._on_utterance_end: Callable | None = None
        self._last_final_transcript: str = ""

    def load(self) -> None:
        import os

        from deepgram import AsyncDeepgramClient

        key = self._api_key or os.environ.get("DEEPGRAM_API_KEY")
        if not key:
            raise RuntimeError(
                "Deepgram API key required. Set DEEPGRAM_API_KEY env var or pass api_key in config."
            )
        self._client = AsyncDeepgramClient(api_key=key)

    async def start_stream(
        self,
        on_transcript: Callable[[str, bool], None],
        on_utterance_end: Callable[[str], None],
    ) -> None:
        """Open Deepgram V1 streaming WebSocket with utterance detection.

        Args:
            on_transcript: called with (text, is_final) on each transcript update.
            on_utterance_end: called with (full_text) when an utterance ends.
        """
        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")

        self._on_transcript = on_transcript
        self._on_utterance_end = on_utterance_end
        self._last_final_transcript = ""

        # Open V1 streaming connection with interim results + utterance end detection
        self._ws_ctx = self._client.listen.v1.connect(
            model=str(self.model),
            encoding="linear16",
            sample_rate="16000",
            language=str(self.language),
            interim_results="true",
            utterance_end_ms="1500",
            vad_events="true",
            smart_format="true",
            punctuate="true",
        )
        self._ws = await self._ws_ctx.__aenter__()

        # Start background task to listen for transcript events
        self._listen_task = asyncio.create_task(self._listen_loop())
        logger.info("Deepgram V1 streaming connection opened (model=%s)", self.model)

    async def _listen_loop(self) -> None:
        """Background loop reading V1 events from the WebSocket."""
        from deepgram.listen.v1.types import (
            ListenV1Results,
            ListenV1UtteranceEnd,
        )

        msg_count = 0
        try:
            async for msg in self._ws:
                msg_count += 1
                msg_type = type(msg).__name__

                if isinstance(msg, ListenV1Results):
                    channel = msg.channel
                    if not channel or not channel.alternatives:
                        logger.debug("Deepgram: empty channel in result #%d", msg_count)
                        continue
                    transcript = channel.alternatives[0].transcript or ""
                    if not transcript:
                        continue

                    is_final = bool(msg.is_final)
                    speech_final = bool(msg.speech_final)
                    logger.info(
                        "Deepgram result #%d: final=%s speech_final=%s text='%s'",
                        msg_count, is_final, speech_final, transcript,
                    )

                    if is_final:
                        self._last_final_transcript = transcript
                        if self._on_transcript:
                            await self._on_transcript(transcript, True)
                        # speech_final means the speaker paused — treat as utterance end
                        if speech_final and self._on_utterance_end:
                            await self._on_utterance_end(transcript)
                            self._last_final_transcript = ""
                    else:
                        # Interim result
                        if self._on_transcript:
                            await self._on_transcript(transcript, False)

                elif isinstance(msg, ListenV1UtteranceEnd):
                    logger.info("Deepgram utterance_end event #%d (pending='%s')", msg_count, self._last_final_transcript)
                    # Utterance end event — fire if we have a pending final transcript
                    if self._last_final_transcript and self._on_utterance_end:
                        await self._on_utterance_end(self._last_final_transcript)
                        self._last_final_transcript = ""
                else:
                    logger.debug("Deepgram msg #%d type=%s", msg_count, msg_type)

            logger.warning("Deepgram listen loop ended normally after %d messages (connection closed?)", msg_count)

        except asyncio.CancelledError:
            logger.info("Deepgram listen loop cancelled after %d messages", msg_count)
        except Exception:
            logger.exception("Deepgram listen loop error after %d messages", msg_count)

    _chunk_count = 0

    async def send_chunk(self, audio_bytes: bytes) -> None:
        """Send raw PCM16 audio bytes to Deepgram."""
        if self._ws is not None:
            self._chunk_count += 1
            if self._chunk_count == 1:
                logger.info("Deepgram: sending first audio chunk (%d bytes)", len(audio_bytes))
            elif self._chunk_count % 200 == 0:
                logger.info("Deepgram: sent %d audio chunks so far", self._chunk_count)
            try:
                await self._ws.send_media(audio_bytes)
            except Exception:
                logger.exception("Deepgram: send_media failed at chunk %d", self._chunk_count)
        else:
            if self._chunk_count == 0:
                logger.warning("Deepgram: send_chunk called but _ws is None")

    async def stop_stream(self) -> None:
        """Close the streaming connection and clean up."""
        if self._ws is not None:
            try:
                await self._ws.send_close_stream()
            except Exception:
                logger.debug("Error sending close_stream", exc_info=True)

        if self._listen_task is not None:
            self._listen_task.cancel()
            try:
                await self._listen_task
            except asyncio.CancelledError:
                pass
            self._listen_task = None

        if self._ws_ctx is not None:
            try:
                await self._ws_ctx.__aexit__(None, None, None)
            except Exception:
                logger.debug("Error closing WS context", exc_info=True)
            self._ws_ctx = None
            self._ws = None

        self._on_transcript = None
        self._on_utterance_end = None
        self._last_final_transcript = ""
        logger.info("Deepgram streaming connection closed")

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> str:
        """Batch transcription fallback using REST API."""
        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load() first.")

        buf = io.BytesIO()
        sf.write(buf, audio, sample_rate, format="WAV")
        wav_bytes = buf.getvalue()

        from deepgram import DeepgramClient

        import os

        key = self._api_key or os.environ.get("DEEPGRAM_API_KEY", "")
        sync_client = DeepgramClient(api_key=key)
        response = sync_client.listen.v1.media.transcribe_file(
            request=wav_bytes,
            model=self.model,
            language=self.language,
            smart_format=True,
        )

        transcript = response.results.channels[0].alternatives[0].transcript
        return transcript.strip()

    def unload(self) -> None:
        if self._client is not None:
            if hasattr(self._client, 'close'):
                self._client.close()
            self._client = None
        self._ws = None
        self._ws_ctx = None

    @property
    def is_loaded(self) -> bool:
        return self._client is not None
