"""Django Channels middleware for WebSocket trace context and tenant_id injection."""

import logging
from urllib.parse import parse_qs
from opentelemetry import trace

logger = logging.getLogger(__name__)

try:
    from varicon_observability.celery_helpers import tenant_context
except ImportError:
    import contextvars

    tenant_context = contextvars.ContextVar("tenant_id", default="unknown-tenant")


class WebSocketTraceMiddleware:
    """
    Django Channels ASGI middleware that:
    1. Extracts tenant_id from WebSocket query params and sets it in context
    2. Creates a span covering the entire WebSocket connection lifecycle
    3. Ensures all logs within the connection have trace_id, span_id, and tenant_id
    """

    def __init__(self, inner):
        self.inner = inner
        self.tracer = trace.get_tracer("varicon.websocket", "1.0.0")

    async def __call__(self, scope, receive, send):
        query_string = scope.get("query_string", b"").decode("utf-8")
        query_params = parse_qs(query_string)

        # Set tenant_id in context for trace correlation
        tenant_id = query_params.get("tenant", [None])[0]
        if tenant_id:
            tenant_context.set(str(tenant_id))

        # Create a span covering the entire WebSocket connection lifecycle
        path = scope.get("path", "/")
        span_name = f"WEBSOCKET {path}"

        with self.tracer.start_as_current_span(
            span_name,
            kind=trace.SpanKind.SERVER,
            attributes={
                "websocket.path": path,
                "websocket.type": scope.get("type", "websocket"),
                "tenant_id": str(tenant_id) if tenant_id else "unknown",
            },
        ) as span:
            ctx = span.get_span_context()
            if ctx.is_valid:
                trace_id = format(ctx.trace_id, "032x")
                span_id = format(ctx.span_id, "016x")
                logger.debug(
                    f"[WS_TRACE] {span_name} trace_id: {trace_id}, "
                    f"span_id: {span_id}, tenant_id: {tenant_id}"
                )

            return await self.inner(scope, receive, send)
