Client

class Client:
    """
    Valor client object for interacting with the api.

    Parameters
    ----------
    connection : ClientConnection, optional
        Option to use an existing connection object.
    """

    def __init__(self, connection: Optional[ClientConnection] = None):
        if not connection:
            connection = get_connection()
        self.conn = connection

    @classmethod
    def connect(
        cls,
        host: str,
        access_token: Optional[str] = None,
        reconnect: bool = False,
    ) -> Client:
        """
        Establishes a connection to the Valor API.

        Parameters
        ----------
        host : str
            The host to connect to. Should start with "http://" or "https://".
        access_token : str
            The access token for the host (if the host requires authentication).
        """
        connect(host=host, access_token=access_token, reconnect=reconnect)
        return cls(get_connection())

    def get_labels(
        self,
        filters: Optional[Filter] = None,
    ) -> List[Label]:
        """
        Gets all labels using an optional filter.

        Parameters
        ----------
        filters : Filter, optional
            Optional constraints to filter by.

        Returns
        ------
        List[valor.Label]
            A list of labels.
        """
        filters = filters if filters is not None else Filter()
        return [
            Label(**label) for label in self.conn.get_labels(filters.to_dict())
        ]

    def get_labels_from_dataset(
        self, dataset: Union[Dataset, str]
    ) -> List[Label]:
        """
        Get all labels associated with a dataset's ground truths.

        Parameters
        ----------
        dataset : valor.Dataset
            The dataset to search by.

        Returns
        ------
        List[valor.Label]
            A list of labels.
        """
        dataset_name = (
            dataset.name if isinstance(dataset, Dataset) else dataset
        )
        return [
            Label(**label)
            for label in self.conn.get_labels_from_dataset(dataset_name)  # type: ignore
        ]

    def get_labels_from_model(self, model: Union[Model, str]) -> List[Label]:
        """
        Get all labels associated with a model's ground truths.

        Parameters
        ----------
        model : valor.Model
            The model to search by.

        Returns
        ------
        List[valor.Label]
            A list of labels.
        """
        model_name = model.name if isinstance(model, Model) else model
        return [
            Label(**label)
            for label in self.conn.get_labels_from_model(model_name)  # type: ignore
        ]

    def create_dataset(
        self,
        dataset: Union[Dataset, dict],
    ) -> None:
        """
        Creates a dataset.

        Parameters
        ----------
        dataset : valor.Dataset
            The dataset to create.
        """
        if isinstance(dataset, Dataset):
            dataset = dataset.encode_value()
        self.conn.create_dataset(dataset)

    def create_groundtruths(
        self,
        dataset: Dataset,
        groundtruths: List[GroundTruth],
        ignore_existing_datums: bool = False,
        timeout: Optional[float] = None,
    ):
        """
        Creates ground truths.

        Parameters
        ----------

        dataset : valor.Dataset
            The dataset to create the ground truth for.
        groundtruths : List[valor.GroundTruth]
            The ground truths to create.
        timeout : float, optional
            The number of seconds the client should wait until raising a timeout.
        ignore_existing_datums : bool, default=False
            If True, will ignore datums that already exist in the backend.
            If False, will raise an error if any datums already exist.
            Default is False.
        """
        groundtruths_json = []
        for groundtruth in groundtruths:
            if not isinstance(groundtruth, GroundTruth):
                raise TypeError(
                    f"Expected ground truth to be of type 'valor.GroundTruth' not '{type(groundtruth)}'."
                )
            if not isinstance(groundtruth.annotations._value, list):
                raise TypeError
            groundtruth_dict = groundtruth.encode_value()
            groundtruth_dict["dataset_name"] = dataset.name
            groundtruths_json.append(groundtruth_dict)
        self.conn.create_groundtruths(
            groundtruths_json,
            timeout=timeout,
            ignore_existing_datums=ignore_existing_datums,
        )

    def get_groundtruth(
        self,
        dataset: Union[Dataset, str],
        datum: Union[Datum, str],
    ) -> Union[GroundTruth, None]:
        """
        Get a particular ground truth.

        Parameters
        ----------
        dataset : Union[Dataset, str]
            The dataset the datum belongs to.
        datum : Union[Datum, str]
            The desired datum.

        Returns
        ----------
        Union[GroundTruth, None]
            The matching ground truth or 'None' if it doesn't exist.
        """
        dataset_name = (
            dataset.name if isinstance(dataset, Dataset) else dataset
        )
        datum_uid = datum.uid if isinstance(datum, Datum) else datum
        try:
            resp = self.conn.get_groundtruth(
                dataset_name=dataset_name, datum_uid=datum_uid  # type: ignore
            )
            resp.pop("dataset_name")
            return GroundTruth.decode_value(resp)
        except ClientException as e:
            if e.status_code == 404:
                return None
            raise e

    def finalize_dataset(self, dataset: Union[Dataset, str]) -> None:
        """
        Finalizes a dataset such that new ground truths cannot be added to it.

        Parameters
        ----------
        dataset : str
            The dataset to be finalized.
        """
        dataset_name = (
            dataset.name if isinstance(dataset, Dataset) else dataset
        )
        return self.conn.finalize_dataset(name=dataset_name)  # type: ignore

    def get_dataset(
        self,
        name: str,
    ) -> Union[Dataset, None]:
        """
        Gets a dataset by name.

        Parameters
        ----------
        name : str
            The name of the dataset to fetch.

        Returns
        -------
        Union[Dataset, None]
            A Dataset with a matching name, or 'None' if one doesn't exist.
        """
        dataset = Dataset.decode_value(
            {
                **self.conn.get_dataset(name),
                "connection": self.conn,
            }
        )
        return dataset

    def get_datasets(
        self,
        filters: Optional[Filter] = None,
    ) -> List[Dataset]:
        """
        Get all datasets, with an option to filter results according to some user-defined parameters.

        Parameters
        ----------
        filters : Filter, optional
            Optional constraints to filter by.

        Returns
        ------
        List[valor.Dataset]
            A list of datasets.
        """
        dataset_list = []
        filters = filters if filters is not None else Filter()
        for kwargs in self.conn.get_datasets(filters.to_dict()):
            dataset = Dataset.decode_value({**kwargs, "connection": self.conn})
            dataset_list.append(dataset)
        return dataset_list

    def get_datums(
        self,
        filters: Optional[Filter] = None,
    ) -> List[Datum]:
        """
        Get all datums using an optional filter.

        Parameters
        ----------
        filters : Filter, optional
            Optional constraints to filter by.

        Returns
        -------
        List[valor.Datum]
            A list datums.
        """

        filters = filters if filters is not None else Filter()
        return [
            Datum.decode_value(datum)
            for datum in self.conn.get_datums(filters.to_dict())
        ]

    def get_datum(
        self,
        dataset: Union[Dataset, str],
        uid: str,
    ) -> Union[Datum, None]:
        """
        Get datum.
        `GET` endpoint.

        Parameters
        ----------
        dataset : valor.Dataset
            The dataset the datum belongs to.
        uid : str
            The UID of the datum.

        Returns
        -------
        valor.Datum
            The requested datum or 'None' if it doesn't exist.
        """
        dataset_name = (
            dataset.name if isinstance(dataset, Dataset) else dataset
        )
        resp = self.conn.get_datum(dataset_name=dataset_name, uid=uid)  # type: ignore
        return Datum.decode_value(resp)

    def get_dataset_status(
        self,
        name: str,
    ) -> Union[TableStatus, None]:
        """
        Get the state of a given dataset.

        Parameters
        ----------
        name : str
            The name of the dataset we want to fetch the state of.

        Returns
        ------
        TableStatus | None
            The state of the dataset, or 'None' if the dataset does not exist.
        """
        try:
            return self.conn.get_dataset_status(name)
        except ClientException as e:
            if e.status_code == 404:
                return None
            raise e

    def get_dataset_summary(
        self,
        name: str,
        *_,
        timeout: Optional[float] = None,
    ) -> DatasetSummary:
        """
        Gets the summary of a dataset.

        Parameters
        ----------
        name : str
            The name of the dataset to create a summary for.

        Returns
        -------
        DatasetSummary
            A dataclass containing the dataset summary.
        """
        return DatasetSummary(
            **self.conn.get_dataset_summary(name, timeout=timeout)
        )

    def delete_dataset(self, name: str, timeout: int = 0) -> None:
        """
        Deletes a dataset.

        Parameters
        ----------
        name : str
            The name of the dataset to be deleted.
        timeout : int
            The number of seconds to wait in order to confirm that the dataset was deleted.
        """
        self.conn.delete_dataset(name)
        if timeout:
            for _ in range(timeout):
                try:
                    self.get_dataset(name)
                except DatasetDoesNotExistError:
                    break
                time.sleep(1)
            else:
                raise TimeoutError(
                    "Dataset wasn't deleted within timeout interval"
                )

    def create_model(
        self,
        model: Union[Model, dict],
    ):
        """
        Creates a model.

        Parameters
        ----------
        model : valor.Model
            The model to create.
        """
        if isinstance(model, Model):
            model = model.encode_value()
        self.conn.create_model(model)

    def create_predictions(
        self,
        dataset: Dataset,
        model: Model,
        predictions: List[Prediction],
        timeout: Optional[float] = None,
    ) -> None:
        """
        Creates predictions.

        Parameters
        ----------
        dataset : valor.Dataset
            The dataset that is being operated over.
        model : valor.Model
            The model making the prediction.
        predictions : List[valor.Prediction]
            The predictions to create.
        timeout : float, optional
            The number of seconds the client should wait until raising a timeout.
        """
        predictions_json = []
        for prediction in predictions:
            if not isinstance(prediction, Prediction):
                raise TypeError(
                    f"Expected prediction to be of type 'valor.Prediction' not '{type(prediction)}'."
                )
            if not isinstance(prediction.annotations._value, list):
                raise TypeError
            prediction_dict = prediction.encode_value()
            prediction_dict["dataset_name"] = dataset.name
            prediction_dict["model_name"] = model.name
            predictions_json.append(prediction_dict)
        self.conn.create_predictions(predictions_json, timeout=timeout)

    def get_prediction(
        self,
        dataset: Union[Dataset, str],
        model: Union[Model, str],
        datum: Union[Datum, str],
    ) -> Union[Prediction, None]:
        """
        Get a particular prediction.

        Parameters
        ----------
        dataset : Union[Dataset, str]
            The dataset the datum belongs to.
        model : Union[Model, str]
            The model that made the prediction.
        datum : Union[Datum, str]
            The desired datum.

        Returns
        ----------
        Union[Prediction, None]
            The matching prediction or 'None' if it doesn't exist.
        """
        dataset_name = (
            dataset.name if isinstance(dataset, Dataset) else dataset
        )
        model_name = model.name if isinstance(model, Model) else model
        datum_uid = datum.uid if isinstance(datum, Datum) else datum

        resp = self.conn.get_prediction(
            dataset_name=dataset_name,  # type: ignore
            model_name=model_name,  # type: ignore
            datum_uid=datum_uid,  # type: ignore
        )
        resp.pop("dataset_name")
        resp.pop("model_name")
        return Prediction.decode_value(resp)

    def finalize_inferences(
        self, dataset: Union[Dataset, str], model: Union[Model, str]
    ) -> None:
        """
        Finalizes a model-dataset pairing such that new predictions cannot be added to it.
        """
        dataset_name = (
            dataset.name if isinstance(dataset, Dataset) else dataset
        )
        model_name = model.name if isinstance(model, Model) else model
        return self.conn.finalize_inferences(
            dataset_name=dataset_name,  # type: ignore
            model_name=model_name,  # type: ignore
        )

    def get_model(
        self,
        name: str,
    ) -> Union[Model, None]:
        """
        Gets a model by name.

        Parameters
        ----------
        name : str
            The name of the model to fetch.

        Returns
        -------
        Union[valor.Model, None]
            A Model with matching name or 'None' if one doesn't exist.
        """
        return Model.decode_value(
            {
                **self.conn.get_model(name),
                "connection": self.conn,
            }
        )

    def get_models(
        self,
        filters: Optional[Filter] = None,
    ) -> List[Model]:
        """
        Get all models using an optional filter.

        Parameters
        ----------
        filters : Filter, optional
            Optional constraints to filter by.

        Returns
        ------
        List[valor.Model]
            A list of models.
        """
        model_list = []
        filters = filters if filters is not None else Filter()
        for kwargs in self.conn.get_models(filters.to_dict()):
            model = Model.decode_value({**kwargs, "connection": self.conn})
            model_list.append(model)
        return model_list

    def get_model_status(
        self,
        dataset_name: str,
        model_name: str,
    ) -> Optional[TableStatus]:
        """
        Get the state of a given model over a dataset.

        Parameters
        ----------
        dataset_name : str
            The name of the dataset that the model is operating over.
        model_name : str
            The name of the model we want to fetch the state of.

        Returns
        ------
        Union[TableStatus, None]
            The state of the model or 'None' if the model doesn't exist.
        """
        try:
            return self.conn.get_model_status(dataset_name, model_name)
        except ClientException as e:
            if e.status_code == 404:
                return None
            raise e

    def get_model_eval_requests(
        self, model: Union[Model, str]
    ) -> List[Evaluation]:
        """
        Get all evaluations that have been created for a model.

        This does not return evaluation results.

        `GET` endpoint.

        Parameters
        ----------
        model : str
            The model to search by.

        Returns
        -------
        List[Evaluation]
            A list of evaluations.
        """
        model_name = model.name if isinstance(model, Model) else model
        return [
            Evaluation(**evaluation, connection=self.conn)
            for evaluation in self.conn.get_model_eval_requests(model_name)  # type: ignore
        ]

    def delete_model(self, name: str, timeout: int = 0) -> None:
        """
        Deletes a model.

        Parameters
        ----------
        name : str
            The name of the model to be deleted.
        timeout : int
            The number of seconds to wait in order to confirm that the model was deleted.
        """
        self.conn.delete_model(name)
        if timeout:
            for _ in range(timeout):
                try:
                    self.get_model(name)
                except ModelDoesNotExistError:
                    break
                time.sleep(1)
            else:
                raise TimeoutError(
                    "Model wasn't deleted within timeout interval"
                )

    def get_evaluations(
        self,
        *,
        evaluation_ids: Optional[List[int]] = None,
        models: Union[List[Model], List[str], None] = None,
        datasets: Union[List[Dataset], List[str], None] = None,
        metrics_to_sort_by: Optional[
            Dict[str, Union[Dict[str, str], str]]
        ] = None,
        timeout: Optional[float] = None,
    ) -> List[Evaluation]:
        """
        Returns all evaluations associated with user-supplied dataset and/or model names.

        Parameters
        ----------
        evaluation_ids : List[int], optional.
            A list of job IDs to return metrics for.
        models : Union[List[valor.Model], List[str]], optional
            A list of model names that we want to return metrics for.
        datasets : Union[List[valor.Dataset], List[str]], optional
            A list of dataset names that we want to return metrics for.
        metrics_to_sort_by : dict[str, str | dict[str, str]], optional
            An optional dict of metric types to sort the evaluations by.
        timeout : float, optional
            The number of seconds the client should wait until raising a timeout.

        Returns
        -------
        List[valor.Evaluation]
            A list of evaluations.
        """
        if isinstance(datasets, list):
            datasets = [  # type: ignore
                element.name if isinstance(element, Dataset) else element
                for element in datasets
            ]
        if isinstance(models, list):
            models = [  # type: ignore
                element.name if isinstance(element, Model) else element
                for element in models
            ]
        return [
            Evaluation(connection=self.conn, **evaluation)
            for evaluation in self.conn.get_evaluations(
                evaluation_ids=evaluation_ids,
                models=models,  # type: ignore
                datasets=datasets,  # type: ignore
                metrics_to_sort_by=metrics_to_sort_by,
                timeout=timeout,
            )
        ]

    def evaluate(
        self,
        request: EvaluationRequest,
        *_,
        allow_retries: bool = False,
        timeout: Optional[float] = None,
    ) -> List[Evaluation]:
        """
        Creates as many evaluations as necessary to fulfill the request.

        Parameters
        ----------
        request : schemas.EvaluationRequest
            The requested evaluation parameters.
        allow_retries : bool, default = False
            Option to retry previously failed evaluations.
        timeout : float, optional
            The number of seconds the client should wait until raising a timeout.

        Returns
        -------
        List[Evaluation]
            A list of evaluations that meet the parameters.
        """
        return [
            Evaluation(**evaluation)
            for evaluation in self.conn.evaluate(
                request.to_dict(),
                allow_retries=allow_retries,
                timeout=timeout,
            )
        ]

    def delete_evaluation(self, evaluation_id: int, timeout: int = 0) -> None:
        """
        Deletes an evaluation.

        Parameters
        ----------
        evaluation_id : int
            The id of the evaluation to be deleted.
        timeout : int, default=0
            The number of seconds to wait in order to confirm that the model was deleted.
        """
        self.conn.delete_evaluation(evaluation_id)
        if timeout:
            for _ in range(timeout):
                try:
                    self.get_evaluations(evaluation_ids=[evaluation_id])
                except EvaluationDoesNotExist:
                    break
                time.sleep(1)
            else:
                raise TimeoutError(
                    "Evaluation wasn't deleted within timeout interval"
                )