from base64 import b64decode
from chitai.services.user import UserService
from litestar.middleware import (
    AbstractAuthenticationMiddleware,
    AuthenticationResult,
    DefineMiddleware,
)
from litestar.connection import ASGIConnection
from litestar.exceptions import NotAuthorizedException, PermissionDeniedException
from chitai.config import settings


class BasicAuthenticationMiddleware(AbstractAuthenticationMiddleware):
    async def authenticate_request(
        self, connection: ASGIConnection
    ) -> AuthenticationResult:
        """Given a request, parse the header for Base64 encoded basic auth credentials."""

        # retrieve the auth header
        auth_header = connection.headers.get("Authorization", None)
        if not auth_header:
            raise NotAuthorizedException()

        username, password = (
            b64decode(auth_header.split("Basic ")[1]).decode().split(":")
        )

        try:
            db_session = settings.alchemy_config.provide_session(
                connection.app.state, connection.scope
            )
            user_service = UserService(db_session)
            user = await user_service.authenticate(username, password)
            return AuthenticationResult(user=user, auth=None)

        except PermissionDeniedException:
            raise NotAuthorizedException()


basic_auth_mw = DefineMiddleware(BasicAuthenticationMiddleware)
