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

__all__ = ["ProgressStatus"]

# Cell

from typing import *

# Internal Cell

from time import sleep
from datetime import datetime, timedelta

from tqdm import tqdm

from fastcore.foundation import patch

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

# Internal Cell

logger = get_logger(__name__)

# Cell


class ProgressStatus:
    """A base class for querying status of a remote operation."""

    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: Rest API endpoint relative url.
            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 if completed, else **False**.
        """
        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.
        """
        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()


# Cell


@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.
    """
    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)
