auth

class KeyCloakAuth(httpx.Auth):
    """Auth class for KeyCloak authentication"""

    def __init__(
        self,
        keycloak_server_url: str,
        realm: str,
        client_id: str,
        username: str,
        password: str,
    ):
        """Initialize the KeyCloakAuth class"""
        self.token_url = (
            f"{keycloak_server_url}/realms/{realm}/protocol/openid-connect/token"
        )
        self.client_id = client_id
        self.username = username
        self.password = password
        self.realm = realm
        self._token = None

    def auth_flow(self, request):
        """Flow for the authentication.

        Note:
            See [httpx documentation](https://www.python-httpx.org/advanced/authentication) on custom authentication schemes
        """
        request.headers["Authorization"] = f"Bearer {self.token}"
        response = yield request
        if response.status_code == 401:
            logger.info("Token expired, retrieving new token")
            self._token = None  # Clear the token
            request.headers["Authorization"] = f"Bearer {self.token}"
            yield request

    @property
    def token(self):
        if self._token is None:
            self._token = self._retrieve_token()
        return self._token

    def _retrieve_token(self):
        # Use the token_url to get the token
        with httpx.Client() as client:
            logger.info(f"Retrieving token from {self.token_url}")
            headers = {"Content-Type": "application/x-www-form-urlencoded"}
            response = client.post(
                self.token_url,
                data={
                    "client_id": self.client_id,
                    "username": self.username,
                    "password": self.password,
                    "scope": "openid email profile",
                    "grant_type": "password",
                },
                headers=headers,
            )
            response.raise_for_status()
            return response.json()["access_token"]

class ApiKeyAuth(httpx.Auth):
    """Auth class for Pixel API key authentication.

    An api key is a long lived, opaque bearer token minted by
    [`create_api_key`][pixel_client.PixelClient.create_api_key]. Unlike a
    Keycloak token it never expires mid-session and cannot be refreshed, so a
    401 means the key was revoked, expired, or lacks the required role.
    """

    def __init__(self, api_key: str):
        """Initialize the ApiKeyAuth class.

        Args:
            api_key: The api key token, as returned once by the create endpoint.
        """
        if not api_key:
            raise ValueError("api_key must be a non-empty string")
        self._api_key = api_key

    def auth_flow(self, request):
        """Flow for the authentication.

        Note:
            See [httpx documentation](https://www.python-httpx.org/advanced/authentication) on custom authentication schemes
        """
        request.headers["Authorization"] = f"Bearer {self._api_key}"
        yield request