"""
NVIDIA realtime STT provider for the Voice Assistant.

Purpose:
    Implement the generic STTAdapter contract using NVIDIA Speech NIM.

Responsibilities:
    - Create and manage NVIDIA realtime transcription sessions.
    - Send canonical PCM audio to the NVIDIA ASR service.
    - Receive provider-specific transcription events.
    - Translate NVIDIA events into provider-independent TranscriptEvent objects.
    - Manage provider connection lifecycle and cleanup.

Non-responsibilities:
    - Managing WebSocket client connections to the Audio Gateway.
    - Managing Voice Assistant sessions or turns.
    - Executing Agent or TTS functionality.
    - Defining the canonical Voice Core audio format.

Architecture:
    Voice Core
        |
        | canonical PCM audio
        v
    STTAdapter
        |
        v
    NvidiaSTTAdapter
        |
        v
    NVIDIA Speech NIM

    NVIDIA-specific HTTP, WebSocket, session, and event details remain
    isolated inside this adapter.
"""

from __future__ import annotations

import asyncio
import base64
import json
from collections.abc import AsyncIterator
from typing import Any

import httpx
import websockets
from websockets.asyncio.client import ClientConnection

from voice_assistant.audio.models import CANONICAL_AUDIO_FORMAT

from .adapter import STTAdapter
from .models import TranscriptEvent


