# 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):
        """Constructs a new ProgressStatus instance.

        Warning:
            Do not construct this object directly by calling the constructor, please use `DataSource.pull` or `DataSource.train` or `Model.predict` functions instead.

        Args:
            relative_url: The relative URL to the rest API endpoint.
        """
        self.relative_url = relative_url

    def is_ready(self) -> bool:
        """A function to check if the method's progress is completed.

        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, sleep_for: Union[int, float] = 5, timeout: int = 0):
        """Blocks execution while waiting for remote action to be completed and displays a progress bar indicating the completion status.

        Args:
            sleep_for: The time interval in seconds between successive API calls to ping the server for fetching the completed steps.
            timeout: The maximum time allowed in seconds for the asynchronous call to complete the process. If the timeout
                exceeds and the process is yet to complete, then the progress_bar will be terminated.

        Raises:
            TimeoutError: in case of 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 < timeout) and (datetime.now() - started_at) > timedelta(
                    seconds=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(sleep_for)

    def wait(self, sleep_for: Union[int, float] = 1, timeout: int = 0):
        raise NotImplementedError()


# Cell


@patch
def wait(self: ProgressStatus, sleep_for: Union[int, float] = 1, timeout: int = 0):
    """Blocks execution while waiting for remote action to be completed.

    Args:
        sleep_for: The time interval in seconds between successive API calls to ping the server for fetching the completed steps.
        timeout: The maximum time allowed in seconds for the asynchronous call to complete the process. If the timeout
            exceeds and the process is yet to complete, then the progress_bar will be terminated.

    Raises:
        TimeoutError: in case of timeout
    """
    started_at = datetime.now()
    while True:
        if (0 < timeout) and (datetime.now() - started_at) > timedelta(seconds=timeout):
            raise TimeoutError()

        if self.is_ready():
            return

        sleep(sleep_for)
