"""ASGI middleware for request context."""

import uuid

from starlette.datastructures import Headers, MutableHeaders
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from owscontext import constants
from owscontext.context.base import (
    get_correlation_id,
    request_context_from_headers,
    reset_correlation_id,
    reset_request_context,
    set_correlation_id,
    set_request_context,
)


class CorrelationIdMiddleware:
    """ASGI middleware that ensures a correlation id is present on every request."""

    def __init__(
        self, app: ASGIApp, header_name: str = constants.HEADER_CORRELATION_ID
    ) -> None:
        """Initialise the middleware with the wrapped ASGI app and header name."""
        self.app = app
        self.header_name = header_name

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Handle an ASGI request, injecting a correlation id from or into headers."""
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        if not (header_value := Headers(scope=scope).get(self.header_name.lower())):
            correlation_id_value = str(uuid.uuid4())
        else:
            correlation_id_value = header_value

        token = set_correlation_id(correlation_id_value)

        async def _send(message: Message) -> None:
            if message["type"] == "http.response.start" and (
                _correlation_id_value := get_correlation_id()
            ):
                headers = MutableHeaders(scope=message)
                headers.append(self.header_name, _correlation_id_value)

            await send(message)

        try:
            await self.app(scope, receive, _send)
        finally:
            reset_correlation_id(token)


class RequestContextMiddleware:
    """ASGI middleware that extracts request context from headers."""

    def __init__(self, app: ASGIApp) -> None:
        """Initialise the middleware with the wrapped ASGI app."""
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Handle an ASGI request, injecting request context from headers."""
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        request_context = request_context_from_headers(headers=Headers(scope=scope))

        token = set_request_context(request_context)

        try:
            await self.app(scope, receive, send)
        finally:
            reset_request_context(token)