import uuid

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

from owslib import constants, context


class CorrelationIdMiddleware:
    def __init__(
        self, app: ASGIApp, header_name: str = constants.HEADER_CORRELATION_ID
    ) -> None:
        self.app = app
        self.header_name = header_name

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        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 = context.correlation_id.set(correlation_id_value)

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

            await send(message)

        try:
            await self.app(scope, receive, send_response)
        finally:
            context.correlation_id.reset(token)

        return


class RequestContextMiddleware:
    def __init__(self, app: ASGIApp, label_profile: bool = False) -> None:
        self.app = app
        self.label_profile = label_profile

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        request_context = context.request_context_from_headers(
            headers=Headers(scope=scope), label_profile=self.label_profile
        )

        token = context.request_context.set(request_context)

        try:
            await self.app(scope, receive, send)
        finally:
            context.request_context.reset(token)

        return