"""Cached static file serving with configurable Cache-Control headers."""

from starlette.staticfiles import StaticFiles
from starlette.types import Receive, Scope, Send


class CachedStaticFiles(StaticFiles):
    """StaticFiles subclass that adds Cache-Control headers to responses."""

    def __init__(self, *args, max_age: int = 86400, **kwargs):
        self.max_age = max_age
        super().__init__(*args, **kwargs)

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        async def send_with_cache(message):
            if message["type"] == "http.response.start":
                headers = list(message.get("headers", []))
                headers.append(
                    (b"cache-control", f"public, max-age={self.max_age}".encode())
                )
                message["headers"] = headers
            await send(message)

        await super().__call__(scope, receive, send_with_cache)
