from contextvars import ContextVar
from typing import Optional, Union

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.websockets import WebSocket

_request_context: ContextVar[Union[Request, None]] = ContextVar("request", default=None)


class RequestMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        _request_context.set(request)

        from core.logger import APILogger

        logger = APILogger(__name__)
        logger.trace("Set new current request")
        response = await call_next(request)
        return response


def set_request(request: Union[Request, WebSocket, None]):
    _request_context.set(request)


def get_request() -> Optional[Request]:
    return _request_context.get()
