from uuid import uuid4

from ariadne.asgi.handlers import GraphQLTransportWSHandler
from starlette.datastructures import MutableHeaders
from starlette.websockets import WebSocket

from api.middleware.request import set_request
from core.logger import APILogger

logger = APILogger("api.ws")


class AuthenticatedGraphQLTransportWSHandler(GraphQLTransportWSHandler):

    def __init__(self, *args, **kwargs):
        super().__init__(
            *args,
            on_connect=self.on_connect,
            on_disconnect=self.on_disconnect,
            **kwargs,
        )

    @staticmethod
    async def on_connect(ws: WebSocket, message: dict):
        ws.state.request_id = uuid4()
        logger.info(f"WebSocket connection {ws.state.request_id}")

        if "Authorization" not in message:
            return True

        mutable_headers = MutableHeaders()
        mutable_headers["Authorization"] = message.get("Authorization", "")
        ws._headers = mutable_headers

        set_request(ws)
        return True

    @staticmethod
    async def on_disconnect(ws: WebSocket):
        logger.debug(f"WebSocket connection {ws.state.request_id} closed")
        return True
