"""Realtime STT client for Riva/NIM WebSocket transcription.

The NIM HTTP docs expose session creation, but the live audio path is a
WebSocket at /v1/realtime?intent=transcription. Audio is sent as base64 PCM16
chunks and transcript events arrive as interim deltas and completed segments.
"""

from __future__ import annotations

import asyncio
import base64
import json
import os
import queue
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from copy import deepcopy
from dataclasses import dataclass
from typing import Callable

from .stt import STTEngine

TextCallback = Callable[[str, bool], None]
StatusCallback = Callable[[str], None]
ErrorCallback = Callable[[Exception], None]


class StreamingSTTError(RuntimeError):
    pass


@dataclass
class RealtimeURLs:
    session_url: str
    websocket_url: str


def realtime_urls(base_url: str) -> RealtimeURLs:
    """Build the HTTP session URL and WebSocket URL from a user base URL."""
    raw = (base_url or "").strip().rstrip("/")
    if not raw:
        raise StreamingSTTError("Missing realtime STT URL.")
    if "://" not in raw:
        raw = "http://" + raw

    parsed = urllib.parse.urlparse(raw)
    path = parsed.path.rstrip("/")
    if path.endswith("/realtime"):
        realtime_path = path
    elif path.endswith("/v1"):
        realtime_path = path + "/realtime"
    else:
        realtime_path = path + "/v1/realtime"

    http_scheme = "https" if parsed.scheme == "https" else "http"
    ws_scheme = "wss" if http_scheme == "https" else "ws"
    netloc = parsed.netloc
    session = urllib.parse.urlunparse((http_scheme, netloc, realtime_path + "/transcription_sessions", "", "", ""))
    ws = urllib.parse.urlunparse((ws_scheme, netloc, realtime_path, "", "intent=transcription", ""))
    return RealtimeURLs(session, ws)


