# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/API_ProgressStatus.ipynb.

# %% auto 0
__all__ = ['ProgressStatus']

# %% ../../notebooks/API_ProgressStatus.ipynb 3
from typing import *

# %% ../../notebooks/API_ProgressStatus.ipynb 4
from time import sleep
from datetime import datetime, timedelta

from tqdm import tqdm

from fastcore.foundation import patch

from airt.components.client import Client
from airt.logger import get_logger, set_level

# %% ../../notebooks/API_ProgressStatus.ipynb 6
logger = get_logger(__name__)

# %% ../../notebooks/API_ProgressStatus.ipynb 9
class ProgressStatus:
    """A base class for querying status of a remote operation.

    Here's an example of using the ProgressStatus class to checking the upload status of the datablob

    Example:
        ```python
        # Importing necessary libraries
        from  airt.client import Client, DataBlob

        # Authenticate
        Client.get_token(username="{fill in username}", password="{fill in password}")

        # Create a datablob
        # In this example, the datablob will be stored in an AWS S3 bucket. The region
        # is set to eu-west-3, feel free to change the cloud provider and the region
        # to suit your needs.
        db = DataBlob.from_s3(
            uri="{fill in uri}",
            cloud_provider="aws",
            region="eu-west-3"
        )

        # Check the ready status of the datablob
        print(db.is_ready())

        # Display the status in a progress bar
        db.progress_bar()

        # Check the ready status of the datablob
        # If the upload is successful, True will be returned
        print(db.is_ready())
        ```
    """

    def __init__(
        self, relative_url: str, sleep_for: Union[int, float] = 5, timeout: int = 0
    ):
        """Constructs a new ProgressStatus instance.

        Warning:
            Do not construct this object directly by calling the constructor, please use either progress_bar,
            is_ready, or wait methods of `DataBlob`, `DataSource`, `Model` or `Prediction` classes instead.

        Args:
            relative_url: Relative URI to query the status of the remote operation.
            sleep_for: The time interval in seconds between successive API calls.
            timeout: The maximum time allowed in seconds for the asynchronous call to complete. If not the
                progressbar will be terminated.

        Raises:
            TimeoutError: in case of connection timeout.
        """
        self.relative_url = relative_url
        self.sleep_for = sleep_for
        self.timeout = timeout

    def is_ready(self) -> bool:
        """Check if the method's progress is complete.

        Returns:
            **True** if the progress is completed, else **False**.

        Here's an example of checking the upload status of the datablob:

        Example:
            ```python
            # Importing necessary libraries
            from  airt.client import Client, DataBlob

            # Authenticate
            Client.get_token(username="{fill in username}", password="{fill in password}")

            # Create a datablob
            # In this example, the datablob will be stored in an AWS S3 bucket. The region
            # is set to eu-west-3, feel free to change the cloud provider and the region
            # to suit your needs.
            db = DataBlob.from_s3(
                uri="{fill in uri}",
                cloud_provider="aws",
                region="eu-west-3"
            )

            # Check the ready status of the datablob
            print(db.is_ready())

            # Display the status in a progress bar
            db.progress_bar()

            # Check the ready status of the datablob
            # If the upload is successful, True will be returned
            print(db.is_ready())
            ```
        """
        response = Client._get_data(relative_url=self.relative_url)
        return response["completed_steps"] == response["total_steps"]

    def progress_bar(self):
        """Blocks the execution and displays a progress bar showing the remote action progress.

        Raises:
            ConnectionError: If the server address is invalid or not reachable.
            TimeoutError: in case of connection timeout.

        Here's an example of checking the upload status of the datablob:

        Example:
            ```python
            # Importing necessary libraries
            from  airt.client import Client, DataBlob

            # Authenticate
            Client.get_token(username="{fill in username}", password="{fill in password}")

            # Create a datablob
            # In this example, the datablob will be stored in an AWS S3 bucket. The region
            # is set to eu-west-3, feel free to change the cloud provider and the region
            # to suit your needs.
            db = DataBlob.from_s3(
                uri="{fill in uri}",
                cloud_provider="aws",
                region="eu-west-3"
            )

            # Check the ready status of the datablob
            print(db.is_ready())

            # Display the status in a progress bar
            db.progress_bar()

            # Check the ready status of the datablob
            # If the upload is successful, True will be returned
            print(db.is_ready())
            ```
        """
        total_steps = Client._get_data(relative_url=self.relative_url)["total_steps"]
        with tqdm(total=total_steps) as pbar:
            started_at = datetime.now()
            while True:
                if (0 < self.timeout) and (datetime.now() - started_at) > timedelta(
                    seconds=self.timeout
                ):
                    raise TimeoutError()

                response = Client._get_data(relative_url=self.relative_url)
                completed_steps = response["completed_steps"]

                pbar.update(completed_steps)

                if completed_steps == total_steps:
                    break

                sleep(self.sleep_for)

    def wait(self):
        raise NotImplementedError()

# %% ../../notebooks/API_ProgressStatus.ipynb 16
@patch
def wait(self: ProgressStatus):
    """Blocks execution while waiting for the remote action to complete.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.
        TimeoutError: in case of timeout.

    Here's an example of preventing further calls from being executed until the datablob upload is complete:

    Example:
        ```python
        # Importing necessary libraries
        from  airt.client import Client, DataBlob

        # Authenticate
        Client.get_token(username="{fill in username}", password="{fill in password}")

        # Create a datablob
        # In this example, the datablob will be stored in an AWS S3 bucket. The region
        # is set to eu-west-3, feel free to change the cloud provider and the region
        # to suit your needs.
        db = DataBlob.from_s3(
            uri="{fill in uri}",
            cloud_provider="aws",
            region="eu-west-3"
        )

        # Further calls to the API will be blocked until the datablob upload is complete.
        db.wait()

        # Print the details of the newly created datablob
        # If the upload is successful, the ready flag should be set to True
        print(db.details())
        ```
    """
    started_at = datetime.now()
    while True:
        if (0 < self.timeout) and (datetime.now() - started_at) > timedelta(
            seconds=self.timeout
        ):
            raise TimeoutError()

        if self.is_ready():
            return

        sleep(self.sleep_for)
