import uuid

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

from audience_common import constants
from audience_common.context import (
    get_correlation_id,
    request_context_from_headers,
    reset_correlation_id,
    reset_request_context,
    set_correlation_id,
    set_request_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 = 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:
    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 = request_context_from_headers(
            headers=Headers(scope=scope), label_profile=self.label_profile
        )

        token = set_request_context(request_context)

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