"""
ASGI middleware for tracing WebSocket connections.
"""
import logging
from typing import Any, Awaitable, Callable

from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode

logger = logging.getLogger(__name__)


class WebSocketTraceMiddleware:
    """
    ASGI middleware that creates a span for each WebSocket connection.
    Wraps the inner ASGI app (e.g. URLRouter for channels).
    """

    def __init__(self, app: Callable[..., Awaitable[None]]):
        self.app = app
        self._tracer = trace.get_tracer(__name__)

    async def __call__(
        self, scope: dict, receive: Callable, send: Callable
    ) -> None:
        if scope.get("type") != "websocket":
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")
        with self._tracer.start_as_current_span(
            f"websocket {path}",
            attributes={
                "websocket.path": path,
                "websocket.type": "server",
            },
        ) as span:
            try:
                await self.app(scope, receive, send)
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise
