"""Drive an ASGI websocket endpoint on the caller's own event loop.

Starlette's ``TestClient`` runs the app in a second thread with its own loop, which
would put the bus and the websockets on different loops - so ``bus.emit`` would take
its thread-safe path and a following ``drain()`` could return before the event was even
queued. In the real app there is only ever one loop, and this harness keeps the tests
that way.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
from collections.abc import AsyncIterator
from typing import Any


class WebSocketSession:
    def __init__(self) -> None:
        self.to_app: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self.from_app: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

    async def receive(self) -> dict[str, Any]:
        return await self.to_app.get()

    async def send(self, message: dict[str, Any]) -> None:
        await self.from_app.put(message)

    async def next_json(self, timeout: float = 2.0) -> dict[str, Any]:
        """The next ``websocket.send`` frame, decoded."""
        while True:
            message = await asyncio.wait_for(self.from_app.get(), timeout)
            if message["type"] == "websocket.send":
                text: str = message["text"]
                return json.loads(text)


@contextlib.asynccontextmanager
async def websocket_connect(app: Any, path: str) -> AsyncIterator[WebSocketSession]:
    session = WebSocketSession()
    scope = {
        "type": "websocket",
        "asgi": {"version": "3.0", "spec_version": "2.3"},
        "http_version": "1.1",
        "scheme": "ws",
        "path": path,
        "raw_path": path.encode(),
        "query_string": b"",
        "root_path": "",
        "headers": [(b"host", b"testserver")],
        "client": ("testclient", 50000),
        "server": ("testserver", 80),
        "subprotocols": [],
        "state": {},
    }
    await session.to_app.put({"type": "websocket.connect"})
    task = asyncio.create_task(app(scope, session.receive, session.send))

    accepted = await asyncio.wait_for(session.from_app.get(), 2.0)
    assert accepted["type"] == "websocket.accept", accepted
    try:
        yield session
    finally:
        await session.to_app.put({"type": "websocket.disconnect", "code": 1000})
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError, Exception):
            await task