Source code for gpp_client.rest.client

"""
REST API client for non-GraphQL requests.
"""

__all__ = ["RESTClient"]

import asyncio
import gzip
import logging
import ssl
from datetime import datetime, timezone

import aiohttp
import certifi

logger = logging.getLogger(__name__)




class RESTClient:
    """
    REST API client to non-GraphQL requests that help with the function of managers and
    coordinators.

    Parameters
    ----------
    base_url : str
        Base URL of the REST API. Derived from GPPClient base_url.
    gpp_token : str
        GPP token to authenticate against the REST API. Same as GPPClient.
    timeout : float
        Timeout for REST API requests in seconds.
    """

    _DEFAULT_TIMEOUT = 30.0  # Seconds.

    def __init__(
        self, base_url: str, gpp_token: str, timeout: float = _DEFAULT_TIMEOUT
    ) -> None:
        self.base_url = base_url
        self.gpp_token = gpp_token
        self._timeout = timeout

        self._session: aiohttp.ClientSession | None = None
        self._lock = asyncio.Lock()

    def _resolve_headers(self) -> dict[str, str]:
        """
        Resolve the headers for the REST API requests.

        Returns
        -------
        dict[str, str]
            Headers dictionary.
        """
        headers = {
            "Content-Type": "text/plain",
            "Authorization": f"Bearer {self.gpp_token}",
        }
        return headers

    def _create_session(self) -> aiohttp.ClientSession:
        """
        Create a new aiohttp client session.

        Returns
        -------
        aiohttp.ClientSession
            Configured aiohttp client session.
        """
        ssl_context = ssl.create_default_context(cafile=certifi.where())
        connector = aiohttp.TCPConnector(ssl=ssl_context)
        return aiohttp.ClientSession(
            base_url=self.base_url,
            timeout=aiohttp.ClientTimeout(total=self._timeout),
            connector=connector,
            headers=self._resolve_headers(),
        )



    async def get_session(self) -> aiohttp.ClientSession:
        """
        Get or create an ``aiohttp.ClientSession`` for making REST API requests.

        Returns
        -------
        aiohttp.ClientSession
            The aiohttp ClientSession instance.
        """
        async with self._lock:
            if self._session is None or self._session.closed:
                self._session = self._create_session()
            return self._session




    async def close(self) -> None:
        """
        Close the session if it exists and is not already closed.
        """
        if self._session and not self._session.closed:
            await self._session.close()


    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()



    async def get_atom_digests(
        self, observation_ids: list[str], accept_gzip: bool = True
    ) -> str:
        """
        Request atom digests for the given observation IDs.

        Parameters
        ----------
        observation_ids : list[str]
            (internal) IDs of the observation to request.

        accept_gzip : bool
            Endpoint allows to return a gzip-compressed representation of the
            complete atoms digest in the case of being too big.

        Returns
        -------
        str
            TSV data as string, one row per line.

        Raises
        ------
        aiohttp.ClientResponseError
            For HTTP error responses.
        aiohttp.ClientError
            For connection or timeout failures.
        """
        headers = {}
        if accept_gzip:
            headers["Accept-Encoding"] = "gzip"

        # Prepare body - one observation ID per line.
        body = "\n".join(observation_ids)

        session = await self.get_session()

        async with session.post(
            "/scheduler/atoms", data=body, headers=headers
        ) as response:
            # Handle different response codes
            if response.status == 400:
                error_text = await response.text()
                raise ValueError(f"Invalid observation IDs: {error_text}")
            elif response.status == 403:
                raise aiohttp.ClientResponseError(
                    request_info=response.request_info,
                    history=response.history,
                    status=response.status,
                    message="Access forbidden - check authentication and permissions",
                )
            elif response.status != 200:
                response.raise_for_status()

            # Handle gzipped response
            content_encoding = response.headers.get("Content-Encoding", "").lower()
            if content_encoding == "gzip":
                try:
                    content = await response.read()
                    return gzip.decompress(content).decode("utf-8")
                except gzip.BadGzipFile:
                    # Server claimed gzip but sent plain text.
                    return await response.text()
            else:
                return await response.text()




    async def get_visibility_changes(self, since: datetime) -> str:
        """
        Request observations and targets with visibility changes since a time.

        Parameters
        ----------
        since : datetime
            Return entities whose visibility-relevant inputs changed at or
            after this time. Naive datetimes are assumed to be UTC.

        Returns
        -------
        str
            TSV data as string, one ``<gid>\\t<iso8601-timestamp>`` per line.

        Raises
        ------
        aiohttp.ClientResponseError
            For HTTP error responses.
        aiohttp.ClientError
            For connection or timeout failures.
        """
        if since.tzinfo is None:
            since = since.replace(tzinfo=timezone.utc)

        session = await self.get_session()

        async with session.get(
            "/scheduler/visibility-changes",
            params={"since": since.astimezone(timezone.utc).isoformat()},
        ) as response:
            response.raise_for_status()
            return await response.text()