"""
Audio Stream

Subscribes to real-time call audio from the DVGateway.
Returns an AsyncIterable[AudioChunk] for easy integration with AI STT adapters.

WebSocket endpoint: /api/v1/ws/stream?linkedid={id}&dir={both|in|out}
Protocol: Binary frames (slin16, 16kHz, 16-bit, little-endian, 640 bytes/frame)
"""

from __future__ import annotations

import asyncio
import json
import uuid
from typing import AsyncIterator

from dvgateway.audio.codec import slin16_to_float32
from dvgateway.transport.ws_pool import WsPool, WsSubscription
from dvgateway.types import AudioChunk, Logger, StreamDir

SAMPLE_RATE = 16000
BYTES_PER_SAMPLE = 2
FRAME_BYTES = 640  # 320 samples × 2 bytes = 20ms at 16kHz


class AudioStream:
    """Async iterable of AudioChunk from a call's real-time audio."""

    def __init__(self, ws_pool: WsPool, linked_id: str, logger: Logger, dir: StreamDir = "both") -> None:
        self._ws_pool = ws_pool
        self._linked_id = linked_id
        self._logger = logger
        self._dir = dir
        self._sub_id = f"audio-{linked_id}-{uuid.uuid4()}"
        self._queue: asyncio.Queue[AudioChunk | None] = asyncio.Queue()
        self._ended = False
        self._started = False

    async def _start(self) -> None:
        if self._started:
            return
        self._started = True

        def on_message(data: bytes) -> None:
            if self._ended:
                return
            import time
            timestamp_ms = time.time() * 1000
            offset = 0
            while offset + FRAME_BYTES <= len(data):
                frame_bytes = data[offset : offset + FRAME_BYTES]
                samples = slin16_to_float32(frame_bytes)
                duration_ms = (len(samples) / SAMPLE_RATE) * 1000
                chunk = AudioChunk(
                    linked_id=self._linked_id,
                    dir="mixed" if self._dir == "both" else self._dir,
                    samples=samples,
                    sample_rate=SAMPLE_RATE,
                    duration_ms=duration_ms,
                    timestamp_ms=timestamp_ms + offset / (SAMPLE_RATE * BYTES_PER_SAMPLE / 1000),
                )
                self._queue.put_nowait(chunk)
                offset += FRAME_BYTES

            # Handle partial frame
            if offset < len(data):
                partial = data[offset:]
                samples = slin16_to_float32(partial)
                duration_ms = (len(samples) / SAMPLE_RATE) * 1000
                chunk = AudioChunk(
                    linked_id=self._linked_id,
                    dir="mixed" if self._dir == "both" else self._dir,
                    samples=samples,
                    sample_rate=SAMPLE_RATE,
                    duration_ms=duration_ms,
                    timestamp_ms=timestamp_ms,
                )
                self._queue.put_nowait(chunk)

        def on_close() -> None:
            self._logger.debug({"linked_id": self._linked_id}, "Audio stream WebSocket closed")
            self._queue.put_nowait(None)

        try:
            await self._ws_pool.subscribe(WsSubscription(
                id=self._sub_id,
                path="/api/v1/ws/stream",
                params={"linkedid": self._linked_id, "dir": self._dir},
                on_message=on_message,
                on_close=on_close,
            ))
        except Exception as e:
            self._logger.error({"linked_id": self._linked_id, "err": str(e)}, "Audio stream subscription failed")
            self._queue.put_nowait(None)

    def unsubscribe(self) -> None:
        self._ended = True
        self._ws_pool.unsubscribe(self._sub_id)
        self._queue.put_nowait(None)

    async def send_thinking_start(self) -> None:
        """Send thinking:start signal to gateway (triggers comfort noise injection)."""
        await self._ws_pool.send_text(self._sub_id, json.dumps({"type": "thinking:start"}))

    async def send_thinking_stop(self) -> None:
        """Send thinking:stop signal to gateway (stops comfort noise injection)."""
        await self._ws_pool.send_text(self._sub_id, json.dumps({"type": "thinking:stop"}))

    def __aiter__(self) -> AsyncIterator[AudioChunk]:
        return self._iterator()

    async def _iterator(self) -> AsyncIterator[AudioChunk]:
        await self._start()
        while True:
            if self._ended:
                return
            entry = await self._queue.get()
            if entry is None:
                return
            yield entry


def create_audio_stream(
    ws_pool: WsPool,
    linked_id: str,
    logger: Logger,
    dir: StreamDir = "both",
) -> AudioStream:
    """Create an async iterable of audio chunks for a given call session."""
    return AudioStream(ws_pool, linked_id, logger, dir)
