# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_management.ipynb.

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

# %% ../nbs/01_management.ipynb 2
from .base import BaseClient

# %% ../nbs/01_management.ipynb 3
class ManagementClient(BaseClient):
    def __init__(self, base_url=None, port=8081):
        super().__init__(base_url)
        self.port = port
        self.base_url = f"{self.base_url}:{self.port}"

    def register_model(
        self,
        url,
        model_name=None,
        handler=None,
        runtime=None,
        batch_size=1,
        max_batch_delay=100,
        initial_workers=0,
        synchronous=False,
        response_timeout=120,
    ):
        data = {
            "url": url,
            "model_name": model_name,
            "handler": handler,
            "runtime": runtime,
            "batch_size": batch_size,
            "max_batch_delay": max_batch_delay,
            "initial_workers": initial_workers,
            "synchronous": synchronous,
            "response_timeout": response_timeout,
        }
        data = self._filter_none_values(data)
        return self._make_request("POST", "/models", json=data)

    def scale_workers(
        self,
        model_name,
        version=None,
        min_worker=1,
        max_worker=None,
        synchronous=False,
        timeout=-1,
    ):
        if version:
            endpoint = f"/models/{model_name}/{version}"
        else:
            endpoint = f"/models/{model_name}"

        params = {
            "min_worker": min_worker,
            "max_worker": max_worker if max_worker is not None else min_worker,
            "synchronous": synchronous,
            "timeout": timeout,
        }

        return self._make_request("PUT", endpoint, data=params)

    def describe_model(self, model_name, version=None, customized=False):
        """
        Returns the model description.
        version :  is optional. if `all` return status of all version of a model. If not provided, the latest version will be returned.
        allowed
        """
        params = {}
        if customized:
            params["customized"] = customized

        if version:
            endpoint = f"/models/{model_name}/{version}"
        else:
            endpoint = f"/models/{model_name}"
        return self._make_request("GET", endpoint, params=params)

    def unregister_model(self, model_name, version=None):
        if version:
            endpoint = f"/models/{model_name}/{version}"
        else:
            endpoint = f"/models/{model_name}"

        return self._make_request("DELETE", endpoint)

    def list_models(self, limit=100, next_page_token=None):
        params = {"limit": limit, "next_page_token": next_page_token}
        params = self._filter_none_values(params)
        return self._make_request("GET", "/models", params=params)

    def api_description(self):
        return self._make_request("OPTIONS", "/")

    def set_default_version(self, model_name, version):
        endpoint = f"/models/{model_name}/{version}/set-default"
        return self._make_request("PUT", endpoint)
