"""client.py
============
The main client module.
"""

import os
import datetime
import time
from requests.exceptions import ConnectionError
from uuid import uuid4
from io import BytesIO
from google.resumable_media.requests import ResumableUpload
from google.resumable_media.common import InvalidResponse, DataCorruption
from tqdm import tqdm
import sys
from .viridis import viridis_rgb as viridis_rgb
from .turbo import turbo_rgb as turbo_rgb
from .issues import IssueManager as IssueManager
from .datasharing import (
    check_urls as check_urls,
    get_mode as get_mode,
    get_errors as get_errors,
)
import re
from typing import Any, Optional, Union, List, Dict, Tuple, Iterable, Callable
from typing_extensions import TypedDict

from .util import (
    _upload_local_files,
    requests_retry,
    assert_valid_name,
    MAX_CHUNK_SIZE,
    raise_resp_exception_error,
)

# All imports here necessary for referenceing and compatiblity
from .dataset import LabeledDataset as LabeledDataset, LabeledFrame as LabeledFrame
from .inference import Inferences as Inferences, InferencesFrame as InferencesFrame
from .frames import UpdateFrame
from .labels import (
    UpdateGTLabelSet,
    UpdateInferenceLabelSet,
)
from .class_map import (
    LabelClassMap as LabelClassMap,
    ClassMapEntry as ClassMapEntry,
    ClassMapUpdateEntry as ClassMapUpdateEntry,
    tableau_colors as tableau_colors,
    orig_label_color_list as orig_label_color_list,
)


class StratifiedMetricsDict(TypedDict):
    name: str
    ordered_values: List[str]


class StratifiedMetricsDefinition:
    """Definitions for stratified metrics given object-level attributes

    Args:
        name (str): The name of this attribute, which should match the attribute on object labels.
        ordered_values (List[str]): The ordered list of valid values to group by.
    """

    def __init__(self, name: str, ordered_values: Iterable[str]) -> None:
        self.name = name
        self.ordered_values = list(ordered_values)  # In case it's a tuple

    def to_dict(self) -> StratifiedMetricsDict:
        return {"name": self.name, "ordered_values": self.ordered_values}


class CustomMetricsDefinition:
    """Definitions for custom user provided metrics.

    Args:
        name (str): The name of this metric.
        metrics_type (str): The metrics type, either 'objective' or 'confusion_matrix'.
    """

    # TODO: Make these literal types too in the type system?
    OBJECTIVE = "objective"
    CONFUSION_MATRIX = "confusion_matrix"

    def __init__(self, name: str, metrics_type: str) -> None:
        valid_metrics_types = set(["objective", "confusion_matrix"])
        self.name = name
        self.metrics_type = metrics_type

    def to_dict(self) -> Dict[str, str]:
        return {"name": self.name, "metrics_type": self.metrics_type}


