"""Stream — WebSocket connection for real-time STT."""

from __future__ import annotations

import asyncio
import json
import logging
import time
from collections.abc import AsyncIterator, Callable

import websockets

from .audio import BYTES_PER_SECOND, load_wav, silence, stream_realtime
from .exceptions import ConnectionError, SessionExpiredError, STTError
from .types import SessionInfo, Transcription

logger = logging.getLogger("dtelecom_stt")

# How long to wait for new transcriptions after audio ends (transcribe_file)
_FILE_DRAIN_TIMEOUT = 5.0
# Trailing silence to flush VAD pipeline
_TRAILING_SILENCE_SECONDS = 2.0


class Stream:
    """WebSocket stream for sending audio and receiving transcriptions.

    Do not instantiate directly — use ``STTClient.session()`` instead.
    """

    def __init__(
        self,
        ws_url: str,
        session_info: SessionInfo,
        client,  # STTClient (avoid circular import)
        language: str,
        auto_extend: bool,
    ):
        self._ws_url = ws_url
        self._info = session_info
        self._client = client
        self._language = language
        self._auto_extend = auto_extend

        self._ws = None
        self._recv_task: asyncio.Task | None = None
        self._queue: asyncio.Queue[Transcription | None] = asyncio.Queue()
        self._callbacks: list[Callable[[Transcription], None]] = []
        self._closed = False
        self._extending = False

    # ── Connection lifecycle ──────────────────────────────────────────

    async def _connect(self):
        """Connect WebSocket, send config, wait for ready."""
        try:
            self._ws = await websockets.connect(self._ws_url)
        except Exception as e:
            raise ConnectionError(f"WebSocket connection failed: {e}") from e

        config = {
            "type": "config",
            "language": self._language,
            "session_key": self._info.session_key,
        }
        await self._ws.send(json.dumps(config))

        try:
            raw = await asyncio.wait_for(self._ws.recv(), timeout=30)
        except asyncio.TimeoutError:
            await self._ws.close()
            raise ConnectionError("Timeout waiting for ready message")

        msg = json.loads(raw)
        if msg.get("type") == "error":
            await self._ws.close()
            raise STTError(f"Server error: {msg.get('message', msg)}")
        if msg.get("type") != "ready":
            await self._ws.close()
            raise STTError(f"Expected ready message, got: {msg}")

        remaining = msg.get("remaining_seconds", self._info.remaining_seconds)
        logger.info("Stream ready, remaining=%.0fs", remaining)

        # Start background receive loop
        self._recv_task = asyncio.create_task(self._recv_loop())

    async def close(self):
        """Close the WebSocket connection gracefully."""
        if self._closed:
            return
        self._closed = True

        # Signal end to async iterator
        self._queue.put_nowait(None)

        if self._recv_task and not self._recv_task.done():
            self._recv_task.cancel()
            try:
                await self._recv_task
            except asyncio.CancelledError:
                pass

        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass

    async def __aenter__(self) -> Stream:
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    # ── Sending audio ─────────────────────────────────────────────────

    async def send_audio(self, data: bytes):
        """Send raw PCM16 audio bytes (16kHz mono).

        Args:
            data: Raw PCM16 audio bytes.
        """
        if self._closed:
            raise STTError("Stream is closed")
        await self._ws.send(data)

    # ── Receiving transcriptions ──────────────────────────────────────

    def on_transcription(self, callback: Callable[[Transcription], None]):
        """Register a callback for incoming transcriptions.

        Args:
            callback: Function called with each Transcription.
        """
        self._callbacks.append(callback)

    async def transcriptions(self) -> AsyncIterator[Transcription]:
        """Async iterator yielding transcriptions as they arrive.

        Stops when the stream is closed or session expires.
        """
        while True:
            item = await self._queue.get()
            if item is None:
                return
            yield item

    async def transcribe_file(self, path: str) -> AsyncIterator[Transcription]:
        """Stream a WAV file and yield transcriptions.

        Loads the file, streams at real-time speed, sends trailing silence
        to flush the VAD pipeline, and yields results.

        Args:
            path: Path to a WAV file (16kHz, mono, PCM16).
        """
        pcm_data, _, _, _ = load_wav(path)
        audio_duration = len(pcm_data) / BYTES_PER_SECOND

        logger.info("Streaming file %s (%.1fs)", path, audio_duration)

        # Stream audio at real-time speed
        await stream_realtime(self._ws, pcm_data)

        # Send trailing silence to flush VAD
        await self._ws.send(silence(_TRAILING_SILENCE_SECONDS))
        await asyncio.sleep(_TRAILING_SILENCE_SECONDS)

        # Drain transcriptions with timeout
        last_recv = time.monotonic()
        while True:
            try:
                item = await asyncio.wait_for(
                    self._queue.get(), timeout=_FILE_DRAIN_TIMEOUT
                )
                if item is None:
                    return
                last_recv = time.monotonic()
                yield item
            except asyncio.TimeoutError:
                # No more transcriptions for _FILE_DRAIN_TIMEOUT seconds
                return

    # ── Background receive loop ───────────────────────────────────────

    async def _recv_loop(self):
        """Background task reading WebSocket messages."""
        try:
            async for raw in self._ws:
                if isinstance(raw, bytes):
                    # Binary messages from server (unexpected but ignore)
                    continue

                msg = json.loads(raw)
                msg_type = msg.get("type")

                if msg_type == "transcription":
                    t = Transcription.from_message(msg)
                    self._queue.put_nowait(t)
                    for cb in self._callbacks:
                        try:
                            cb(t)
                        except Exception:
                            logger.exception("Transcription callback error")

                elif msg_type == "session_expiring":
                    remaining = msg.get("remaining_seconds", 0)
                    logger.warning("Session expiring, %.0fs remaining", remaining)
                    if self._auto_extend and not self._extending:
                        asyncio.create_task(self._auto_extend_session())

                elif msg_type == "session_extended":
                    remaining = msg.get("remaining_seconds", 0)
                    logger.info("Session extended, %.0fs remaining", remaining)

                elif msg_type == "session_expired":
                    logger.error("Session expired")
                    self._queue.put_nowait(None)
                    raise SessionExpiredError("Session time exhausted")

                elif msg_type == "error":
                    logger.error("Server error: %s", msg.get("message", msg))

        except websockets.ConnectionClosed:
            logger.info("WebSocket closed")
        except SessionExpiredError:
            raise
        except Exception:
            logger.exception("Receive loop error")
        finally:
            # Signal end to async iterator
            if not self._closed:
                self._queue.put_nowait(None)

    async def _auto_extend_session(self):
        """Buy more time when session is about to expire."""
        self._extending = True
        try:
            result = await self._client._extend_session(self._info.session_id, 5)
            logger.info(
                "Auto-extended session: +%d min, remaining=%.0fs, $%s",
                result.get("minutes_added", 5),
                result.get("remaining_seconds", 0),
                result.get("price_usd", "?"),
            )
        except Exception:
            logger.exception("Auto-extend failed")
        finally:
            self._extending = False