class RivaRealtimeStreamer:
    """Owns one live microphone -> Riva WebSocket transcription session."""

    def __init__(
        self,
        engine: STTEngine,
        *,
        device: str = "",
        language: str = "",
        sample_rate: int = 16000,
        chunk_frames: int = 1600,
        on_text: TextCallback | None = None,
        on_status: StatusCallback | None = None,
        on_error: ErrorCallback | None = None,
    ):
        self.engine = engine
        self.device = device
        self.language = language.strip()
        self.sample_rate = sample_rate
        self.chunk_frames = chunk_frames
        self.on_text = on_text
        self.on_status = on_status
        self.on_error = on_error
        self._stop = threading.Event()
        self._ready = threading.Event()
        self._thread: threading.Thread | None = None
        self._error: Exception | None = None

    def start(self) -> None:
        if self._thread is not None:
            return
        self._thread = threading.Thread(target=self._run_thread, daemon=True)
        self._thread.start()
        self._ready.wait(timeout=5.0)
        if self._error:
            raise self._error

    def stop(self) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=1.0)

    def _run_thread(self) -> None:
        try:
            asyncio.run(self._run())
        except Exception as exc:  # noqa: BLE001 - surfaced to daemon/UI
            self._error = exc
            self._ready.set()
            if self.on_error:
                self.on_error(exc)

    async def _run(self) -> None:
        try:
            import sounddevice as sd
            import websockets
        except ModuleNotFoundError as exc:
            raise StreamingSTTError(
                "Realtime streaming needs the Python packages 'websockets' and 'sounddevice'."
            ) from exc

        urls = realtime_urls(self.engine.url)
        session = self._create_session(urls.session_url)
        headers = self._auth_headers()
        websocket = await self._connect_websocket(websockets, urls.websocket_url, headers)
        audio_q: queue.Queue[bytes] = queue.Queue(maxsize=60)

        def audio_cb(indata, _frames, _time_info, status):
            if status and self.on_status:
                self.on_status(str(status))
            try:
                audio_q.put_nowait(bytes(indata))
            except queue.Full:
                pass

        try:
            await self._initialize_websocket(websocket, session)
            self._ready.set()
            if self.on_status:
                self.on_status("Streaming")

            with sd.RawInputStream(
                samplerate=self.sample_rate,
                channels=1,
                dtype="int16",
                blocksize=self.chunk_frames,
                device=self._resolve_device(sd),
                callback=audio_cb,
            ):
                send_task = asyncio.create_task(self._send_audio(websocket, audio_q))
                recv_task = asyncio.create_task(self._receive_text(websocket))
                await send_task
                try:
                    await asyncio.wait_for(recv_task, timeout=8.0)
                except asyncio.TimeoutError:
                    recv_task.cancel()
        finally:
            await websocket.close()

    def _auth_headers(self) -> dict[str, str]:
        if not self.engine.api_key_env:
            return {}
        key = os.environ.get(self.engine.api_key_env)
        return {"Authorization": f"Bearer {key}"} if key else {}

    def _create_session(self, url: str) -> dict:
        body = b"{}"
        headers = {"Content-Type": "application/json", **self._auth_headers()}
        req = urllib.request.Request(url, data=body, headers=headers, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=10) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", "replace")[:300]
            raise StreamingSTTError(f"HTTP {exc.code} from {url}: {detail}") from exc
        except (urllib.error.URLError, json.JSONDecodeError, OSError) as exc:
            raise StreamingSTTError(f"Cannot initialize realtime STT session at {url}: {exc}") from exc

    async def _connect_websocket(self, websockets, url: str, headers: dict[str, str]):
        try:
            return await websockets.connect(url, additional_headers=headers or None)
        except TypeError:
            return await websockets.connect(url, extra_headers=headers or None)

    async def _initialize_websocket(self, websocket, session: dict) -> None:
        first = json.loads(await asyncio.wait_for(websocket.recv(), timeout=5.0))
        if first.get("type") != "conversation.created":
            raise StreamingSTTError(f"Unexpected realtime greeting: {first}")

        updated = deepcopy(session)
        updated["input_audio_format"] = "pcm16"
        updated.setdefault("input_audio_params", {})
        updated["input_audio_params"]["sample_rate_hz"] = self.sample_rate
        updated["input_audio_params"]["num_channels"] = 1
        updated.setdefault("recognition_config", {})
        updated["recognition_config"]["max_alternatives"] = 1
        if self.language:
            updated.setdefault("input_audio_transcription", {})
            updated["input_audio_transcription"]["language"] = self.language
        if self.engine.model:
            updated.setdefault("input_audio_transcription", {})
            updated["input_audio_transcription"]["model"] = self.engine.model

        await websocket.send(json.dumps({"type": "transcription_session.update", "session": updated}))
        response = json.loads(await asyncio.wait_for(websocket.recv(), timeout=5.0))
        if response.get("type") != "transcription_session.updated":
            raise StreamingSTTError(f"Realtime session update failed: {response}")

    async def _send_audio(self, websocket, audio_q: queue.Queue[bytes]) -> None:
        while not self._stop.is_set():
            try:
                chunk = await asyncio.to_thread(audio_q.get, True, 0.2)
            except queue.Empty:
                continue
            await self._send_chunk(websocket, chunk)

        drain_until = time.monotonic() + 0.3
        while time.monotonic() < drain_until:
            try:
                chunk = audio_q.get_nowait()
            except queue.Empty:
                break
            await self._send_chunk(websocket, chunk)
        await websocket.send(json.dumps({"type": "input_audio_buffer.done"}))

    async def _send_chunk(self, websocket, chunk: bytes) -> None:
        audio = base64.b64encode(chunk).decode("ascii")
        await websocket.send(json.dumps({"type": "input_audio_buffer.append", "audio": audio}))
        await websocket.send(json.dumps({"type": "input_audio_buffer.commit"}))

    async def _receive_text(self, websocket) -> None:
        while True:
            try:
                raw = await asyncio.wait_for(websocket.recv(), timeout=1.0)
            except asyncio.TimeoutError:
                continue
            event = json.loads(raw)
            kind = event.get("type", "")
            if kind == "conversation.item.input_audio_transcription.delta":
                text = (event.get("delta") or "").strip()
                if text and self.on_text:
                    self.on_text(text, False)
            elif kind == "conversation.item.input_audio_transcription.completed":
                text = (event.get("transcript") or "").strip()
                if text and self.on_text:
                    self.on_text(text, True)
                if event.get("is_last_result"):
                    return
            elif kind == "conversation.item.input_audio_transcription.failed":
                raise StreamingSTTError(event.get("error", {}).get("message", "Transcription failed."))
            elif kind == "error":
                err = event.get("error", {})
                raise StreamingSTTError(err.get("message") or str(err) or "Realtime STT error.")

    def _resolve_device(self, sd):
        if not self.device:
            return None
        try:
            for i, d in enumerate(sd.query_devices()):
                if d["max_input_channels"] > 0 and self.device in d["name"]:
                    return i
        except Exception:  # noqa: BLE001
            pass
        return None