afnio.tellurio.client

class InvalidAPIKeyError(Exception):
    """Exception raised when the API key is invalid."""

    pass

class TellurioClient:
    """A client for interacting with the
    [Tellurio Studio](https://platform.tellurio.ai/) backend.

    This client provides methods for authenticating with the backend, making HTTP
    requests (GET, POST, DELETE), and verifying API keys. It is designed to simplify
    communication with the [Tellurio Studio](https://platform.tellurio.ai/) platform.
    """

    def __init__(self, base_url: str = None, port: int = None):
        """Initializes the `TellurioClient` instance.

        Args:
            base_url: The base URL of the Tellurio backend. If not
                provided, it defaults to the value of the
                `TELLURIO_BACKEND_HTTP_BASE_URL` environment variable
                or `"https://platform.tellurio.ai"`.
            port: The port number for the backend. If not provided,
                it defaults to the value of the `TELLURIO_BACKEND_HTTP_PORT`
                environment variable or 443.
        """
        self.base_url = base_url or os.getenv(
            "TELLURIO_BACKEND_HTTP_BASE_URL", "https://platform.tellurio.ai"
        )
        self.port = port or os.getenv("TELLURIO_BACKEND_HTTP_PORT", 443)
        self.url = f"{self.base_url}:{self.port}"
        self.service_name = os.getenv(
            "KEYRING_SERVICE_NAME", "Tellurio"
        )  # Service name for keyring
        self.api_key = None

    def login(self, api_key: str = None, relogin: bool = False) -> dict:
        """Logs in the user using an API key and verifies its validity.

        Credential resolution order:

        1. If `api_key` is provided, it is used.
        2. Otherwise, if the `TELLURIO_API_KEY` environment variable is set,
            it is used.
        3. Otherwise, if not relogin, attempts to load a stored API key from
            the keyring.

        If authentication succeeds and the API key was provided directly (not via
        keyring), it is stored in the keyring for future use.

        Args:
            api_key: The user's API key. If not provided, the method
                attempts to use the `TELLURIO_API_KEY` environment variable, then the
                keyring.
            relogin: If True, forces a re-login and requires a new API key.

        Returns:
            A dictionary containing the user's email and username.

        Raises:
            afnio.tellurio.client.InvalidAPIKeyError: If the API key is invalid.
            ValueError: If the API key is invalid or not provided during re-login.
        """
        # Use the provided API key if passed, otherwise check env var, then keyring
        if api_key:
            self.api_key = api_key
        elif os.getenv("TELLURIO_API_KEY"):
            self.api_key = os.getenv("TELLURIO_API_KEY")
            logger.info("Using API key from TELLURIO_API_KEY environment variable.")
        elif not relogin:
            username = load_username()
            self.api_key = keyring.get_password(self.service_name, username)
            if self.api_key:
                logger.info("Using stored API key from local keyring.")
            else:
                logger.error("No API key found in local keyring.")
                raise ValueError("API key is required for the first login.")
        else:
            logger.error("No API key provided for re-login.")
            raise ValueError("API key is required for re-login.")

        # Verify the API key
        response_data = self._verify_api_key()
        if response_data:
            email = response_data.get("email", "unknown user")
            username = response_data.get("username", "unknown user")
            logger.debug(f"API key is valid for user '{username}'.")

            # Save the API key securely only if it was provided and is valid
            if api_key or os.getenv("TELLURIO_API_KEY"):
                if _is_keyring_usable():
                    keyring.set_password(
                        self.service_name, response_data["username"], self.api_key
                    )
                    logger.info(
                        "API key provided and stored securely in local keyring."
                    )
                    save_username(response_data["username"])
                else:
                    logger.info(
                        "Keyring is not available; skipping secure storage of API key."
                    )

            return {
                "email": email,
                "username": username,
            }
        else:
            logger.warning("Invalid API key. Please provide a valid API key.")
            if relogin:
                raise InvalidAPIKeyError("Re-login failed due to invalid API key.")
            raise InvalidAPIKeyError("Login failed due to invalid API key.")

    def get(self, endpoint: str) -> httpx.Response:
        """Makes a GET request to the specified endpoint.

        Args:
            endpoint: The API endpoint (relative to the base URL).

        Returns:
            The HTTP response object.

        Raises:
            ValueError: If a network error occurs or an unexpected error happens.
        """
        url = f"{self.url}{endpoint}"
        headers = {
            "Authorization": f"Api-Key {self.api_key}",
            "Accept": "*/*",
        }

        try:
            with httpx.Client() as client:
                response = client.get(url, headers=headers)
            return response
        except httpx.RequestError as e:
            logger.error(f"Network error occurred while making GET request: {e}")
            raise ValueError("Network error occurred. Please check your connection.")
        except Exception as e:
            logger.error(f"An unexpected error occurred: {e}")
            raise ValueError("An unexpected error occurred. Please try again later.")

    def post(self, endpoint: str, json: dict) -> httpx.Response:
        """Makes a POST request to the specified endpoint.

        Args:
            endpoint: The API endpoint (relative to the base URL).
            json: The JSON payload to send in the request.

        Returns:
            The HTTP response object.

        Raises:
            ValueError: If a network error occurs or an unexpected error happens.
        """
        url = f"{self.url}{endpoint}"
        headers = {
            "Authorization": f"Api-Key {self.api_key}",
            "Content-Type": "application/json",
        }

        try:
            with httpx.Client() as client:
                response = client.post(url, headers=headers, json=json)
            return response
        except httpx.RequestError as e:
            logger.error(f"Network error occurred while making POST request: {e}")
            raise ValueError("Network error occurred. Please check your connection.")
        except Exception as e:
            logger.error(f"An unexpected error occurred: {e}")
            raise ValueError("An unexpected error occurred. Please try again later.")

    def patch(self, endpoint: str, json: dict) -> httpx.Response:
        """Makes a PATCH request to the specified endpoint.

        Args:
            endpoint: The API endpoint (relative to the base URL).
            json: The JSON payload to send in the request.

        Returns:
            The HTTP response object.

        Raises:
            ValueError: If a network error occurs or an unexpected error happens.
        """
        url = f"{self.url}{endpoint}"
        headers = {
            "Authorization": f"Api-Key {self.api_key}",
            "Content-Type": "application/json",
        }

        try:
            with httpx.Client() as client:
                response = client.patch(url, headers=headers, json=json)
            return response
        except httpx.RequestError as e:
            logger.error(f"Network error occurred while making PATCH request: {e}")
            raise ValueError("Network error occurred. Please check your connection.")
        except Exception as e:
            logger.error(f"An unexpected error occurred: {e}")
            raise ValueError("An unexpected error occurred. Please try again later.")

    def delete(self, endpoint: str) -> httpx.Response:
        """Makes a DELETE request to the specified endpoint.

        Args:
            endpoint: The API endpoint (relative to the base URL).

        Returns:
            The HTTP response object.

        Raises:
            ValueError: If a network error occurs or an unexpected error happens.
        """
        url = f"{self.url}{endpoint}"
        headers = {
            "Authorization": f"Api-Key {self.api_key}",
            "Accept": "*/*",
        }

        try:
            with httpx.Client() as client:
                response = client.delete(url, headers=headers)
            return response
        except httpx.RequestError as e:
            logger.error(f"Network error occurred while making DELETE request: {e}")
            raise ValueError("Network error occurred. Please check your connection.")
        except Exception as e:
            logger.error(f"An unexpected error occurred: {e}")
            raise ValueError("An unexpected error occurred. Please try again later.")

    def _verify_api_key(self) -> dict:
        """Verifies the validity of the API key
        by calling the `/api/v0/verify-api-key/` endpoint.

        Returns:
            A dictionary containing the user's email, username and a message \
            indicating if the API key is valid, None otherwise.
        """
        endpoint = "/api/v0/verify-api-key/"
        try:
            response = self.get(endpoint)

            if response.status_code == 200:
                try:
                    data = response.json()
                    logger.debug(f"API key verification successful: {data}")
                    return data
                except ValueError:
                    logger.error("Failed to parse JSON response from backend.")
                    return None
            elif response.status_code == 401:
                logger.warning("API key is invalid or missing.")
            else:
                logger.error(f"Error: {response.status_code} - {response.text}")
        except ValueError as e:
            logger.error(f"Error during API key verification: {e}")
            raise

        return None