class Client:
    """Client class that interacts with the Aquarium REST API.

    Args:
        api_endpoint (str, optional): The API endpoint to hit. Defaults to "https://illume.aquariumlearning.com/api/v1".
    """

    _creds_token: Optional[str]
    _creds_app_id: Optional[str]
    _creds_app_key: Optional[str]
    _creds_api_key: Optional[str]
    api_endpoint: str

    def __init__(
        self,
        *,
        api_endpoint: str = "https://illume.aquariumlearning.com/api/v1",
        **kwargs: Any,
    ) -> None:
        self._creds_token = None
        self._creds_app_id = None
        self._creds_app_key = None
        self._creds_api_key = None
        self.api_endpoint = api_endpoint

    def _get_creds_headers(self) -> Dict[str, str]:
        """Get appropriate request headers for the currently set credentials.

        Raises:
            Exception: No credentials set.

        Returns:
            dict: Dictionary of headers
        """
        if self._creds_token:
            return {"Authorization": "Bearer {token}".format(token=self._creds_token)}
        elif self._creds_api_key:
            return {"x-illume-api-key": self._creds_api_key}
        elif self._creds_app_id and self._creds_app_key:
            return {
                "x-illume-app": self._creds_app_id,
                "x-illume-key": self._creds_app_key,
            }
        else:
            raise Exception("No credentials set.")

    def set_credentials(
        self,
        *,
        token: Optional[str] = None,
        app_id: Optional[str] = None,
        app_key: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> None:
        """Set credentials for the client.

        Args:
            api_key (str, optional): A string for a long lived API key. Defaults to None.
            token (str, optional): A JWT providing auth credentials. Defaults to None.
            app_id (str, optional): Application ID string. Defaults to None.
            app_key (str, optional): Application secret key. Defaults to None.

        Raises:
            Exception: Invalid credential combination provided.
        """
        if api_key is not None:
            self._creds_api_key = api_key
        elif token is not None:
            self._creds_token = token
        elif app_id is not None and app_key is not None:
            self._creds_app_id = app_id
            self._creds_app_key = app_key
        else:
            raise Exception(
                "Please provide either an api_key, token, or app_id and app_key"
            )

    def _format_error_logs(self, raw_error_logs: List[Dict[str, str]]) -> List[str]:
        """Format error log data into strings.

        Args:
            raw_error_logs (list[dict]): Error log data.

        Returns:
            list[str]: list of string formatted error messages.
        """
        formatted_lines = []
        for raw in raw_error_logs:
            formatted_lines.append(
                f"    {raw.get('aquarium_dataflow_step', '')}: {raw.get('msg', '')}"
            )
        return formatted_lines

    def get_issue_manager(self, project_id: str) -> IssueManager:
        """Get an issue manager object.

        Args:
            project_id (str): Project ID to manage.

        Returns:
            IssueManager: The issue manager object.
        """
        return IssueManager(self, project_id)

    def get_projects(self) -> List[Dict[str, Any]]:
        """Get info about existing projects

        Returns:
            list of dict: Project Info
        """
        r = requests_retry.get(
            self.api_endpoint + "/projects", headers=self._get_creds_headers()
        )

        raise_resp_exception_error(r)
        result: List[Dict[str, Any]] = r.json()
        return result

    def get_project(self, project_id: str) -> Dict[str, Any]:
        """Get detailed info about a specific project

        Returns:
            Dict[str, Any]: detailed info about a project
        """
        r = requests_retry.get(
            self.api_endpoint + "/projects/" + project_id,
            headers=self._get_creds_headers(),
        )

        raise_resp_exception_error(r)
        result: Dict[str, Any] = r.json()
        return result

    def delete_project(self, project_id: str) -> None:
        """Mark a project for deletion

        Args:
            project_id (str): project_id
        """
        if not self.project_exists(project_id):
            raise Exception("Project {} does not exist.".format(project_id))

        url = self.api_endpoint + "/projects/{}".format(project_id)
        r = requests_retry.delete(url, headers=self._get_creds_headers())

        raise_resp_exception_error(r)

    def project_exists(self, project_id: str) -> bool:
        """Checks whether a project exists.

        Args:
            project_id (str): project_id

        Returns:
            bool: Does project exist
        """
        projects = self.get_projects()
        existing_project_ids = [project["id"] for project in projects]
        return project_id in existing_project_ids

    def create_project(
        self,
        project_id: str,
        label_class_map: LabelClassMap,
        primary_task: Optional[str] = None,
        secondary_labels: Optional[Any] = None,
        frame_links: Optional[List[str]] = None,
        label_links: Optional[List[str]] = None,
        default_camera_target: Optional[List[float]] = None,
        default_camera_position: Optional[List[float]] = None,
        custom_metrics: Optional[
            Union[CustomMetricsDefinition, List[CustomMetricsDefinition]]
        ] = None,
        max_shown_categories: Optional[int] = None,
        stratified_metrics: Optional[List[StratifiedMetricsDefinition]] = None,
        include_no_gt: Optional[bool] = None,
        metrics_confidence_threshold: Optional[float] = None,
        metrics_iou_threshold: Optional[float] = None,
        external_metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Create a new project via the REST API.

        Args:
            project_id (str): project_id
            label_class_map (LabelClassMap): The label class map used to interpret classifications.
            primary_task (Optional[str], optional): Any specific primary task for a non-object detection or classification task. Can be '2D_SEMSEG' or 'CLASSIFICATION' or None.
            secondary_labels ([type], optional): List of secondary labels in classification tasks
            frame_links (Optional[List[str]], optional): List of string keys for links between frames
            label_links (Optional[List[str]], optional): List of string keys for links between labels
            default_camera_target (Optional[List[float]], optional): For 3D scenes, the default camera target
            default_camera_position (Optional[List[float]], optional): For 3D scenes, the default camera position
            custom_metrics (Optional[ Union[CustomMetricsDefinition, List[CustomMetricsDefinition]] ], optional): Defines which custom metrics exist for this project, defaults to None.
            max_shown_categories (Optional[int], optional): For categorical visualizations, set the maximum shown simultaneously. Max 100.
            stratified_metrics (Optional[List[StratifiedMetricsDefinition]], optional): Defines what object-level attributes to stratify metrics over.
            metrics_confidence_threshold (Optional[float], optional): In order to calculate metrics + confusion matrices, Aquarium uses this threshold (in combination with IOU) to match your ground truth (GT) labels with your inference labels. Defaults to 0.1 if not specified.
            metrics_iou_threshold(Optional[float], optional): In order to calculate metrics + confusion matrices, Aquarium uses this threshold (in combination with confidence) to match your ground truth (GT) labels with your inference labels. Defaults to 0.5 if not specified.
            external_metadata (Optional[Dict[str, Any]], optional): A JSON object that can be used to attach metadata to the project itself
        """

        assert_valid_name(project_id)

        if not isinstance(label_class_map, LabelClassMap):
            raise Exception("label_class_map must be a LabelClassMap")

        if not label_class_map.entries:
            raise Exception("label_class_map must have at least one class")

        dumped_classmap = [x.to_dict() for x in label_class_map.entries]
        payload = {"project_id": project_id, "label_class_map": dumped_classmap}

        if primary_task is not None:
            payload["primary_task"] = primary_task
        if secondary_labels is not None:
            dumped_secondary_labels = []
            for raw in secondary_labels:
                dumped_classmap = [x.to_dict() for x in raw["label_class_map"].entries]
                raw["label_class_map"] = dumped_classmap
                dumped_secondary_labels.append(raw)

            payload["secondary_labels"] = dumped_secondary_labels
        if frame_links is not None:
            if not isinstance(frame_links, list):
                raise Exception("frame_links must be a list of strings")
            payload["frame_links"] = frame_links
        if label_links is not None:
            if not isinstance(label_links, list):
                raise Exception("label_links must be a list of strings")
            payload["label_links"] = label_links
        if default_camera_position is not None:
            if not isinstance(default_camera_position, list):
                raise Exception("default_camera_position must be a list of floats")
            payload["default_camera_position"] = default_camera_position
        if default_camera_target is not None:
            if not isinstance(default_camera_target, list):
                raise Exception("default_camera_target must be a list of floats")
            payload["default_camera_target"] = default_camera_target
        if custom_metrics is not None:
            if isinstance(custom_metrics, CustomMetricsDefinition):
                custom_metrics = [custom_metrics]

            if (
                not custom_metrics
                or (not isinstance(custom_metrics, list))
                or (not isinstance(custom_metrics[0], CustomMetricsDefinition))
            ):
                raise Exception(
                    "custom_metrics must be a CustomMetricsDefinition or list of CustomMetricsDefinition."
                )

            serializable_custom = [x.to_dict() for x in custom_metrics]
            payload["custom_metrics"] = serializable_custom

        if stratified_metrics is not None:
            if (
                not stratified_metrics
                or (not isinstance(stratified_metrics, list))
                or (not isinstance(stratified_metrics[0], StratifiedMetricsDefinition))
            ):
                raise Exception(
                    "stratified_metrics must be a list of StratifiedMetricsDefinition."
                )
            serializable_strat = [x.to_dict() for x in stratified_metrics]
            payload["stratified_metrics"] = serializable_strat

        if max_shown_categories is not None:
            if not isinstance(max_shown_categories, int):
                raise Exception("max_shown_categories must be an int")
            if max_shown_categories < 1 or max_shown_categories > 100:
                raise Exception("max_shown_categories must be between 1 and 100")
            payload["max_shown_categories"] = max_shown_categories

        if include_no_gt is not None:
            payload["include_no_gt"] = include_no_gt

        if metrics_confidence_threshold is not None:
            if not isinstance(metrics_confidence_threshold, float):
                raise Exception("metrics_confidence_threshold must be a float")
            payload["confidence_threshold"] = metrics_confidence_threshold

        if metrics_iou_threshold is not None:
            if not isinstance(metrics_iou_threshold, float):
                raise Exception("metrics_iou_threshold must be a float")
            payload["iou_threshold"] = metrics_iou_threshold

        if external_metadata is not None:
            if not isinstance(external_metadata, dict) or (
                external_metadata and not isinstance(next(iter(external_metadata)), str)
            ):
                raise Exception("external_metadata must be a dict with string keys")
            payload["external_metadata"] = external_metadata

        r = requests_retry.post(
            self.api_endpoint + "/projects",
            headers=self._get_creds_headers(),
            json=payload,
        )
        raise_resp_exception_error(r)

    def _preview_frame_dict(
        self, project_id: str, both_frames_dict: Dict[str, Optional[Dict[str, Any]]]
    ) -> None:
        """Generate preview with both dataset frame and inference frame as dict

        Args:
            project_id (str): name of project to preview frame with
            both_frames_dict (dict): Dictionary containing the labeled and inference frame
        """
        api_path = "/projects/{}/preview_frame".format(project_id)

        preview_frame_api_root = self.api_endpoint + api_path

        r = requests_retry.post(
            preview_frame_api_root,
            headers=self._get_creds_headers(),
            json=both_frames_dict,
        )
        response_data = r.json()
        if response_data.get("preview_frame_uuid"):
            print("Please visit the following url to preview your frame in the webapp")
            url = (
                self.api_endpoint[:-7]
                + api_path
                + "/"
                + response_data["preview_frame_uuid"]
            )
            print(url)
        else:
            raise Exception(
                "Preview URL could not be constructed by server. "
                "Please make sure you're logged in and check frame data accordingly."
            )

    def preview_frame(
        self,
        project_id: str,
        labeled_frame: LabeledFrame,
        inference_frame: Optional[InferencesFrame] = None,
    ) -> None:
        """prints out a URL that lets you preview a provided frame in the web browser
        Useful for debugging data and image url issues.

        Args:
            project_id (str): Name of project to be associated for this frame preview (for label class association)
            labeled_frame (LabeledFrame): Labeled Frame desired for preview in web-app
            inference_frame (Optional[InferencesFrame], optional): Labeled Inference Desired for preview in web-app. Defaults to None.
        """

        both_frames: Dict[str, Optional[Dict[str, Any]]] = {}
        both_frames["labeled_frame"] = labeled_frame.to_dict()
        both_frames["inference_frame"] = (
            inference_frame.to_dict() if inference_frame else None
        )
        self._preview_frame_dict(project_id, both_frames)

    def update_project_metadata(
        self, project_id: str, external_metadata: Dict[str, Any]
    ) -> None:
        """Update project metadata

        Args:
            project_id (str): The project id.
            external_metadata (Dict[Any, Any]): The new metadata
        """
        if not isinstance(external_metadata, dict) or (
            external_metadata and not isinstance(next(iter(external_metadata)), str)
        ):
            raise Exception("external_metadata must be a dict with string keys")

        payload = {"external_metadata": external_metadata}
        r = requests_retry.post(
            f"{self.api_endpoint}/projects/{project_id}/metadata",
            headers=self._get_creds_headers(),
            json=payload,
        )
        raise_resp_exception_error(r)

    def update_label_class_map_colors(
        self, project_id: str, changed_label_classes: List[ClassMapUpdateEntry]
    ) -> None:
        """Updates label class colors of a specific project

        Args:
            project_id (str): The project id.
            changed_label_classes (List[ClassMapUpdateEntry]): The list of label classes with changed colors. Must be a subset of the project's overall label class map
        """
        project = self.get_project(project_id)
        label_class_map_dict_by_id = {
            label_class["id"]: label_class for label_class in project["label_class_map"]
        }
        label_class_map_dict_by_name = {
            label_class["name"]: label_class
            for label_class in project["label_class_map"]
        }

        seen_ids = set()
        dumped_changes = []
        for entry in changed_label_classes:
            if not entry.class_id:
                known_label_class = label_class_map_dict_by_name.get(entry.name)
                if not known_label_class:
                    raise Exception(
                        f"Label class with name {entry.name} could not be found in the project's label class map. "
                        "This method only allows changing an existing label class map. "
                        "To append new label classes please use create_project."
                    )
                entry.class_id = known_label_class["id"]

            known_label_class = label_class_map_dict_by_id.get(entry.class_id)

            if not known_label_class:
                raise Exception(
                    f"Label class with id {entry.class_id} could not be found in the project's label class map. "
                    "This method only allows changing an existing label class map. "
                    "To append new label classes please use create_project."
                )

            if entry.class_id in seen_ids:
                raise Exception(
                    f"Label class with id {entry.class_id} ({known_label_class['name']}) has multiple change entries. "
                    "Please consolidate into one change."
                )

            seen_ids.add(entry.class_id)
            dumped_entry = entry.to_dict()

            # lazy shallow None check to avoid overwriting with a null
            dumped_patch_entry = {
                k: v for k, v in dumped_entry.items() if v is not None
            }

            dumped_changes.append(dumped_patch_entry)

        payload = {"label_class_map": dumped_changes}
        r = requests_retry.patch(
            f"{self.api_endpoint}/projects/{project_id}/label_class_map",
            headers=self._get_creds_headers(),
            json=payload,
        )
        raise_resp_exception_error(r)

    def get_datasets(
        self, project_id: str, include_archived: Optional[bool]
    ) -> List[Dict[str, Any]]:
        """Get existing datasets for a project.

        Args:
            project_id (str): The project id.

        Returns:
            list: A list of dataset info for the project.
        """
        datasets_api_root = self.api_endpoint + "/projects/{}/datasets".format(
            project_id
        )
        params = {"include_archived": include_archived} if include_archived else {}
        r = requests_retry.get(
            datasets_api_root, headers=self._get_creds_headers(), params=params
        )
        raise_resp_exception_error(r)
        result: List[Dict[str, Any]] = r.json()
        return result

    def delete_dataset(self, project_id: str, dataset_id: str) -> None:
        """Mark a dataset for deletion

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
        """
        if not self.dataset_exists(project_id, dataset_id):
            raise Exception("Dataset {} does not exist.".format(dataset_id))

        url = self.api_endpoint + "/projects/{}/datasets/{}".format(
            project_id, dataset_id
        )
        r = requests_retry.delete(url, headers=self._get_creds_headers())

        raise_resp_exception_error(r)

    def dataset_exists(self, project_id: str, dataset_id: str) -> bool:
        """Check if a dataset exists.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id

        Returns:
            bool: Whether the dataset already exists.
        """
        datasets = self.get_datasets(project_id, include_archived=True)
        existing_dataset_ids = [dataset.get("id") for dataset in datasets]
        return dataset_id in existing_dataset_ids

    def is_dataset_processed(self, project_id: str, dataset_id: str) -> bool:
        """Check if a dataset is fully processed.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id

        Returns:
            bool: If the dataset is done processing.
        """
        endpoint_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/is_processed".format(project_id, dataset_id)
        )

        r = requests_retry.get(endpoint_path, headers=self._get_creds_headers())
        raise_resp_exception_error(r)
        parsed: Dict[str, bool] = r.json()
        return parsed["processed"]

    def get_dataset_ingest_error_logs(
        self, project_id: str, dataset_id: str
    ) -> List[Dict[str, Any]]:
        """Get ingest error log entries for a dataset.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id

        Returns:
            list[dict]: List of error entries
        """
        endpoint_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/ingest_error_logs".format(
                project_id, dataset_id
            )
        )

        r = requests_retry.get(endpoint_path, headers=self._get_creds_headers())
        raise_resp_exception_error(r)
        parsed: List[Dict[str, Any]] = r.json()
        return parsed

    def current_dataset_process_state(
        self, project_id: str, dataset_id: str
    ) -> Tuple[str, float]:
        """Current processing state of a dataset.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id

        Returns:
            Tuple[str, float]: semantic name of state of processing, percent done of job
        """
        endpoint_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/process_state".format(project_id, dataset_id)
        )

        r = requests_retry.get(endpoint_path, headers=self._get_creds_headers())
        raise_resp_exception_error(r)
        parsed = r.json()
        return parsed["current_state"], parsed["percent_done"]

    def current_abstract_dataset_process_step_status(
        self, project_id: str, dataset_id: str
    ) -> Dict[str, Any]:
        """Returns the process steps statuses for a given dataset or inferenceset

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id

        Returns:
            Dict[str, Any]: A set of process step statuses that exist for a given abstract dataset
        """
        endpoint_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/process_step_status".format(
                project_id, dataset_id
            )
        )

        r = requests_retry.get(endpoint_path, headers=self._get_creds_headers())
        raise_resp_exception_error(r)
        parsed: Dict[str, Any] = r.json()
        return parsed

    @staticmethod
    def parse_normalize_process_step_status(
        process_step_status_payload: Dict[str, Any]
    ) -> str:
        result: str = process_step_status_payload["process_step_statuses"]["normalize"][
            "status"
        ]
        return result

    def inferences_exists(
        self, project_id: str, dataset_id: str, inferences_id: str
    ) -> bool:
        """Check if a set of inferences exists.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            inferences_id (str): inferences_id

        Returns:
            bool: Whether the inferences id already exists.
        """
        # TODO: FIXME: We need a first class model for inferences,
        # not just name gluing
        inferences_dataset_id = "_".join(["inferences", dataset_id, inferences_id])
        datasets = self.get_datasets(project_id, include_archived=True)
        existing_dataset_ids = [dataset.get("id") for dataset in datasets]
        return inferences_dataset_id in existing_dataset_ids

    def is_inferences_processed(
        self, project_id: str, dataset_id: str, inferences_id: str
    ) -> bool:
        """Check if a set of inferences is fully processed.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            inferences_id(str): inferences_id

        Returns:
            bool: If the inferences set is done processing.
        """
        endpoint_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/inferences/{}/is_processed".format(
                project_id, dataset_id, inferences_id
            )
        )

        r = requests_retry.get(endpoint_path, headers=self._get_creds_headers())
        raise_resp_exception_error(r)
        parsed = r.json()
        processed: bool = parsed["processed"]
        return processed

    def get_inferences_ingest_error_logs(
        self, project_id: str, dataset_id: str, inferences_id: str
    ) -> List[Dict[str, Any]]:
        """Get ingest error log entries for an inference set.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            inferences_id(str): inferences_id

        Returns:
            list[dict]: List of error entries
        """
        endpoint_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/inferences/{}/ingest_error_logs".format(
                project_id, dataset_id, inferences_id
            )
        )

        r = requests_retry.get(endpoint_path, headers=self._get_creds_headers())
        raise_resp_exception_error(r)
        parsed: List[Dict[str, Any]] = r.json()
        return parsed

    def current_inferences_process_state(
        self, project_id: str, dataset_id: str, inferences_id: str
    ) -> Tuple[str, float]:
        """current processing state of inferences.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            inferences_id(str): inferences_id

        Returns:
            Tuple[str, float]: semantic name of state of processing, percent done of job
        """
        endpoint_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/inferences/{}/process_state".format(
                project_id, dataset_id, inferences_id
            )
        )

        r = requests_retry.get(endpoint_path, headers=self._get_creds_headers())
        raise_resp_exception_error(r)
        parsed = r.json()
        return parsed["current_state"], parsed["percent_done"]

    def upload_asset_from_filepath(
        self, project_id: str, dataset_id: str, filepath: str
    ) -> str:
        """Upload an asset from a local file path.
        This is useful in cases where you have data on your local machine that you want to mirror in aquarium.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            filepath (str): The filepath to grab the assset data from

        Returns:
            str: The URL to the mirrored asset.
        """

        get_upload_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/get_upload_url".format(project_id, dataset_id)
        )

        upload_filename = os.path.basename(filepath)
        upload_filename = "{}_{}".format(str(uuid4()), upload_filename)

        params = {"upload_filename": upload_filename}
        upload_url_resp = requests_retry.get(
            get_upload_path, headers=self._get_creds_headers(), params=params
        )

        raise_resp_exception_error(upload_url_resp)
        urls = upload_url_resp.json()
        put_url = urls["put_url"]
        download_url: str = urls["download_url"]

        with open(filepath, "rb") as f:
            upload_resp = requests_retry.put(put_url, data=f)

        raise_resp_exception_error(upload_resp)
        return download_url

    def upload_asset_from_url(
        self, project_id: str, dataset_id: str, source_url: str
    ) -> str:
        """Upload an asset from a private url.
        This is useful in cases where you have data easily accessible on your network that you want to mirror in aquarium.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            source_url (str): The source url to grab the assset data from

        Returns:
            str: The URL to the mirrored asset.
        """
        get_upload_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/get_upload_url".format(project_id, dataset_id)
        )

        upload_filename = os.path.basename(source_url)
        upload_filename = "{}_{}".format(str(uuid4()), upload_filename)

        params = {"upload_filename": upload_filename}
        upload_url_resp = requests_retry.get(
            get_upload_path, headers=self._get_creds_headers(), params=params
        )

        raise_resp_exception_error(upload_url_resp)
        urls = upload_url_resp.json()
        put_url = urls["put_url"]
        download_url: str = urls["download_url"]

        dl_resp = requests_retry.get(source_url)
        payload = BytesIO(dl_resp.content)

        upload_resp = requests_retry.put(put_url, data=payload)

        raise_resp_exception_error(upload_resp)
        return download_url

    def _upload_rows_from_files(
        self,
        project_id: str,
        dataset_id: str,
        upload_prefix: str,
        upload_suffix: str,
        file_names: List[str],
        delete_after_upload: bool = True,
    ) -> List[str]:
        # Get upload / download URLs
        get_upload_path = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/get_upload_url".format(project_id, dataset_id)
        )

        download_urls = _upload_local_files(
            file_names,
            get_upload_path,
            self._get_creds_headers(),
            upload_prefix,
            upload_suffix,
            delete_after_upload=delete_after_upload,
        )

        return download_urls

    def add_new_frames(
        self,
        project_id: str,
        dataset_id: str,
        frames: List[LabeledFrame],
        batch_name: Optional[str] = None,
    ) -> None:
        raise Exception("Not Implemented")

    def update_dataset_frames(
        self,
        project_id: str,
        dataset_id: str,
        update: List[UpdateFrame],
        batch_name: Optional[str] = None,
    ) -> None:
        raise Exception("Not Implemented")

    def update_dataset_labels(
        self,
        project_id: str,
        dataset_id: str,
        update: List[UpdateGTLabelSet],
        batch_name: Optional[str] = None,
    ) -> None:
        raise Exception("Not Implemented")

    def update_inference_labels(
        self,
        project_id: str,
        dataset_id: str,
        update: List[UpdateInferenceLabelSet],
        batch_name: Optional[str] = None,
    ) -> None:
        raise Exception("Not Implemented")

    def create_dataset(
        self,
        project_id: str,
        dataset_id: str,
        data_url: Optional[str] = None,
        embeddings_url: Optional[str] = None,
        dataset: Optional[LabeledDataset] = None,
        wait_until_finish: bool = False,
        wait_timeout: datetime.timedelta = datetime.timedelta(hours=2),
        embedding_distance_metric: str = "euclidean",
        preview_first_frame: bool = False,
        delete_cache_files_after_upload: bool = True,
        check_first_frame: bool = True,
        external_metadata: Optional[Dict[str, Any]] = None,
        batch_name: Optional[str] = None,
    ) -> None:
        """Create a dataset with the provided data urls.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            data_url (Optional[str], optional): A URL to the serialized dataset entries.
            embeddings_url (Optional[str], optional): A URL to the serialized dataset embeddings. Defaults to None.
            dataset (Optional[LabeledDataset], optional): The LabeledDataset to upload.
            wait_until_finish (bool, optional): Block until the dataset processing job finishes. This generally takes at least 5 minutes, and scales with the size of the dataset. Defaults to False.
            wait_timeout (datetime.timedelta, optional): Maximum time to wait for. Defaults to 2 hours.
            embedding_distance_metric (str, optional): Distance metric to use for embedding layout. Can be a member of ['euclidean', 'cosine']. Defaults to 'euclidean'.
            preview_first_frame (bool, optional): preview the first frame of the dataset in the webapp before continuing. Requires interaction.
            delete_cache_files_after_upload (bool, optional): flag to turn off automatic deletion of cache files after upload. Useful for ipython notebook users that reload/re-attempt uploads. Defaults to True.
            external_metadata (Optional[Dict[str, Any]], optional): A JSON object that can be used to attach metadata to the dataset itself
        """

        assert_valid_name(dataset_id)

        if embedding_distance_metric not in ["euclidean", "cosine"]:
            raise Exception("embedding_distance_metric must be euclidean or cosine.")

        if not isinstance(wait_timeout, datetime.timedelta):
            raise Exception("wait_timeout must be a datetime.timedelta object")

        if self.dataset_exists(project_id, dataset_id):
            raise Exception("Dataset already exists.")

        if dataset:
            dataset._flush_to_disk()

            project_info = self.get_project(project_id)

            dataset._validate_frames(project_info)

            if len(dataset._temp_frame_file_names) == 0:
                raise Exception("Cannot create dataset with 0 frames")

            if preview_first_frame:
                first_frame_dict = dataset.get_first_frame_dict()
                self._preview_frame_dict(
                    project_id,
                    {"labeled_frame": first_frame_dict, "inference_frame": None},
                )
                user_response = input(
                    "Please vist above URL to see your Preview frame.\n\n"
                    "Press ENTER to continue or type `exit` and press ENTER "
                    "to cancel dataset upload.\n"
                )
                if user_response == "exit":
                    print("Canceling dataset upload!")
                    return

            if check_first_frame:
                first_frame_dict = dataset.get_first_frame_dict()
                sensor_data = first_frame_dict.get("sensor_data")
                if sensor_data is None:
                    raise Exception("First frame contains zero sensor/media inputs.")

                for sensor in sensor_data:
                    urls = list(sensor.get("data_urls").values())
                    local_results = check_urls(urls, None)
                    r = requests_retry.post(
                        self.api_endpoint + "/datasets/verify_urls",
                        headers=self._get_creds_headers(),
                        json={"urls": urls},
                    )
                    server_results = r.json()
                    mode = get_mode(urls, local_results, server_results)
                    errors = get_errors(mode, dataset)

                    if len(errors) > 0:
                        print(f"WARNING! SOME URL ACCESS ERRORS FOUND!")
                        print(f"Errors found: {errors}")
                    else:
                        print("No errors found when checking first frame!")

            project_info = self.get_project(project_id)
            project_label_set = set(
                [label_class["name"] for label_class in project_info["label_class_map"]]
            )
            extra_dataset_labels = dataset._label_classes_set - project_label_set
            if len(extra_dataset_labels) > 0:
                raise Exception(
                    f"Dataset contains labels {extra_dataset_labels} not in project Label Class Map {project_label_set}. "
                    f"Label Class Maps can be updated using append-only logic as specified here. "
                    f"https://aquarium.gitbook.io/aquarium-changelog/changelog-entries/2021-01-22#updating-label-class-maps"
                )

            print("Uploading Labeled Dataset...")
            upload_prefix = "{}_data".format(str(uuid4()))
            upload_suffix = ".jsonl"
            final_urls = self._upload_rows_from_files(
                project_id,
                dataset_id,
                upload_prefix,
                upload_suffix,
                dataset._temp_frame_file_names,
                delete_after_upload=delete_cache_files_after_upload,
            )
        elif data_url:
            final_urls = [
                data_url,
            ]
        else:
            raise Exception("Please provide either a data_url or dataset argument")

        datasets_api_root = self.api_endpoint + "/projects/{}/datasets".format(
            project_id
        )
        payload = {
            "dataset_id": dataset_id,
            "data_url": final_urls,
            "embedding_distance_metric": embedding_distance_metric,
            "embedding_upload_version": 1,
        }

        if external_metadata is not None:
            if not isinstance(external_metadata, dict) or (
                external_metadata and not isinstance(next(iter(external_metadata)), str)
            ):
                raise Exception("external_metadata must be a dict with string keys")
            payload["external_metadata"] = external_metadata

        if dataset and dataset._temp_frame_embeddings_file_names:
            print("Uploading Labeled Dataset Embeddings...")
            upload_prefix = "{}_embeddings".format(str(uuid4()))
            upload_suffix = ".arrow"
            uploaded_urls = self._upload_rows_from_files(
                project_id,
                dataset_id,
                upload_prefix,
                upload_suffix,
                dataset._temp_frame_embeddings_file_names,
                delete_after_upload=delete_cache_files_after_upload,
            )
            if uploaded_urls:  # not empty list
                payload["embeddings_url"] = uploaded_urls
        elif embeddings_url:
            payload["embeddings_url"] = [
                embeddings_url,
            ]

        # TODO: Should we just make it not optional, because very few (if any?)
        # people are still using the old data url flow
        if dataset is not None:
            dataset._cleanup_temp_dir()

        print("Dataset Processing Initiating...")
        r = requests_retry.post(
            datasets_api_root, headers=self._get_creds_headers(), json=payload
        )
        raise_resp_exception_error(r)
        print("Dataset Processing Initiated Successfully")

        if wait_until_finish:
            print(f"Dataset Processing is waiting on workers to spin up...")
            normalize_status = "PENDING"
            while normalize_status == "PENDING":
                time.sleep(10)
                normalize_status = Client.parse_normalize_process_step_status(
                    self.current_abstract_dataset_process_step_status(
                        project_id, dataset_id
                    )
                )
                full_process_state, _ = self.current_dataset_process_state(
                    project_id, dataset_id
                )
                if full_process_state == "FAILED":
                    print("Dataset processing has failed. Exiting...")
                    return
            print(f"Dataset Processing Workers have been spun up.")
            with tqdm(
                total=100.0,
                file=sys.stdout,
                unit_scale=True,
                desc="Dataset Processing Progress",
            ) as pbar:  # type: tqdm[Any]
                start_time = datetime.datetime.now()
                processing_state = "PENDING"
                display_processing_state: Callable[
                    [str], str
                ] = lambda state: f"Dataset Processing Status: {state}"
                pbar.write(display_processing_state(processing_state))
                while (datetime.datetime.now() - start_time) < wait_timeout:
                    (
                        new_processing_state,
                        new_percent_done,
                    ) = self.current_dataset_process_state(project_id, dataset_id)
                    pbar.update(new_percent_done - pbar.n)
                    if new_processing_state != processing_state:
                        processing_state = new_processing_state
                        pbar.write(display_processing_state(processing_state))
                    processed = self.is_dataset_processed(
                        project_id=project_id, dataset_id=dataset_id
                    )
                    if processed:
                        pbar.update(100.0 - pbar.n)
                        pbar.close()
                        print("Dataset is fully processed.")
                        break
                    if processing_state == "FAILED":
                        pbar.update(100.0 - pbar.n)
                        pbar.close()
                        print("Dataset processing has failed. Exiting...")
                        raw_logs = self.get_dataset_ingest_error_logs(
                            project_id, dataset_id
                        )
                        formatted = self._format_error_logs(raw_logs)
                        for entry in formatted:
                            print(entry)

                        break
                    else:
                        time.sleep(10)

                if datetime.datetime.now() - start_time >= wait_timeout:
                    pbar.close()
                    print("Exceeded timeout waiting for job completion.")

    def update_dataset_metadata(
        self, project_id: str, dataset_id: str, external_metadata: Dict[str, Any]
    ) -> None:
        """Update dataset metadata

        Args:
            project_id (str): The project id.
            dataset_id (str): The dataset id.
            external_metadata (Dict[Any, Any]): The new metadata
        """
        if not isinstance(external_metadata, dict) or (
            external_metadata and not isinstance(next(iter(external_metadata)), str)
        ):
            raise Exception("external_metadata must be a dict with string keys")

        payload = {"external_metadata": external_metadata}
        r = requests_retry.post(
            f"{self.api_endpoint}/projects/{project_id}/datasets/{dataset_id}/metadata",
            headers=self._get_creds_headers(),
            json=payload,
        )
        raise_resp_exception_error(r)

    # TODO: Check if the inferences already exist
    def create_inferences(
        self,
        project_id: str,
        dataset_id: str,
        inferences_id: str,
        data_url: Optional[str] = None,
        embeddings_url: Optional[str] = None,
        inferences: Optional[Inferences] = None,
        wait_until_finish: bool = False,
        wait_timeout: datetime.timedelta = datetime.timedelta(hours=2),
        embedding_distance_metric: str = "euclidean",
        delete_cache_files_after_upload: bool = True,
        external_metadata: Optional[Dict[str, Any]] = None,
        batch_name: Optional[str] = None,
    ) -> None:
        """Create an inference set with the provided data urls.

        Args:
            project_id (str): project_id
            dataset_id (str): dataset_id
            inferences_id (str): A unique identifier for this set of inferences.
            data_url (Optional[str], optional): A URL to the serialized inference entries.
            embeddings_url (Optional[str], optional): A URL to the serialized inference embeddings. Defaults to None.
            inferences (Optional[Inferences], optional): The inferences to upload.
            wait_until_finish (bool, optional): Block until the dataset processing job finishes. This generally takes at least 5 minutes, and scales with the size of the dataset. Defaults to False.
            wait_timeout (datetime.timedelta, optional): Maximum time to wait for. Defaults to 2 hours.
            embedding_distance_metric (str, optional): Distance metric to use for embedding layout. Can be a member of ['euclidean', 'cosine']. Defaults to 'euclidean'.
            delete_cache_files_after_upload (bool, optional): flag to turn off automatic deletion of cache files after upload. Useful for ipython notebook users that reload/re-attempt uploads. Defaults to True.
            external_metadata (Optional[Dict[str, Any]], optional): A JSON object that can be used to attach metadata to the inferences itself
        """

        assert_valid_name(dataset_id)
        assert_valid_name(inferences_id)
        queue_after_dataset = True

        if embedding_distance_metric not in ["euclidean", "cosine"]:
            raise Exception("embedding_distance_metric must be euclidean or cosine.")

        if not isinstance(wait_timeout, datetime.timedelta):
            raise Exception("wait_timeout must be a datetime.timedelta object")

        if not self.dataset_exists(project_id, dataset_id):
            raise Exception(f"Dataset {dataset_id} does not exist")
        if (
            not self.is_dataset_processed(project_id, dataset_id)
            and not queue_after_dataset
        ):
            raise Exception(f"Dataset {dataset_id} is not fully processed")
        # If the dataset is already processed, ignore the queue arg
        if self.is_dataset_processed(project_id, dataset_id) and queue_after_dataset:
            queue_after_dataset = False

        if self.inferences_exists(project_id, dataset_id, inferences_id):
            raise Exception("Inferences already exists.")

        inferences_api_root = (
            self.api_endpoint
            + "/projects/{}/datasets/{}/inferences".format(project_id, dataset_id)
        )

        if inferences:
            inferences._flush_to_disk()

            project_info = self.get_project(project_id)
            project_label_set = set(
                [label_class["name"] for label_class in project_info["label_class_map"]]
            )
            extra_inferences_labels = inferences._label_classes_set - project_label_set
            if len(extra_inferences_labels) > 0:
                raise Exception(
                    f"Dataset contains labels {extra_inferences_labels} not in project Label Class Map {project_label_set}"
                )

            inferences._validate_frames(project_info)
            print("Uploading Inferences...")

            upload_prefix = "{}_data".format(str(uuid4()))
            upload_suffix = ".jsonl"
            final_urls = self._upload_rows_from_files(
                project_id,
                dataset_id,
                upload_prefix,
                upload_suffix,
                inferences._temp_frame_file_names,
                delete_after_upload=delete_cache_files_after_upload,
            )
        elif data_url:
            final_urls = [
                data_url,
            ]
        else:
            raise Exception("Please provide either a data_url or dataset argument")

        payload = {
            "inferences_id": inferences_id,
            "data_url": final_urls,
            "embedding_distance_metric": embedding_distance_metric,
            "embedding_upload_version": 1,
            "queue_after_dataset": queue_after_dataset,
        }

        if external_metadata is not None:
            if not isinstance(external_metadata, dict) or (
                external_metadata and not isinstance(next(iter(external_metadata)), str)
            ):
                raise Exception("external_metadata must be a dict with string keys")
            payload["external_metadata"] = external_metadata

        if inferences and inferences._temp_frame_embeddings_file_names:
            print("Uploading Inference Embeddings...")
            upload_prefix = "{}_embeddings".format(str(uuid4()))
            upload_suffix = ".arrow"
            uploaded_urls = self._upload_rows_from_files(
                project_id,
                dataset_id,
                upload_prefix,
                upload_suffix,
                inferences._temp_frame_embeddings_file_names,
                delete_after_upload=delete_cache_files_after_upload,
            )
            if uploaded_urls:  # not empty list
                payload["embeddings_url"] = uploaded_urls
        elif embeddings_url:
            payload["embeddings_url"] = [
                embeddings_url,
            ]

        # TODO: Should we just make it not optional, because very few (if any?)
        # people are still using the old data url flow
        if inferences is not None:
            inferences._cleanup_temp_dir()

        print(f"Inferences Processing Queuing...")

        r = requests_retry.post(
            inferences_api_root, headers=self._get_creds_headers(), json=payload
        )
        raise_resp_exception_error(r)
        print(f"Inferences Processing Queued Successfully")

        if wait_until_finish:
            total_inferences_id = f"inferences_{dataset_id}_{inferences_id}"
            print(f"Inferences Processing is waiting on workers to spin up...")
            normalize_status = "PENDING"
            while normalize_status == "PENDING":
                time.sleep(10)
                normalize_status = Client.parse_normalize_process_step_status(
                    self.current_abstract_dataset_process_step_status(
                        project_id, total_inferences_id
                    )
                )
                full_process_state, _ = self.current_inferences_process_state(
                    project_id, dataset_id, inferences_id
                )
                if full_process_state == "FAILED":
                    print("Inferences processing has failed. Exiting...")
                    return
            print(f"Inferences Processing Workers have been spun up.")
            with tqdm(
                total=100.0,
                file=sys.stdout,
                unit_scale=True,
                desc="Inferences Processing Progress",
            ) as pbar:  # type: tqdm[Any]
                start_time = datetime.datetime.now()
                processing_state = "PENDING"
                display_processing_state: Callable[
                    [str], str
                ] = lambda state: f"Inferences Processing State: {state}"
                pbar.write(display_processing_state(processing_state))
                while (datetime.datetime.now() - start_time) < wait_timeout:
                    (
                        new_processing_state,
                        new_percent_done,
                    ) = self.current_inferences_process_state(
                        project_id, dataset_id, inferences_id
                    )
                    pbar.update(new_percent_done - pbar.n)
                    if new_processing_state != processing_state:
                        processing_state = new_processing_state
                        pbar.write(display_processing_state(processing_state))
                    processed = self.is_inferences_processed(
                        project_id=project_id,
                        dataset_id=dataset_id,
                        inferences_id=inferences_id,
                    )
                    if processed:
                        pbar.update(100.0 - pbar.n)
                        pbar.close()
                        print("Inferences are fully processed.")
                        break
                    if processing_state == "FAILED":
                        pbar.update(100.0 - pbar.n)
                        pbar.close()
                        print("Inferences processing has failed. Exiting...")
                        raw_logs = self.get_inferences_ingest_error_logs(
                            project_id, dataset_id, inferences_id
                        )
                        formatted = self._format_error_logs(raw_logs)
                        for entry in formatted:
                            print(entry)

                        break
                    else:
                        time.sleep(10)

                if datetime.datetime.now() - start_time >= wait_timeout:
                    pbar.close()
                    print("Exceeded timeout waiting for job completion.")

    # Even though this is the same implementation as `update_dataset_metadata`, we split it out
    # because users have been working with inference-specific functions
    def update_inferences_metadata(
        self, project_id: str, inferences_id: str, external_metadata: Dict[str, Any]
    ) -> None:
        """Update dataset metadata

        Args:
            project_id (str): The project id.
            inferences_id (str): The inferences id.
            external_metadata (Dict[Any, Any]): The new metadata
        """
        if not isinstance(external_metadata, dict) or (
            external_metadata and not isinstance(next(iter(external_metadata)), str)
        ):
            raise Exception("external_metadata must be a dict with string keys")

        payload = {"external_metadata": external_metadata}
        r = requests_retry.post(
            f"{self.api_endpoint}/projects/{project_id}/datasets/{inferences_id}/metadata",
            headers=self._get_creds_headers(),
            json=payload,
        )
        raise_resp_exception_error(r)
