from __future__ import annotations

import json
from collections.abc import AsyncIterator, Awaitable, Callable
from types import TracebackType
from typing import Any, Protocol

from .models import ASRControl, ASRStart, ASRTranscriptEvent


class ASRSocket(Protocol):
    async def send(self, message: str | bytes) -> Any: ...
    async def recv(self) -> str | bytes: ...
    async def close(self) -> Any: ...


class ASRConnector(Protocol):
    def __call__(self, url: str, **kwargs: Any) -> Any: ...


class StreamingASRSession:
    """Thin binary-audio/JSON-control contract for the ASR service on port 8101."""

    def __init__(
        self,
        *,
        url: str,
        start: ASRStart,
        connector: ASRConnector | None = None,
        open_timeout_s: float = 0.75,
        max_size: int = 2**20,
    ) -> None:
        self.url = url
        self.start = start
        self.connector = connector
        self.open_timeout_s = open_timeout_s
        self.max_size = max_size
        self._connection_context: Any = None
        self._socket: ASRSocket | None = None
        self._closed = False

    async def __aenter__(self) -> "StreamingASRSession":
        connector = self.connector
        if connector is None:
            from websockets.asyncio.client import connect

            connector = connect
        self._connection_context = connector(
            self.url,
            open_timeout=self.open_timeout_s,
            max_size=self.max_size,
            ping_interval=20,
            ping_timeout=10,
        )
        self._socket = await self._connection_context.__aenter__()
        await self._socket.send(self.start.model_dump_json(exclude_none=True))
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        self._closed = True
        if self._connection_context is not None:
            await self._connection_context.__aexit__(exc_type, exc, traceback)

    async def send_audio(self, pcm: bytes) -> None:
        if not pcm:
            return
        socket = self._require_socket()
        await socket.send(pcm)

    async def commit(self) -> None:
        await self._require_socket().send(ASRControl(type="input.commit").model_dump_json())

    async def cancel(self) -> None:
        if self._socket is None or self._closed:
            return
        await self._socket.send(ASRControl(type="input.cancel").model_dump_json())

    async def events(self) -> AsyncIterator[ASRTranscriptEvent]:
        socket = self._require_socket()
        while not self._closed:
            raw = await socket.recv()
            if isinstance(raw, bytes):
                yield ASRTranscriptEvent(type="error", detail="ASR returned an unexpected binary frame")
                continue
            event = ASRTranscriptEvent.model_validate_json(raw)
            yield event
            if event.type in {"input.transcript.final", "error"}:
                return

    def _require_socket(self) -> ASRSocket:
        if self._socket is None or self._closed:
            raise RuntimeError("ASR session is not open")
        return self._socket


class StreamingASRBridge:
    def __init__(
        self,
        url: str,
        *,
        connector: ASRConnector | None = None,
        open_timeout_s: float = 0.75,
    ) -> None:
        self.url = url
        self.connector = connector
        self.open_timeout_s = open_timeout_s

    def session(self, start: ASRStart) -> StreamingASRSession:
        return StreamingASRSession(
            url=self.url,
            start=start,
            connector=self.connector,
            open_timeout_s=self.open_timeout_s,
        )