# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/API_Client.ipynb (unless otherwise specified).

__all__ = ["Client"]

# Cell

from typing import *

# Internal Cell

import os
import importlib
import urllib.parse
import json
import secrets

from fastcore.foundation import patch
import pandas as pd

import airt
from ..logger import get_logger, set_level
from ..helper import get_data, post_data, delete_data, get_base_url
from ..constant import (
    SERVICE_USERNAME,
    SERVICE_PASSWORD,
    SERVER_URL,
    CLIENT_NAME,
    SERVICE_TOKEN,
)

# Internal Cell

logger = get_logger(__name__)

# Internal Cell


def _get_credentials(
    username: Optional[str] = None, password: Optional[str] = None
) -> Tuple[(str, str)]:
    """Returns the value for username and password.

    If username is **None**, retrive the value from AIRT_SERVICE_USERNAME environment variable.
    If password is **None**, retrive the value from AIRT_SERVICE_PASSWORD environment variable.

    Args:
        username: Username for your developer account.
        password: Password for your developer account.

    Returns:
        The values for username and password as a tuple.

    Raises:
        Key Error, if the environment variables are not set.
    """

    username = username if username is not None else os.environ.get(SERVICE_USERNAME)

    password = password if password is not None else os.environ.get(SERVICE_PASSWORD)

    if not username and not password:
        raise KeyError(
            f"The username and password are neither passed as parameters nor set in the environment variables "
            f"`{SERVICE_USERNAME}` and `{SERVICE_PASSWORD}`."
        )

    elif not username:
        raise KeyError(
            f"The username is neither passed as parameter nor set in the environment variable {SERVICE_USERNAME}."
        )

    elif not password:
        raise KeyError(
            f"The password is neither passed as parameter nor set in the environment variable {SERVICE_PASSWORD}."
        )

    return (username, password)


# Cell


