# #region TraceContextMiddlewareModule [C:3] [TYPE Module] [SEMANTICS fastapi, middleware, trace, context, request]
# @defgroup Core Module group.
# @BRIEF Raw ASGI middleware that seeds a trace_id for every incoming HTTP request.
#           Optionally extracts X-Trace-ID header for cross-service trace propagation.
#           Implemented as raw ASGI middleware (not BaseHTTPMiddleware) to avoid
#           contextvar isolation caused by BaseHTTPMiddleware's internal task switching
#           via anyio.create_task_group() — see [RCA TRACE-001] for details.
# @LAYER Core
# @RELATION DEPENDS_ON -> [CotLoggerModule]
# @RELATION CALLED_BY -> [AppModule]
# @INVARIANT X-Trace-ID header UUID is validated as version 4 explicitly (not via
#            uuid.UUID(hex=..., version=4) which silently mutates non-v4 UUIDs).
#            Non-UTF8 header bytes are caught via UnicodeDecodeError. See [FIX TRACE-002].
# @PRE  FastAPI app instance with ASGI middleware support.
# @POST  Every HTTP request gets a trace_id via seed_trace_id().
#        If a valid UUID4 X-Trace-ID header is present, that value is used.
#        Non-v4 UUIDs, non-UUID values, and malformed UTF-8 headers fall back to
#        seed_trace_id() (graceful degradation).
# @SIDE_EFFECT  Sets ContextVar _trace_id for the current asyncio task context.
#        Lifetime is bounded by the asyncio task — no explicit cleanup needed since
#        each request runs in its own task.
# @RATIONALE BaseHTTPMiddleware (Starlette 0.50.0) uses anyio.create_task_group() internally,
#            which creates separate asyncio tasks for dispatch and call_next. ContextVars set
#            in dispatch() are not visible to outer middlewares (like log_requests) because
#            they run in different async contexts. Converting to raw ASGI middleware ensures
#            trace_id is seeded in the root task context before any BaseHTTPMiddleware runs,
#            making it visible to ALL middleware layers. See [RCA TRACE-001].
# @REJECTED Keeping BaseHTTPMiddleware with workaround (e.g. manually passing trace_id via
#            request.state) was rejected — it would require modifying every middleware that
#            needs trace_id access. Raw ASGI middleware is the minimal fix.
# @REJECTED uuid.UUID(hex=..., version=4) validation rejected in [FIX TRACE-002] — the
#            version=N parameter silently MUTATES the UUID's version nibble instead of
#            validating it. A UUID v1 passed as X-Trace-ID would be corrupted to a different
#            UUID, breaking cross-service trace correlation. We now parse without version
#            and check parsed.version == 4 explicitly.

import uuid

from starlette.types import ASGIApp, Receive, Scope, Send

from ..cot_logger import get_trace_id, seed_trace_id, set_trace_id


# #region TraceContextMiddleware [C:2] [TYPE Class] [SEMANTICS middleware, trace, context, asgi]
# @defgroup Core Module group.
# @BRIEF Raw ASGI middleware that seeds a trace_id per request. Runs in the root task
#        context, ensuring all downstream BaseHTTPMiddleware layers see the trace_id.
class TraceContextMiddleware:

    # #region TraceContextMiddleware.__init__ [TYPE Function]
# @ingroup Core
    # @BRIEF Store the inner ASGI app.
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    # #endregion TraceContextMiddleware.__init__

    # #region TraceContextMiddleware.__call__ [C:2] [TYPE Function] [SEMANTICS asgi, call, trace, seed, header]
# @ingroup Core
    # @BRIEF ASGI __call__ that seeds trace_id and injects X-Trace-Id response header.
    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        # Extract X-Trace-ID from raw ASGI scope headers
        incoming_trace_id = None
        for header_name, header_value in scope.get("headers", []):
            if header_name.lower() == b"x-trace-id":
                incoming_trace_id = header_value.decode("utf-8", errors="replace")
                break

        if incoming_trace_id:
            try:
                parsed = uuid.UUID(hex=incoming_trace_id)
                if parsed.version == 4:
                    set_trace_id(incoming_trace_id)
                else:
                    seed_trace_id()
            except (ValueError, AttributeError):
                seed_trace_id()
        else:
            seed_trace_id()

        # Capture trace_id for response header injection
        response_trace_id = get_trace_id()

        async def send_with_trace(message: dict) -> None:
            if message["type"] == "http.response.start":
                headers = list(message.get("headers", []))
                headers.append(
                    (b"x-trace-id", response_trace_id.encode("ascii"))
                )
                message["headers"] = headers
            await send(message)

        await self.app(scope, receive, send_with_trace)

    # #endregion TraceContextMiddleware.__call__


# #endregion TraceContextMiddleware
# #endregion TraceContextMiddlewareModule
