"""Middleware for making it easier to do distributed tracing."""


class TracingMiddleware:
    """Middleware for propagating distributed tracing context using LangSmith.

    This middleware checks for the 'langsmith-trace' header and propagates the
    tracing context if present. It does not start new traces by default.
    It is designed to work with ASGI applications.

    Attributes:
        app: The ASGI application being wrapped.
    """

    def __init__(self, app):
        """Initialize the middleware."""
        from langsmith.run_helpers import tracing_context  # type: ignore

        self._with_headers = tracing_context
        self.app = app

    async def __call__(self, scope: dict, receive, send):
        """Handle incoming requests and propagate tracing context if applicable.

        Args:
            scope: A dict containing ASGI connection scope.
            receive: An awaitable callable for receiving ASGI events.
            send: An awaitable callable for sending ASGI events.

        If the request is HTTP and contains the 'langsmith-trace' header,
        it propagates the tracing context before calling the wrapped application.
        Otherwise, it calls the application directly without modifying the context.
        """
        if scope["type"] == "http" and "headers" in scope:
            headers = dict(scope["headers"])
            if b"langsmith-trace" in headers:
                with self._with_headers(parent=headers):
                    await self.app(scope, receive, send)
                return
        await self.app(scope, receive, send)