"""
OpenAI Realtime Speech-to-Speech Adapter

Implements a direct speech-to-speech pipeline using the OpenAI Realtime API.
Audio goes in, audio comes out with minimal latency (~300ms glass-to-glass).

Architecture:
  DVGateway audio in (16kHz PCM)
      ↓ upsample to 24kHz
  OpenAI Realtime API (WebSocket)
      ↓ gpt-4o-realtime processes speech end-to-end
  DVGateway audio out (24kHz → 16kHz PCM)

Features:
- Sub-300ms end-to-end latency (vs ~500ms for cascaded pipeline)
- Native voice activity detection (server-side VAD)
- Input and output transcription

Model Reference (2026-03):
  gpt-4o-realtime-preview      — Latest stable realtime model
  gpt-4o-mini-realtime-preview — Cost-efficient realtime (audio 1.5)

API Reference: https://platform.openai.com/docs/api-reference/realtime
"""

from __future__ import annotations

import abc
import asyncio
import base64
import json
import time
from dataclasses import dataclass, field
from typing import Any, AsyncIterable, Callable, Literal
from urllib.parse import quote

import websockets
import websockets.client

from dvgateway.audio.codec import float32_to_slin16, resample, slin16_to_float32
from dvgateway.types import AudioChunk, TranscriptResult

REALTIME_API_URL = "wss://api.openai.com/v1/realtime"
OPENAI_SAMPLE_RATE = 24000
DV_SAMPLE_RATE = 16000
INPUT_CHUNK_MS = 20
INPUT_SAMPLES_24K = round(OPENAI_SAMPLE_RATE * INPUT_CHUNK_MS / 1000)

OpenAIRealtimeVoice = Literal[
    "alloy", "echo", "fable", "onyx", "nova", "shimmer",
    "ash", "ballad", "coral", "sage", "verse",
]


@dataclass
class OpenAIRealtimeTurnDetectionOptions:
    mode: Literal["server_vad", "none"] = "server_vad"
    threshold: float = 0.5
    silence_duration_ms: int = 200
    prefix_padding_ms: int = 300


class RealtimeSpeechAdapter(abc.ABC):
    """Combined speech-to-speech adapter interface."""

    @abc.abstractmethod
    async def start_session(self, linked_id: str, audio_in: AsyncIterable[AudioChunk]) -> None: ...

    @abc.abstractmethod
    def on_audio_output(self, handler: Callable[[bytes, str], None]) -> None: ...

    @abc.abstractmethod
    def on_transcript(self, handler: Callable[[TranscriptResult], None]) -> None: ...

    @abc.abstractmethod
    def on_error(self, handler: Callable[[Exception, str | None], None]) -> None: ...

    @abc.abstractmethod
    async def stop(self, linked_id: str | None = None) -> None: ...


class OpenAIRealtimeAdapter(RealtimeSpeechAdapter):
    """OpenAI Realtime speech-to-speech adapter."""

    def __init__(
        self,
        api_key: str = "",
        model: str = "gpt-4o-realtime-preview",
        voice: OpenAIRealtimeVoice = "alloy",
        instructions: str = "You are a helpful voice assistant. Keep answers concise and conversational.",
        turn_detection: OpenAIRealtimeTurnDetectionOptions | None = None,
        input_transcription: bool = True,
        temperature: float = 0.8,
        max_response_tokens: int | str = "inf",
    ) -> None:
        self._api_key = api_key
        self._model = model
        self._voice = voice
        self._instructions = instructions
        self._turn_detection = turn_detection or OpenAIRealtimeTurnDetectionOptions()
        self._input_transcription = input_transcription
        self._temperature = temperature
        self._max_response_tokens = max_response_tokens

        self._audio_output_handler: Callable[[bytes, str], None] | None = None
        self._transcript_handler: Callable[[TranscriptResult], None] | None = None
        self._error_handler: Callable[[Exception, str | None], None] | None = None
        self._sessions: dict[str, websockets.client.WebSocketClientProtocol] = {}  # type: ignore[type-arg]
        self._stopped = False

    def on_audio_output(self, handler: Callable[[bytes, str], None]) -> None:
        self._audio_output_handler = handler

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

    def on_error(self, handler: Callable[[Exception, str | None], None]) -> None:
        self._error_handler = handler

    async def start_session(self, linked_id: str, audio_in: AsyncIterable[AudioChunk]) -> None:
        self._stopped = False
        ws = await self._connect_realtime(linked_id)
        self._sessions[linked_id] = ws

        # Send session config
        self._send_event(ws, {
            "type": "session.update",
            "session": {
                "modalities": ["audio", "text"],
                "voice": self._voice,
                "instructions": self._instructions,
                "input_audio_format": "pcm16",
                "output_audio_format": "pcm16",
                "temperature": self._temperature,
                "max_response_output_tokens": self._max_response_tokens,
                **({"input_audio_transcription": {"model": "whisper-1"}} if self._input_transcription else {}),
                "turn_detection": self._build_turn_detection(),
            },
        })

        # Pipe audio in background
        asyncio.ensure_future(self._pipe_audio_in(linked_id, ws, audio_in))

        # Read events
        try:
            async for message in ws:
                if self._stopped or linked_id not in self._sessions:
                    break
                try:
                    event = json.loads(message if isinstance(message, str) else message.decode("utf-8"))
                    self._handle_server_event(event, linked_id)
                except (json.JSONDecodeError, UnicodeDecodeError):
                    pass
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            self._sessions.pop(linked_id, None)

    async def stop(self, linked_id: str | None = None) -> None:
        if linked_id:
            ws = self._sessions.pop(linked_id, None)
            if ws:
                await ws.close(1000, "session ended")
        else:
            self._stopped = True
            for lid, ws in list(self._sessions.items()):
                await ws.close(1000, "adapter stopped")
            self._sessions.clear()

    async def _connect_realtime(self, linked_id: str) -> websockets.client.WebSocketClientProtocol:  # type: ignore[type-arg]
        url = f"{REALTIME_API_URL}?model={quote(self._model)}"
        ws = await websockets.client.connect(
            url,
            additional_headers={
                "Authorization": f"Bearer {self._api_key}",
                "OpenAI-Beta": "realtime=v1",
            },
        )
        return ws

    async def _pipe_audio_in(
        self,
        linked_id: str,
        ws: websockets.client.WebSocketClientProtocol,  # type: ignore[type-arg]
        audio_in: AsyncIterable[AudioChunk],
    ) -> None:
        async for chunk in audio_in:
            if self._stopped or linked_id not in self._sessions:
                break

            # Upsample from 16kHz to 24kHz
            samples_24k = resample(chunk.samples, DV_SAMPLE_RATE, OPENAI_SAMPLE_RATE)

            # Split into 20ms frames
            offset = 0
            while offset < len(samples_24k):
                end = min(offset + INPUT_SAMPLES_24K, len(samples_24k))
                frame = samples_24k[offset:end]
                offset = end

                pcm16 = float32_to_slin16(frame)
                b64 = base64.b64encode(pcm16).decode("ascii")

                self._send_event(ws, {
                    "type": "input_audio_buffer.append",
                    "audio": b64,
                })

        # Signal end of input if not using server VAD
        if self._turn_detection.mode == "none":
            self._send_event(ws, {"type": "input_audio_buffer.commit"})
            self._send_event(ws, {"type": "response.create"})

    def _handle_server_event(self, event: dict[str, Any], linked_id: str) -> None:
        event_type = event.get("type", "")

        if event_type == "response.audio.delta":
            delta = event.get("delta")
            if not delta:
                return
            pcm_24k = base64.b64decode(delta)
            samples_24k = slin16_to_float32(pcm_24k)
            samples_16k = resample(samples_24k, OPENAI_SAMPLE_RATE, DV_SAMPLE_RATE)
            pcm_16k = float32_to_slin16(samples_16k)
            if self._audio_output_handler:
                self._audio_output_handler(pcm_16k, linked_id)

        elif event_type == "conversation.item.input_audio_transcription.completed":
            transcript = event.get("transcript", "")
            if transcript and self._transcript_handler:
                self._transcript_handler(TranscriptResult(
                    linked_id=linked_id,
                    text=transcript,
                    is_final=True,
                    speaker="customer",
                    timestamp_ms=time.time() * 1000,
                ))

        elif event_type == "response.audio_transcript.done":
            transcript = event.get("transcript", "")
            if not transcript:
                part = event.get("part", {})
                transcript = part.get("transcript", "")
            if transcript and self._transcript_handler:
                self._transcript_handler(TranscriptResult(
                    linked_id=linked_id,
                    text=transcript,
                    is_final=True,
                    speaker="agent",
                    timestamp_ms=time.time() * 1000,
                ))

        elif event_type == "error":
            err_info = event.get("error", {})
            msg = err_info.get("message", "Unknown realtime API error")
            code = err_info.get("code", "")
            code_str = f" [{code}]" if code else ""
            if self._error_handler:
                self._error_handler(RuntimeError(f"OpenAI Realtime{code_str}: {msg}"), linked_id)

    def _build_turn_detection(self) -> dict[str, object]:
        td = self._turn_detection
        if td.mode == "none":
            return {"type": "none"}
        return {
            "type": "server_vad",
            "threshold": td.threshold,
            "silence_duration_ms": td.silence_duration_ms,
            "prefix_padding_ms": td.prefix_padding_ms,
            "create_response": True,
        }

    def _send_event(self, ws: websockets.client.WebSocketClientProtocol, event: dict[str, object]) -> None:  # type: ignore[type-arg]
        asyncio.ensure_future(ws.send(json.dumps(event)))
