# Generated by smithplates HTTP codegen. Do not edit.

from __future__ import annotations

import json
from typing import Any, Self
from urllib.parse import quote

import websockets
from generated.example.server_event import ServerEvent
from pydantic import TypeAdapter


class StreamApiWebsocketClient:
    def __init__(self, base_url: str) -> None:
        self._base_url = base_url.rstrip("/")

    async def connect_stream_events(self, stream_id: str) -> StreamEventsConnection:
        websocket = await websockets.connect(f"{self._base_url}/streams/{quote(str(stream_id), safe='')}/events/ws")
        return StreamEventsConnection(websocket)


class StreamEventsConnection:
    def __init__(self, websocket: Any) -> None:
        self._websocket = websocket
        self._output_adapter: TypeAdapter[ServerEvent] = TypeAdapter(ServerEvent)

    async def receive(self) -> ServerEvent:
        raw = await self._websocket.recv()
        return self._output_adapter.validate_python(json.loads(raw))

    async def close(self) -> None:
        await self._websocket.close()

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

    async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
        await self.close()

    def __aiter__(self) -> Self:
        return self

    async def __anext__(self) -> ServerEvent:
        try:
            return await self.receive()
        except websockets.ConnectionClosed:
            raise StopAsyncIteration from None