class Client:
    """A class for authenticating and accessing the airt service.

    To access the airt service, you must create a developer account. Please contact us by email info@airt.ai to get one.

    Upon successful verification, you will receive the username/password for the developer account in an email.

    Finally, you need an application token to access all the APIs in airt service. Please call the `get_token` method with the username/password to get one. You
    can either pass the username, password, and server address as parameters to the `get_token` method or store the same in the **AIRT_SERVICE_USERNAME**,
    **AIRT_SERVICE_PASSWORD**, and **AIRT_SERVER_URL** environment variables.

    Upon successful authentication, the airt services will be available to access.

    As an additional layer of security, we also support Multi-Factor Authentication (MFA) and Single sign-on (SSO) for generating tokens.

    Multi-Factor Authentication (MFA) can be used to help protect your account from unauthorized access by requiring
    you to enter an additional code when you request a new token.

    Once authenticated successfully, activating MFA for your account is a simple two step process:

    1. Enable MFA for your account by calling `User.enable_mfa` method which will generate a QR code. You can
    then use an authenticator app, such as Google Authenticator to scan the QR code.

    2. Activate the MFA by calling `User.activate_mfa` method and passing the dynamically generated six-digit
    verification code from the authenticator app.

    Once MFA is successfully activated, you need to pass the dynamically generated six-digit verification code
    along with the username/password to `get_token` method for generating new tokens.

    Once authenticated successfully, you can also enable Single sign-on (SSO) for your account. Currently, we support only Google and
    Github as the external authentication providers (SSO). More authentication providers will be supported in the upcoming releases.

    Authenticating using Single sign-on (SSO) is also a three-step process:

    1. Enable the SSO provider by calling the `User.enable_sso` method with a valid SSO provider and an email address.

    2. To get the token, you must have to complete the SSO authentication with the provider. Calling the `Client.get_token` method with a valid SSO provider will give you an authorization URL. Please copy and paste it into your preferred browser and initiate the authentication and authorization process with the SSO provider.

    3. Once the authentication is successful, calling the `Client.set_sso_token` method gets a new developer token and will implicitly use in all the interactions with the airt server.
    """

    server: Optional[str] = None
    auth_token: Optional[str] = None
    sso_authorization_url: Optional[str] = None

    def __init__(
        self,
        server: str,
        auth_token: str,
        sso_authorization_url: Optional[str] = None,
    ):
        Client.server = server
        Client.auth_token = auth_token
        Client.sso_authorization_url = sso_authorization_url

    @classmethod
    def get_token(  # type: ignore
        cls,
        *,
        username: str = None,
        password: str = None,
        server: str = None,
        sso_provider: Optional[str] = None,
        otp: Optional[int] = None,
    ) -> Optional[str]:
        """Get application token for airt service from a username/password pair.

        This methods validates the developer credentials and returns an auth token. The returned auth
        token is implicitly used in all the interactions with the server.

        In case the token is requested using Single sign-on (SSO), then an authorization URL will be
        returned to the user.

        The user must have to complete the authentication from the external SSO provider within 10
        minutes. Else the SSO login will be a timeout and the user have to request the token again.

        In case of MFA enabled user, along with the username/password the dynamically generated
        six-digit verification code also needs to be passed when requesting an auth token.

        Args:
            username: Username for the developer account. If None (default value), then the value from
                **AIRT_SERVICE_USERNAME** environment variable is used.
            password: Password for the developer account. If None (default value), then the value from
                **AIRT_SERVICE_PASSWORD** environment variable is used.
            server: The airt server uri. If None (default value), then the value from **AIRT_SERVER_URL** environment variable
                is used. If the variable is not set as well, then the default public server will be used. Please leave this
                setting to default unless you are running the service in your own server (please email us to info@airt.ai
                for that possibility).
            sso_provider: Name of the Single sign-on (SSO) provider. Please pass this parameter only if you have successfully
                enabled SSO for this provider.
            otp: Dynamically generated six-digit verification code from the authenticator app.

        Returns:
            The authorization url if the token is requested using Single sign-on (SSO).

        Raises:
            ValueException: If the username/password pair does not match.
            ConnectionError: If the server address is invalid or not reachable.
            KeyError: If username/password is neither passed as parameters nor stored in environment variables.
        """
        cls.server = get_base_url(server)

        username, password = _get_credentials(username, password)

        if otp is not None:
            password = json.dumps({"password": password, "user_otp": otp})

        if sso_provider is None:
            response = post_data(
                url=f"{cls.server}/token",
                data=dict(username=username, password=password),
                token=None,
            )

            cls.auth_token = response["access_token"]
        else:
            response = post_data(
                url=f"{cls.server}/sso/initiate",
                data=json.dumps(  # type: ignore
                    dict(
                        username=username, password=password, sso_provider=sso_provider
                    )
                ),
                token=None,
            )

            cls.sso_authorization_url = response["authorization_url"]
            return cls.sso_authorization_url

    @classmethod
    def set_sso_token(cls):
        """Set the application token generated using Single sign-on (SSO).

        The token set using this method will be implicitly used in all the interactions with the server.

        Please call this method only if you successfully enabled and completed the login with the Single
        sign-on (SSO) provider. If not, please call the `get_token` method with an appropriate
        sso_provider to initiate the SSO authentication.
        """
        quoted_authorization_url = urllib.parse.quote(cls.sso_authorization_url)
        response = get_data(
            url=f"{cls.server}/sso/token/?authorization_url={quoted_authorization_url}",
            token=None,
        )

        cls.auth_token = response["access_token"]

    @classmethod
    def set_token(cls, token: str = None, server: str = None):
        """Set application token for airt service.

        The token set using this method will be implicitly used in all the interactions with the server.

        Please call this method only if you already have a token. If not, please call the `get_token` method to generate one.

        Args:
            token: The application token generated from a username/password pair. If None (default value), then the value from
                **AIRT_SERVICE_TOKEN** environment variable is used.
            server: The airt server uri. If None (default value), then the value from **AIRT_SERVER_URL** environment variable
                is used. If the variable is not set as well, then the default public server will be used. Please leave this
                setting to default unless you are running the service in your own server (please email us to info@airt.ai
                for that possibility).

        An example to set the token:

        ```python
        Client.set_token()
        ```
        """

        auth_token = token if token is not None else os.environ.get(SERVICE_TOKEN)

        if not auth_token:
            raise KeyError(
                f"The token is neither passed as parameter nor set in the environment variable {SERVICE_TOKEN}."
            )

        cls.auth_token = auth_token
        cls.server = get_base_url(server)

    @staticmethod
    def version() -> dict:
        """Return the client and server versions.

        Returns:
            A dict containing the client and server versions.

        Raises:
            ConnectionError: If the server address is invalid or not reachable.

        An example to get the client and server versions:

        ```python
        Client.version()
        ```
        """

        response = Client._get_data(relative_url=f"/version")

        version = {
            # nosemgrep: python.lang.security.audit.non-literal-import.non-literal-import
            "client": importlib.import_module(CLIENT_NAME).__version__,
            "server": response["airt_service"],
        }

        return version

    @classmethod
    def _get_server_url_and_token(cls) -> Tuple[Optional[str], Optional[str]]:
        """Fetch the server URL and the auth token.

        Returns:
            A tuple containing server URL and auth token.
        """

        cls.server = get_base_url(cls.server)

        return cls.server, cls.auth_token

    @classmethod
    def _post_data(
        cls, relative_url: str, data: Optional[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Make a POST request.

        This method will implicitly add the server base URL and the token for every request.

        Args:
            relative_url: The relative URL of the server's API endpoint.
            data: A Dictionary object to send in the body of the POST request.

        Returns:
            Response body as a dictionary.

        Raises:
            ConnectionError: If the server is not reachable.
            ValueError: If the response code is not in range of 200 - 399.
        """

        server, auth_token = Client._get_server_url_and_token()

        return post_data(
            url=f"{server}{relative_url}",
            data=data,
            token=auth_token,
        )

    @classmethod
    def _get_data(cls, relative_url: str) -> Any:
        """Make a GET request.

        This method will implicitly add the server base URL and the token for every request.

        Args:
            relative_url: The relative URL of the API endpoint.

        Returns:
            A dictionary that encapsulates the response body.

        Raises:
            ConnectionError: If the server is not reachable.
            ValueError: If the response code is not in range of 200 - 399.
        """

        server, auth_token = Client._get_server_url_and_token()

        return get_data(url=f"{server}{relative_url}", token=auth_token)

    @classmethod
    def _delete_data(cls, relative_url: str) -> Dict[str, Any]:
        """Make a DELETE request.

        This method will implicitly add the server base URL and the token for every request.

        Args:
            relative_url: The relative URL of the API endpoint.

        Returns:
            A dictionary that encapsulates the response body.

        Raises:
            ConnectionError: If the server is not reachable.
            ValueError: If the response code is not in range of 200 - 399.
        """

        server, auth_token = Client._get_server_url_and_token()

        return delete_data(url=f"{server}{relative_url}", token=auth_token)
