from __future__ import annotations

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

NO_CACHE = b"no-cache, must-revalidate"
NO_STORE = b"no-store"
IMMUTABLE = b"public, max-age=31536000, immutable"


def _cache_control_for_path(path: str) -> bytes | None:
    if path in ("/sw.js", "/index.html", "/manifest.json"):
        return NO_CACHE
    if path == "/version.json":
        return NO_STORE
    if path.startswith("/assets/"):
        return IMMUTABLE
    if path.endswith((".js", ".css")) and not path.startswith("/assets/"):
        return NO_CACHE
    return None


class CacheControlMiddleware:
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

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

        path = scope.get("path", "")
        cache_control = _cache_control_for_path(path)

        async def send_wrapper(message: dict) -> None:
            if message["type"] == "http.response.start" and cache_control is not None:
                headers = list(message.get("headers", []))
                headers = [(k, v) for k, v in headers if k.lower() != b"cache-control"]
                headers.append((b"cache-control", cache_control))
                message = {**message, "headers": headers}
            await send(message)

        await self.app(scope, receive, send_wrapper)