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


class FlashMessageMiddleware:
    """Pure ASGI middleware — avoids the extra asyncio.Task that BaseHTTPMiddleware creates."""

    def __init__(self, app: ASGIApp):
        self.app = app

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

        request = Request(scope)

        # Initialize flash messages from session
        if not hasattr(request.state, "flash_messages"):
            request.state.flash_messages = request.session.get("flash_messages", [])

        async def send_wrapper(message):
            # Persist flash messages back to session before response headers are sent
            if message["type"] == "http.response.start":
                if hasattr(request.state, "flash_messages") and request.state.flash_messages:
                    request.session["flash_messages"] = request.state.flash_messages
                    request.state.flash_messages = []
            await send(message)

        await self.app(scope, receive, send_wrapper)


class FlashMessage:

    def __init__(self, message, category='info', exception=None):
        self.message = message
        self.category = category
        self.exception = exception

    def to_dict(self):
        return {
            "message": self.message,
            "category": self.category,
            "exception": self.exception.__dict__ if isinstance(self.exception, Exception) else None,
        }

    @staticmethod
    def add(request: Request, message: str, category: str = "info", exception=None):
        if not hasattr(request.state, "flash_messages"):
            request.state.flash_messages = []

        flash = FlashMessage(message, category, exception)
        request.state.flash_messages.append(flash.to_dict())
        return flash

    @staticmethod
    def get_flash_messages(request: Request):
        flash_messages = request.state.flash_messages if hasattr(request.state, "flash_messages") else []
        # Clear the messages from the session after they are accessed
        request.session.pop("flash_messages", None)
        request.state.flash_messages = []  # Clear after use
        return flash_messages
