from time import perf_counter

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


class ResponseTimeMiddleware:
    def __init__(self, app: ASGIApp, header_name: str = "Elapsed") -> None:
        self.app = app
        self.header_name = header_name

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        start = perf_counter()

        if scope["type"] not in ("http", "websocket"):
            await self.app(scope, receive, send)
            return

        async def send_wrapper(message: Message) -> None:
            if message["type"] == "http.response.start":
                headers = MutableHeaders(scope=message)
                elapsed = perf_counter() - start
                headers.append(self.header_name, f"{elapsed:.20f}")
            await send(message)

        await self.app(scope, receive, send_wrapper)