class NvidiaSTTAdapter(STTAdapter):
    """Realtime STT adapter for NVIDIA Speech NIM."""

    def __init__(
        self,
        server: str = "127.0.0.1:9000",
        language_code: str = "pt-BR",
        automatic_punctuation: bool = True,
    ) -> None:
        self.server = server
        self.language_code = language_code
        self.automatic_punctuation = automatic_punctuation

        self._http: httpx.AsyncClient | None = None
        self._websocket: ClientConnection | None = None
        self._receiver_task: asyncio.Task[None] | None = None

        self._events: asyncio.Queue[TranscriptEvent | None] = asyncio.Queue()

        self._started = False
        self._stopping = False
        self._aborted = False
        self._abort_lock = asyncio.Lock()


    async def start_session(self) -> None:
        if self._started:
            raise RuntimeError("STT session is already started.")

        self._http = httpx.AsyncClient(timeout=10.0)

        try:
            session = await self._create_session()
            self._websocket = await self._connect()

            await self._initialize_session(session)

            self._started = True
            self._stopping = False
            self._aborted = False

            self._receiver_task = asyncio.create_task(
                self._receive_events()
            )

        except Exception:
            await self._cleanup()
            raise

    async def send_audio(self, audio: bytes) -> None:
        if not self._started or self._websocket is None:
            raise RuntimeError("STT session has not been started.")

        if not audio:
            raise ValueError("Audio chunk cannot be empty.")

        await self._send_json(
            {
                "type": "input_audio_buffer.append",
                "audio": base64.b64encode(audio).decode("utf-8"),
            }
        )

        await self._send_json(
            {
                "type": "input_audio_buffer.commit",
            }
        )

    async def stop_session(self) -> None:
        if not self._started or self._websocket is None:
            return

        self._stopping = True

        try:
            await self._send_json(
                {
                    "type": "input_audio_buffer.done",
                }
            )

            if self._receiver_task is not None:
                await self._receiver_task

        finally:
            await self._cleanup()


    async def abort_session(self) -> None:
        """
        Abort the NVIDIA transcription session safely.

        Purpose:
            Stop transcription immediately without waiting for a final
            transcript, including when multiple abort requests overlap.

        Responsibilities:
            - Serialize concurrent abort requests.
            - Stop accepting new audio.
            - Cancel and await the receiver task.
            - Release WebSocket and HTTP resources.
            - Discard pending transcription events.
            - Publish one end-of-stream marker.

        Non-responsibilities:
            - Managing Gateway WebSocket connections.
            - Updating VoicePipeline turn state.
            - Cancelling Agent or TTS execution.
        """
        async with self._abort_lock:
            # A completed abort must not run cleanup again.
            if self._aborted:
                return

            # Mark the session as aborted before releasing resources.
            self._aborted = True
            self._stopping = True
            self._started = False

            try:
                await self._cleanup()

            finally:
                # Remove transcripts and termination markers that may
                # have been queued before the abort completed.
                while True:
                    try:
                        self._events.get_nowait()
                    except asyncio.QueueEmpty:
                        break

                # Publish one termination marker for the event consumer.
                self._events.put_nowait(None)


    def events(self) -> AsyncIterator[TranscriptEvent]:
        return self._event_stream()

    async def _event_stream(
        self,
    ) -> AsyncIterator[TranscriptEvent]:
        while True:
            event = await self._events.get()

            if event is None:
                break

            yield event

    async def _create_session(self) -> dict[str, Any]:
        if self._http is None:
            raise RuntimeError("HTTP client is not initialized.")

        response = await self._http.post(
            f"http://{self.server}/v1/realtime/transcription_sessions"
        )

        response.raise_for_status()

        session = response.json()

        if not isinstance(session, dict):
            raise RuntimeError(
                "NVIDIA NIM returned an invalid session response."
            )

        return session

    async def _connect(self) -> ClientConnection:
        return await websockets.connect(
            f"ws://{self.server}"
            "/v1/realtime"
            "?intent=transcription"
        )

    async def _initialize_session(
        self,
        session: dict[str, Any],
    ) -> None:
        if self._websocket is None:
            raise RuntimeError("NVIDIA ASR WebSocket is not connected.")

        response = await self._websocket.recv()
        event = self._decode_event(response)

        if event.get("type") != "conversation.created":
            raise RuntimeError(
                "Unexpected first NIM event: "
                f"{event.get('type')!r}"
            )

        session["input_audio_format"] = "pcm16"

        input_audio_transcription = dict(
            session.get("input_audio_transcription", {})
        )

        input_audio_transcription["language"] = self.language_code
        session["input_audio_transcription"] = (
            input_audio_transcription
        )

        session["input_audio_params"] = {
            "sample_rate_hz": CANONICAL_AUDIO_FORMAT.sample_rate,
            "num_channels": CANONICAL_AUDIO_FORMAT.channels,
        }

        recognition_config = dict(
            session.get("recognition_config", {})
        )

        recognition_config[
            "enable_automatic_punctuation"
        ] = self.automatic_punctuation

        session["recognition_config"] = recognition_config

        await self._send_json(
            {
                "type": "transcription_session.update",
                "session": session,
            }
        )

        response = await self._websocket.recv()
        event = self._decode_event(response)

        if event.get("type") != "transcription_session.updated":
            raise RuntimeError(
                "NVIDIA NIM did not confirm session update: "
                f"{event.get('type')!r}"
            )

    async def _receive_events(self) -> None:
        websocket = self._websocket

        if websocket is None:
            return

        try:
            async for raw_message in websocket:
                # Stop processing provider events after an explicit abort.
                if self._aborted:
                    break
                event = self._decode_event(raw_message)
                event_type = event.get("type", "")

                if (
                    event_type
                    == "conversation.item.input_audio_transcription.delta"
                ):
                    text = event.get("delta", "")

                    if text:
                        await self._events.put(
                            TranscriptEvent(
                                type="partial",
                                text=text,
                            )
                        )

                elif (
                    event_type
                    == "conversation.item.input_audio_transcription.completed"
                ):
                    text = event.get("transcript", "")

                    if not text:
                        continue

                    is_last_result = event.get(
                        "is_last_result",
                        False,
                    )

                    await self._events.put(
                        TranscriptEvent(
                            type=(
                                "final"
                                if is_last_result
                                else "partial"
                            ),
                            text=text,
                        )
                    )

                    if is_last_result:
                        break

                elif event_type.lower().startswith("error"):
                    error = event.get("error", {})

                    raise RuntimeError(
                        error.get(
                            "message",
                            "Unknown NVIDIA ASR error.",
                        )
                    )

        except websockets.ConnectionClosed as exc:
            if not self._stopping:
                raise RuntimeError(
                    "NVIDIA ASR WebSocket closed unexpectedly."
                ) from exc

        finally:
            if not self._aborted:
                await self._events.put(None)

    async def _send_json(self, message: dict[str, Any]) -> None:
        if self._websocket is None:
            raise RuntimeError(
                "NVIDIA ASR WebSocket is not connected."
            )

        await self._websocket.send(
            json.dumps(
                message,
                separators=(",", ":"),
            )
        )

    @staticmethod
    def _decode_event(
        message: str | bytes,
    ) -> dict[str, Any]:
        if isinstance(message, bytes):
            message = message.decode("utf-8")

        try:
            event = json.loads(message)
        except json.JSONDecodeError as exc:
            raise RuntimeError(
                "Received invalid JSON from NVIDIA ASR."
            ) from exc

        if not isinstance(event, dict):
            raise RuntimeError(
                "Received invalid event from NVIDIA ASR."
            )

        return event

    async def _cleanup(self) -> None:
        """
        Release NVIDIA transport resources.

        Lifecycle flags are controlled by start_session(),
        stop_session(), and abort_session().
        """
        self._started = False

        current_task = asyncio.current_task()

        if self._receiver_task is not None:
            if self._receiver_task is not current_task:
                if not self._receiver_task.done():
                    self._receiver_task.cancel()

                try:
                    await self._receiver_task
                except asyncio.CancelledError:
                    pass

            self._receiver_task = None

        if self._websocket is not None:
            try:
                await self._websocket.close()
            finally:
                self._websocket = None

        if self._http is not None:
            await self._http.aclose()
            self._http = None