import base64
import datetime
import gzip
import json
import logging
import math
import os.path
import time
import webbrowser
from contextlib import contextmanager
from dataclasses import dataclass, field
from os import PathLike
from pathlib import Path
from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union, Set, ContextManager, Tuple

import rich.progress
from dataclasses_json import config, LetterCase, DataClassJsonMixin
from pathvalidate import sanitize_filepath

import dagshub.common.config
from dagshub.common import rich_console
from dagshub.common.analytics import send_analytics_event
from dagshub.common.helpers import prompt_user, http_request, log_message
from dagshub.common.rich_util import get_rich_progress
from dagshub.common.util import lazy_load, multi_urljoin, to_timestamp, exclude_if_none
from dagshub.data_engine.client.models import (
    PreprocessingStatus,
    MetadataFieldSchema,
    ScanOption,
    autogenerated_columns,
    DatasetResult,
)
from dagshub.data_engine.dtypes import MetadataFieldType
from dagshub.data_engine.model.datapoint import Datapoint
from dagshub.data_engine.model.errors import (
    WrongOperatorError,
    WrongOrderError,
    DatasetFieldComparisonError,
    FieldNotFoundError,
    DatasetNotFoundError,
)
from dagshub.data_engine.model.metadata_field_builder import MetadataFieldBuilder
from dagshub.data_engine.model.query import QueryFilterTree
from dagshub.data_engine.model.schema_util import metadataTypeLookup, metadataTypeLookupReverse
from dagshub.data_engine.model.datasource_state import DatasourceState

if TYPE_CHECKING:
    from dagshub.data_engine.model.query_result import QueryResult
    import fiftyone as fo
    import pandas
    import mlflow
    import mlflow.entities
else:
    plugin_server_module = lazy_load("dagshub.data_engine.voxel_plugin_server.server")
    fo = lazy_load("fiftyone")
    mlflow = lazy_load("mlflow")

logger = logging.getLogger(__name__)

DEFAULT_MLFLOW_ARTIFACT_NAME = "datasource.dagshub.json"


@dataclass
class DatapointMetadataUpdateEntry(DataClassJsonMixin):
    url: str
    key: str
    value: str
    valueType: MetadataFieldType = field(metadata=config(encoder=lambda val: val.value))
    allowMultiple: bool = False


@dataclass
class DatapointDeleteMetadataEntry(DataClassJsonMixin):
    datapointId: str
    key: str


@dataclass
class DatapointDeleteEntry(DataClassJsonMixin):
    datapointId: str


@dataclass
class Field:
    """
    Class used to define custom fields for use in \
    :func:`Datasource.select() <dagshub.data_engine.model.datasource.Datasource.select>` or in filtering.

    Example of filtering on old data from a field::

        t = datetime.now() - timedelta(days=2)
        q = ds[Field("size", as_of=t)] > 500
        q.all()
    """

    field_name: str
    """The database field where the values are stored. In other words, where to get the values from."""
    as_of: Optional[Union[float, datetime.datetime]] = None
    """
    If defined, the data in this field would be shown as of this moment in time.

    Accepts either a datetime object, or a UTC timestamp.
    """
    alias: Optional[str] = None
    """
    How the returned custom data field should be named.

    Useful when you're comparing the same field at multiple points in time::

        yesterday = datetime.now() - timedelta(days=1)

        ds.select(
            Field("value", alias="value_today"),
            Field("value", as_of=yesterday, alias="value_yesterday")
        ).all()
    """

    @property
    def as_of_timestamp(self) -> Optional[int]:
        if self.as_of is None:
            return None
        return to_timestamp(self.as_of)

    def to_dict(self, ds: "Datasource") -> Dict[str, Any]:
        if not ds.has_field(self.field_name):
            raise FieldNotFoundError(self.field_name)

        res_dict: Dict[str, Union[str, int, None]] = {"name": self.field_name}
        if self.as_of is not None:
            res_dict["asOf"] = self.as_of_timestamp
        if self.alias:
            res_dict["alias"] = self.alias
        return res_dict


class Datasource:
    def __init__(
        self,
        datasource: "DatasourceState",
        query: Optional["DatasourceQuery"] = None,
        from_dataset: Optional["DatasetState"] = None,
    ):
        self._source = datasource
        if query is None:
            query = DatasourceQuery()
        self._query = query
        # a per datasource context used for dict update syntax
        self._implicit_update_ctx: Optional[MetadataContextManager] = None
        # this ref marks if source is currently used in
        # meta-data update 'with' block
        self._explicit_update_ctx: Optional[MetadataContextManager] = None
        self.assigned_dataset = from_dataset

    @property
    def has_explicit_context(self):
        return self._explicit_update_ctx is not None

    def _get_source_implicit_metadata_entries(self):
        if self._implicit_update_ctx is not None:
            return self._implicit_update_ctx.get_metadata_entries()
        else:
            return []

    @property
    def source(self) -> "DatasourceState":
        return self._source

    def clear_query(self, reset_to_dataset=True):
        """
        Clear the attached query.

        Args:
            reset_to_dataset: If ``True`` and this Datasource was saved as a dataset, reset to the query in the dataset,
                instead of clearing the query completely.
        """
        if reset_to_dataset and self.assigned_dataset is not None and self.assigned_dataset.query is not None:
            self._query = self.assigned_dataset.query.__deepcopy__()
        else:
            self._query = DatasourceQuery()

    def __deepcopy__(self, memodict={}) -> "Datasource":
        res = Datasource(self._source, self._query.__deepcopy__())
        res.assigned_dataset = self.assigned_dataset

        # Carry over the update context, that way we'll keep track of the metadata being uploaded
        res._implicit_update_ctx = self._implicit_update_ctx

        return res

    def get_query(self) -> "DatasourceQuery":
        return self._query

    @property
    def annotation_fields(self) -> List[str]:
        """Return all fields that have the annotation meta tag set"""
        return [f.name for f in self.fields if f.is_annotation()]

    def serialize_gql_query_input(self) -> Dict:
        """
        Serialize the query of this Datasource for use in GraphQL querying (e.g. getting datapoints)

        :meta private:
        """
        return self._query.to_dict()

    def _deserialize_from_gql_result(self, query_dict: Dict):
        """
        Imports query information from ``query_dict``
        """
        self._query = DatasourceQuery.from_dict(query_dict)

    def load_from_dataset(
        self, dataset_id: Optional[Union[str, int]] = None, dataset_name: Optional[str] = None, change_query=True
    ):
        """
        Imports query information from a dataset with the specified id or name. Either of id or name could be specified

        Args:
            dataset_id: ID of the dataset
            dataset_name: Name of the dataset
            change_query: Whether to change the query of this object to the query in the dataset

        :meta private:
        """
        datasets = self.source.client.get_datasets(id=dataset_id, name=dataset_name)
        if not datasets:
            raise DatasetNotFoundError(self.source.repo, dataset_id, dataset_name)
        dataset_state = DatasetState.from_gql_dataset_result(datasets[0])
        self.load_from_dataset_state(dataset_state, change_query)

    def load_from_dataset_state(self, dataset_state: "DatasetState", change_query=True):
        """
        Imports query information from a :class:`~dagshub.data_engine.model.dataset_state.DatasetState`

        Args:
            dataset_state: State to load
            change_query: If false, only assigns the dataset.
                If true, also changes the query to be the query of the dataset

        :meta private:
        """
        if self.source.id != dataset_state.datasource_id:
            raise RuntimeError(
                "Dataset belongs to a different datasource "
                f"(This datasource: {self.source.id}, Dataset's datasource: {dataset_state.datasource_id})"
            )
        self.assigned_dataset = dataset_state
        if change_query:
            self._query = dataset_state.query

    def sample(self, start: Optional[int] = None, end: Optional[int] = None):
        if start is not None:
            logger.warning("Starting slices is not implemented for now")
        return self._source.client.sample(self, end, include_metadata=True)

    def head(self, size=100) -> "QueryResult":
        """
        Executes the query and returns a :class:`.QueryResult` object containing first ``size`` datapoints

        Args:
            size: how many datapoints to get. Default is 100
        """
        self._check_preprocess()
        send_analytics_event("Client_DataEngine_DisplayTopResults", repo=self.source.repoApi)
        return self._source.client.head(self, size)

    def all(self) -> "QueryResult":
        """
        Executes the query and returns a :class:`.QueryResult` object containing all datapoints
        """
        self._check_preprocess()
        return self._source.client.get_datapoints(self)

    def select(self, *selected: Union[str, Field]) -> "Datasource":
        """
        Select which fields should appear in the query result.

        If you want to query older versions of metadata,
        use :class:`Field` objects with ``as_of`` set to your desired time.

        By default, only the defined fields are returned.
        If you want to return all existing fields plus whatever additional fields you define,
        add ``"*"`` into the arguments.

        Args:
            selected: Fields you want to select. Can be either of:

                - Name of the field to select: ``"field"``.
                - ``"*"`` to select all the fields in the datasource.
                - :class:`Field` object.

        Example::

            t = datetime.now() - timedelta(hours=24)
            q1 = ds.select("*", Field("size", as_of=t, alias="size_asof_24h_ago"))
            q1.all()
        """

        include_all = False
        selects = []
        for s in selected:
            if isinstance(s, Field):
                selects.append(s.to_dict(self))
            else:
                if s != "*":
                    selects.append({"name": s})
                else:
                    include_all = True

        if include_all:
            aliases = set([s["alias"] for s in selects if "alias" in s])
            for f in self.fields:
                if f.name in aliases:
                    raise ValueError(
                        f"Alias {f.name} can't be used, because * was specified and "
                        f"a field with that name already exists"
                    )
                selects.append({"name": f.name})

        new_ds = self.__deepcopy__()
        new_ds.get_query().select = selects
        return new_ds

    def as_of(self, time: Union[float, datetime.datetime]) -> "Datasource":
        """
        Get a snapshot of the datasource's state as of ``time``.

        Args:
            time: At which point in time do you want to get data from.\
                Either a UTC timestamp or a ``datetime`` object.

        In the following example, you will get back datapoints that were created no later than yesterday AND \
        had their size at this point bigger than 5 bytes::

            t = datetime.now() - timedelta(hours=24)
            q1 = (ds["size"] > 5).as_of(t)
            q1.all()

        .. note::
            If used with :func:`select`, the ``as_of`` set on the fields takes precedence
            over the global query ``as_of`` set here.
        """
        new_ds = self.__deepcopy__()

        new_ds._query.as_of = to_timestamp(time)
        return new_ds

    def order_by(self, *args: Union[str, Tuple[str, Union[bool, str]]]) -> "Datasource":
        """
        Sort the query result by the specified fields.
        Any previously set order will be overwritten.

        Args:
            Fields to sort by. Can be either of:
                - Name of the field to sort by: ``"field"``.
                - A tuple of ``(field_name, ascending)``: ``("field", True)``.
                - A tuple of ``(field_name, "asc"|"desc")``: ``("field", "asc")``.

        Examples::

            ds.order_by("size").all()                   # Order by ascending size
            ds.order_by(("date", "desc"), "size).all()  # Order by descending date, then ascending size
        """
        new_ds = self.__deepcopy__()
        orders = []
        for arg in args:
            if isinstance(arg, str):
                orders.append({"field": arg, "order": "ASC"})
            else:
                if len(arg) != 2:
                    raise RuntimeError(
                        f"Invalid sort argument {arg}, must be a tuple (<field>, 'asc'|'desc'|True|False)"
                    )
                if isinstance(arg[1], bool):
                    order = "ASC" if arg[1] else "DESC"
                elif isinstance(arg[1], str) and arg[1].upper() in ["ASC", "DESC"]:
                    order = arg[1].upper()
                else:
                    raise RuntimeError(f"Invalid sort argument {arg}, second value must be 'asc'|'desc'|True|False")
                orders.append({"field": arg[0], "order": order})
        new_ds.get_query().order_by = orders
        return new_ds

    def _check_preprocess(self):
        self.source.get_from_dagshub()
        if (
            self.source.preprocessing_status == PreprocessingStatus.IN_PROGRESS
            or self.source.preprocessing_status == PreprocessingStatus.AUTO_SCAN_IN_PROGRESS
        ):
            logger.warning(
                f"Datasource {self.source.name} is currently in the progress of rescanning. "
                f"Values might change if you requery later"
            )

    def metadata_field(self, field_name: str) -> MetadataFieldBuilder:
        """
        Returns a builder for a metadata field.
        The builder can be used to change properties of a field or create a new field altogether.
        Note that fields get automatically created when you upload new metadata to the Data Engine,
        so it's not necessary to create fields with this function.

        Example of creating a new annotation field::

            ds.metadata_field("annotation").set_type(dtypes.LabelStudioAnnotation).apply()

        .. note::
            New fields have to have their type defined using ``.set_type()`` before doing anything else

        Example of marking an existing field as an annotation field::

            ds.metadata_field("existing-field").set_annotation().apply()

        Args:
            field_name: Name of the field that you want to create/change

        """
        return MetadataFieldBuilder(self, field_name)

    def apply_field_changes(self, field_builders: List[MetadataFieldBuilder]):
        """
        Applies one or multiple metadata field builders
        that can be constructed using the :func:`metadata_field()` function.
        """
        self.source.client.update_metadata_fields(self, [builder.schema for builder in field_builders])
        self.source.get_from_dagshub()

    @property
    def implicit_update_context(self) -> "MetadataContextManager":
        if not self._implicit_update_ctx:
            self._implicit_update_ctx = MetadataContextManager(self)

        return self._implicit_update_ctx

    def upload_metadata_of_implicit_context(self):
        """
        commit meta data changes done in dictionary assignment context
        :meta private:
        """
        if self._implicit_update_ctx:
            try:
                self._upload_metadata(self._get_source_implicit_metadata_entries())
            finally:
                self._implicit_update_ctx = None

    def metadata_context(self) -> ContextManager["MetadataContextManager"]:
        """
        Returns a metadata context, that you can upload metadata through using its
        :func:`~MetadataContextManager.update_metadata()` function.
        Once the context is exited, all metadata is uploaded in one batch::

            with df.metadata_context() as ctx:
                ctx.update_metadata("file1", {"key1": True, "key2": "value"})

        """

        # Need to have the context manager inside a wrapper to satisfy MyPy + PyCharm type hinter
        @contextmanager
        def func():
            self.source.get_from_dagshub()
            send_analytics_event("Client_DataEngine_addEnrichments", repo=self.source.repoApi)
            ctx = MetadataContextManager(self)

            self._explicit_update_ctx = ctx
            yield ctx
            try:
                self._upload_metadata(ctx.get_metadata_entries() + self._get_source_implicit_metadata_entries())
            finally:
                self._implicit_update_ctx = None
                self._explicit_update_ctx = None

        return func()

    def upload_metadata_from_dataframe(self, df: "pandas.DataFrame", path_column: Optional[Union[str, int]] = None):
        """
        Upload metadata from a pandas dataframe.

        All columns are uploaded as metadata, and the path of every datapoint is taken from `path_column`.

        Args:
            df (`pandas.DataFrame <https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html>`_):
                DataFrame with metadata
            path_column: Column with the datapoints' paths. Can either be the name of the column, or its index.
                If not specified, the first column is used.
        """
        self.source.get_from_dagshub()
        send_analytics_event("Client_DataEngine_addEnrichmentsWithDataFrame", repo=self.source.repoApi)
        self._upload_metadata(self._df_to_metadata(df, path_column, multivalue_fields=self._get_multivalue_fields()))

    def _get_multivalue_fields(self) -> Set[str]:
        res = set()
        for col in self.source.metadata_fields:
            if col.multiple:
                res.add(col.name)
        return res

    def _df_to_metadata(
        self, df: "pandas.DataFrame", path_column: Optional[Union[str, int]] = None, multivalue_fields=set()
    ) -> List[DatapointMetadataUpdateEntry]:
        res: List[DatapointMetadataUpdateEntry] = []
        if path_column is None:
            path_column = df.columns[0]
        elif type(path_column) is str:
            if path_column not in df.columns:
                raise RuntimeError(f"Column {path_column} does not exist in the dataframe")
        elif type(path_column) is int:
            path_column = df.columns[path_column]

        # objects are actually mixed and not guaranteed to be string, but this should cover most use cases
        if df.dtypes[path_column] != "object":
            raise RuntimeError(f"Column {path_column} doesn't have strings")

        field_value_types = {f.name: f.valueType for f in self.fields}

        for _, row in df.iterrows():
            datapoint = row[path_column]
            for key, val in row.items():
                if key == path_column:
                    continue
                key = str(key)
                if key in autogenerated_columns:
                    continue
                if val is None:
                    continue
                # ONLY FOR PANDAS: since pandas doesn't distinguish between None and NaN, don't upload it
                if type(val) is float and math.isnan(val):
                    continue
                if type(val) is list:
                    if key not in multivalue_fields:
                        multivalue_fields.add(key)
                        # Promote all the existing uploading metadata to multivalue
                        for update_entry in res:
                            if update_entry.key == key:
                                update_entry.allowMultiple = True
                    for sub_val in val:
                        value_type = field_value_types.get(key)
                        if value_type is None:
                            value_type = metadataTypeLookup[type(sub_val)]
                            field_value_types[key] = value_type
                        # Don't override bytes if they're not bytes - probably just undownloaded values
                        if value_type == MetadataFieldType.BLOB and type(sub_val) is not bytes:
                            continue
                        # Pandas quirk - integers are floats on the backend
                        if value_type == MetadataFieldType.INTEGER:
                            sub_val = int(sub_val)
                        if type(sub_val) is bytes:
                            sub_val = MetadataContextManager.wrap_bytes(sub_val)
                        res.append(
                            DatapointMetadataUpdateEntry(
                                url=datapoint, key=key, value=str(sub_val), valueType=value_type, allowMultiple=True
                            )
                        )
                else:
                    value_type = field_value_types.get(key)
                    if value_type is None:
                        value_type = metadataTypeLookup[type(val)]
                        field_value_types[key] = value_type
                    # Don't override bytes if they're not bytes - probably just undownloaded values
                    if value_type == MetadataFieldType.BLOB and type(val) is not bytes:
                        continue
                    # Pandas quirk - integers are floats on the backend
                    if value_type == MetadataFieldType.INTEGER:
                        val = int(val)
                    if type(val) is bytes:
                        val = MetadataContextManager.wrap_bytes(val)
                    res.append(
                        DatapointMetadataUpdateEntry(
                            url=datapoint,
                            key=key,
                            value=str(val),
                            valueType=value_type,
                            allowMultiple=key in multivalue_fields,
                        )
                    )
        return res

    def delete_source(self, force: bool = False):
        """
        Delete the record of this datasource along with all datapoints.

        .. warning::
            This is a destructive operation! If you delete the datasource,
            all the datapoints and metadata will be removed.

        Args:
            force: Skip the confirmation prompt
        """
        prompt = (
            f'You are about to delete datasource "{self.source.name}" for repo "{self.source.repo}"\n'
            f"This will remove the datasource and ALL datapoints "
            f"and metadata records associated with the source."
        )
        if not force:
            user_response = prompt_user(prompt)
            if not user_response:
                print("Deletion cancelled")
                return
        self.source.client.delete_datasource(self)

    def scan_source(self, options: Optional[List[ScanOption]] = None):
        """
        This function fires a call to the backend to rescan the datapoints.
        Call this function whenever you uploaded new files and want them to appear when querying the datasource,
        or if you changed existing file contents and want their metadata to be updated.

        DagsHub periodically rescans all datasources, this function is a way to make a scan happen as soon as possible.

        Notes about automatically scanned metadata:
            1. Only new datapoints (files) will be added.
               If files were removed from the source, their metadata will still remain,
               and they will still be returned from queries on the datasource.
               An API to actively remove metadata will be available soon.
            2. Some metadata fields will be automatically scanned and updated by DagsHub based on this scan -
               the list of automatic metadata fields is growing frequently!

        Args:
            options: List of scanning options. If not sure, leave empty.
        """
        logger.debug("Rescanning datasource")
        self.source.client.scan_datasource(self, options=options)

    def _upload_metadata(self, metadata_entries: List[DatapointMetadataUpdateEntry]):
        progress = get_rich_progress(rich.progress.MofNCompleteColumn())

        upload_batch_size = dagshub.common.config.dataengine_metadata_upload_batch_size
        total_entries = len(metadata_entries)
        total_task = progress.add_task(f"Uploading metadata (batch size {upload_batch_size})...", total=total_entries)

        with progress:
            for start in range(0, total_entries, upload_batch_size):
                entries = metadata_entries[start : start + upload_batch_size]
                logger.debug(f"Uploading {len(entries)} metadata entries...")
                self.source.client.update_metadata(self, entries)
                progress.update(total_task, advance=upload_batch_size)
            progress.update(total_task, completed=total_entries, refresh=True)

        # Update the status from dagshub, so we get back the new metadata columns
        self.source.get_from_dagshub()

    def delete_metadata_from_datapoints(self, datapoints: List[Datapoint], fields: List[str]):
        """
        Delete metadata from datapoints.
        The deleted values can be accessed using versioned query with time set before the deletion

        Args:
            datapoints: datapoints to delete metadata from
            fields: fields to delete
        """
        metadata_entries = []
        for d in datapoints:
            for n in fields:
                metadata_entries.append(DatapointDeleteMetadataEntry(datapointId=d.datapoint_id, key=n))
        self.source.client.delete_metadata_for_datapoint(self, metadata_entries)

    def delete_datapoints(self, datapoints: List[Datapoint], force: bool = False):
        """
        Delete datapoints.

        - These datapoints will no longer show up in queries.
        - Does not delete the datapoint's file, only removing the data from the datasource.
        - You can still query these datapoints and associated metadata with \
        versioned queries whose time is before deletion time.
        - You can re-add these datapoints to the datasource by uploading new metadata to it with, for example, \
        :func:`Datasource.metadata_context <dagshub.data_engine.model.datasource.Datasource.metadata_context>`. \
        This will create a new datapoint with new id and new metadata records.
        - Datasource scanning will *not* add these datapoints back.

        Args:
            datapoints: list of datapoints objects to delete
            force: Skip the confirmation prompt
        """
        dps_str = "\n\t".join([""] + [d.path for d in datapoints])
        prompt = (
            f"You are about to delete the following datapoint(s): {dps_str}\n"
            f"This will remove the datapoint and metadata from unversioned queries, "
            f"but won't delete the underlying file."
        )
        if not force:
            user_response = prompt_user(prompt)
            if not user_response:
                print("Deletion cancelled")
                return

        self.source.client.delete_datapoints(
            self, [DatapointDeleteEntry(datapointId=d.datapoint_id) for d in datapoints]
        )

    def save_dataset(self, name: str) -> "Datasource":
        """
        Save the dataset, which is a combination of datasource + query, on the backend.
        That way you can persist and share your queries.
        You can get the dataset back later by calling :func:`.datasets.get_dataset()`

        Args:
            name: Name of the dataset

        Returns:
            A datasource object with the dataset assigned to it
        """
        send_analytics_event("Client_DataEngine_QuerySaved", repo=self.source.repoApi)

        self.source.client.save_dataset(self, name)
        log_message(f"Dataset {name} saved")

        copy_with_ds_assigned = self.__deepcopy__()
        copy_with_ds_assigned.load_from_dataset(dataset_name=name, change_query=False)
        return copy_with_ds_assigned

    def log_to_mlflow(
        self, artifact_name=DEFAULT_MLFLOW_ARTIFACT_NAME, run: Optional["mlflow.entities.Run"] = None
    ) -> "mlflow.Entities.Run":
        """
        Logs the current datasource state to MLflow as an artifact.

        Args:
            artifact_name: Name of the artifact that will be stored in the MLflow run.
            run: MLflow run to save to. If ``None``, uses the active MLflow run or creates a new run.

        Returns:
            Run to which the artifact was logged.
        """
        if run is None:
            run = mlflow.active_run()
            if run is None:
                run = mlflow.start_run()
        mlflow.MlflowClient().log_dict(run.info.run_id, self._to_dict(), artifact_name)
        log_message(f'Saved the datasource state to MLflow (run "{run.info.run_name}") as "{artifact_name}"')
        return run

    def save_to_file(self, path: Union[str, PathLike] = ".") -> Path:
        """
        Saves a JSON file representing the current state of datasource or dataset.
        Useful for connecting code versions to the datasource used for training.

        .. note::
            Does not save the actual contents of the datasource/dataset, only the query.

        Args:
            path: Where to save the file. If path is an existing folder, saves to ``<path>/<ds_name>.json``.

        Returns:
            The path to the saved file
        """
        path = Path(path)
        if path.is_dir():
            if self.assigned_dataset is not None:
                name = self.assigned_dataset.dataset_name
            else:
                name = self.source.name
            path = path / f"{name}.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        res = self._to_dict()
        with open(path, "w") as file:
            json.dump(res, file, indent=4, sort_keys=True)
        log_message(f"Datasource saved to '{path}'")

        return path

    def _serialize(self) -> "DatasourceSerializedState":
        res = DatasourceSerializedState(
            repo=self.source.repo,
            datasource_id=self.source.id,
            datasource_name=self.source.name,
            query=self._query,
            timestamp=datetime.datetime.now().timestamp(),
            modified=self.is_query_different_from_dataset,
        )
        if self.assigned_dataset is not None:
            res.dataset_id = self.assigned_dataset.dataset_id
            res.dataset_name = self.assigned_dataset.dataset_name
        return res

    def _to_dict(self) -> Dict:
        res = self._serialize().to_dict()
        # Skip Nones in the result
        res = {k: v for k, v in res.items() if v is not None}
        return res

    @property
    def is_query_different_from_dataset(self) -> Optional[bool]:
        """
        Is the current query of the object different from the one in the assigned dataset.

        If no dataset is assigned, returns ``None``.
        """
        if self.assigned_dataset is None:
            return None
        return self._query.to_dict() != self.assigned_dataset.query.to_dict()

    @staticmethod
    def load_from_serialized_state(state_dict: Dict) -> "Datasource":
        """
        Load a Datasource that was saved with :func:`save_to_file`

        Args:
            state_dict: Serialized JSON object
        """

        state = DatasourceSerializedState.from_dict(state_dict)

        ds_state = DatasourceState(repo=state.repo, name=state.datasource_name, id=state.datasource_id)
        ds_state.get_from_dagshub()
        ds = Datasource(ds_state)
        ds._query = state.query

        if state.dataset_id is not None:
            ds.load_from_dataset(state.dataset_id, state.dataset_name, change_query=False)

        return ds

    def to_voxel51_dataset(self, **kwargs) -> "fo.Dataset":
        """
        Refer to :func:`QueryResult.to_voxel51_dataset() \
        <dagshub.data_engine.model.query_result.QueryResult.to_voxel51_dataset>`\
        for documentation.
        """
        return self.all().to_voxel51_dataset(**kwargs)

    @property
    def default_dataset_location(self) -> Path:
        """
        Default location where datapoint files are stored.

        On UNIX-likes the path is ``~/dagshub/datasets/<repo_name>/<datasource_id>``

        On Windows the path is ``C:\\Users\\<user>\\dagshub\\datasets\\<repo_name>\\<datasource_id>``
        """
        return Path(
            sanitize_filepath(os.path.join(Path.home(), "dagshub", "datasets", self.source.repo, str(self.source.id)))
        )

    def visualize(self, **kwargs) -> "fo.Session":
        """
        Visualize the whole datasource using
        :func:`QueryResult.visualize() <dagshub.data_engine.model.query_result.QueryResult.visualize>`.

        Read the function docs for kwarg documentation.
        """
        return self.all().visualize(**kwargs)

    @property
    def fields(self) -> List[MetadataFieldSchema]:
        return self.source.metadata_fields

    def annotate(self, fields_to_embed=None, fields_to_exclude=None) -> Optional[str]:
        """
        Sends all datapoints in the datasource for annotation in Label Studio.

        Args:
            fields_to_embed: list of meta-data columns that will show up in Label Studio UI.
             if not specified all will be displayed.
            fields_to_exclude: list of meta-data columns that will not show up in Label Studio UI

        .. note::
            This will send ALL datapoints in the datasource for annotation.
            It's recommended to not send a huge amount of datapoints to be annotated at once, to avoid overloading
            the Label Studio workspace.
            Use :func:`QueryResult.annotate() <dagshub.data_engine.model.query_result.QueryResult.annotate>`
            to annotate a result of a query with less datapoints.
            Alternatively, use a lower level :func:`send_datapoints_to_annotation` function

        :return: Link to open Label Studio in the browser
        """
        return self.all().annotate(fields_to_exclude=fields_to_exclude, fields_to_embed=fields_to_embed)

    def send_to_annotation(self):
        """
        deprecated, see :func:`annotate()`

        :meta private:
        """
        return self.annotate()

    def send_datapoints_to_annotation(
        self,
        datapoints: Union[List[Datapoint], "QueryResult", List[Dict]],
        open_project=True,
        ignore_warning=False,
        fields_to_exclude=None,
        fields_to_embed=None,
    ) -> Optional[str]:
        """
        Sends datapoints for annotation in Label Studio.

        Args:
            datapoints: Either of:

                - A :class:`.QueryResult`
                - List of :class:`.Datapoint` objects
                - List of dictionaries. Each dictionary should have fields ``id`` and ``download_url``.
                    ``id`` is the ID of the datapoint in the datasource.

            open_project: Automatically open the created Label Studio project in the browser.
            ignore_warning: Suppress the prompt-warning if you try to annotate too many datapoints at once.
            fields_to_embed: list of meta-data columns that will show up in Label Studio UI.
             if not specified all will be displayed.
            fields_to_exclude: list of meta-data columns that will not show up in Label Studio UI
        Returns:
            Link to open Label Studio in the browser
        """
        for f in (fields_to_embed or []) + (fields_to_exclude or []):
            if not self.has_field(f):
                raise FieldNotFoundError(f)

        if len(datapoints) == 0:
            logger.warning("No datapoints provided to be sent to annotation")
            return None
        elif len(datapoints) > dagshub.common.config.recommended_annotate_limit and not ignore_warning:
            force = prompt_user(
                f"You are attempting to annotate {len(datapoints)} datapoints at once - it's "
                f"recommended to only annotate up to "
                f"{dagshub.common.config.recommended_annotate_limit} "
                f"datapoints at a time."
            )
            if not force:
                return ""

        req_data = {
            "datasource_id": self.source.id,
            "datapoints": [],
            "ls_meta_excludes": fields_to_exclude,
            "ls_meta_includes": fields_to_embed,
        }

        for dp in datapoints:
            req_dict = {}
            if type(dp) is dict:
                req_dict["id"] = dp["datapoint_id"]
                req_dict["download_url"] = dp["download_url"]
            else:
                req_dict["id"] = dp.datapoint_id
                req_dict["download_url"] = dp.download_url
            req_data["datapoints"].append(req_dict)

        init_url = multi_urljoin(self.source.repoApi.data_engine_url, "annotations/init")
        resp = http_request("POST", init_url, json=req_data, auth=self.source.repoApi.auth)

        if resp.status_code != 200:
            logger.error(f"Error while sending request for annotation: {resp.content}")
            return None
        link = resp.json()["link"]

        # Do a raw print so it works in colab/jupyter
        print("Open the following link to start working on your annotation project:")
        print(link)

        if open_project:
            webbrowser.open_new_tab(link)
        return link

    def _launch_annotation_workspace(self):
        try:
            start_workspace_url = multi_urljoin(self.source.repoApi.annotations_url, "start")
            http_request("POST", start_workspace_url, auth=self.source.repoApi.auth)
        except:  # noqa
            pass

    def wait_until_ready(self, max_wait_time=300, fail_on_timeout=True):
        """
        Blocks until the datasource preprocessing is complete.

        Useful when you have just created a datasource and the initial scanning hasn't finished yet.

        Args:
            max_wait_time: Maximum time to wait in seconds
            fail_on_timeout: Whether to raise a RuntimeError or log a warning if the scan does not complete on time
        """

        # Start LS workspace to save time later in the flow
        self._launch_annotation_workspace()

        start = time.time()
        if max_wait_time:
            rich_console.log(f"Maximum waiting time set to {int(max_wait_time / 60)} minutes")
        spinner = rich_console.status("Waiting for datasource preprocessing to complete...")
        with spinner:
            while True:
                self.source.get_from_dagshub()
                if self.source.preprocessing_status == PreprocessingStatus.READY:
                    return

                if self.source.preprocessing_status == PreprocessingStatus.FAILED:
                    raise RuntimeError("Datasource preprocessing failed")

                if max_wait_time is not None and (time.time() - start) > max_wait_time:
                    if fail_on_timeout:
                        raise RuntimeError(
                            f"Time limit of {max_wait_time} seconds reached before processing was completed."
                        )
                    else:
                        logger.warning(
                            f"Time limit of {max_wait_time} seconds reached before processing was completed."
                        )
                        return

                time.sleep(1)

    def has_field(self, field_name: str) -> bool:
        """
        Checks if a metadata field ``field_name`` exists in the datasource.
        """
        reserved_searchable_fields = ["path"]
        fields = (f.name for f in self.fields)
        return field_name in reserved_searchable_fields or field_name in fields

    def __repr__(self):
        res = f"Datasource {self.source.name}"
        res += f"\n\tRepo: {self.source.repo}, path: {self.source.path}"
        if self.assigned_dataset:
            res += (
                f"\n\tAssigned Dataset: {self.assigned_dataset.dataset_name} "
                f"(Query was modified: {self.is_query_different_from_dataset})"
            )
        res += f"\n\t{self._query}"
        res += "\n\tFields:"
        for f in self.fields:
            res += f"\n\t\t{f}"
        return res + "\n"

    """ FUNCTIONS RELATED TO QUERYING
    These are functions that overload operators on the DataSet, so you can do pandas-like filtering
        ds = Dataset(...)
        queried_ds = ds[ds["value"] == 5]
    """

    def __getitem__(self, other: Union[slice, str, "Datasource", "Field"]):
        # Slicing - get items from the slice
        if type(other) is slice:
            return self.sample(other.start, other.stop)

        # Otherwise we're doing querying
        new_ds = self.__deepcopy__()
        if isinstance(other, (str, Field)):
            query_field: Field = Field(other) if type(other) is str else other
            other_query = QueryFilterTree(
                query_field.field_name,
                query_field.as_of_timestamp,
            )
            if not self.has_field(query_field.field_name):
                raise FieldNotFoundError(query_field.field_name)
            if self._query.filter.is_empty:
                new_ds._query.filter = other_query
            else:
                new_ds._query.compose("and", other_query)
            return new_ds
        # "index" is a datasource with a query - return the datasource inside
        # Example:
        #   ds = Dataset()
        #   filtered_ds = ds[ds["aaa"] > 5]
        #   filtered_ds2 = filtered_ds[filtered_ds["bbb"] < 4]
        # filtered_ds2 will be "aaa" > 5 AND "bbb" < 4
        if isinstance(other, Datasource):
            return other

    def __gt__(self, other: object):
        self._test_not_comparing_other_ds(other)
        if not isinstance(other, (int, float, str)):
            raise NotImplementedError
        return self.add_query_op("gt", other)

    def __ge__(self, other: object):
        self._test_not_comparing_other_ds(other)
        if not isinstance(other, (int, float, str)):
            raise NotImplementedError
        return self.add_query_op("ge", other)

    def __le__(self, other: object):
        self._test_not_comparing_other_ds(other)
        if not isinstance(other, (int, float, str)):
            raise NotImplementedError
        return self.add_query_op("le", other)

    def __lt__(self, other: object):
        self._test_not_comparing_other_ds(other)
        if not isinstance(other, (int, float, str)):
            raise NotImplementedError
        return self.add_query_op("lt", other)

    def __eq__(self, other: object):
        self._test_not_comparing_other_ds(other)
        if other is None:
            return self.is_null()
        if not isinstance(other, (int, float, str)):
            raise NotImplementedError
        return self.add_query_op("eq", other)

    def __ne__(self, other: object):
        self._test_not_comparing_other_ds(other)
        if other is None:
            return self.is_not_null()
        if not isinstance(other, (int, float, str)):
            raise NotImplementedError
        return self.add_query_op("eq", other).add_query_op("not")

    def __invert__(self):
        return self.add_query_op("not")

    def __contains__(self, item):
        raise WrongOperatorError("Use `ds.contains(a)` for querying instead of `a in ds`")

    def contains(self, item: str):
        """
        Check if the filtering field contains the specified string item.

        :meta private:
        """
        if type(item) is not str:
            return WrongOperatorError(f"Cannot use contains with non-string value {item}")
        self._test_not_comparing_other_ds(item)
        return self.add_query_op("contains", item)

    def is_null(self):
        """
        Check if the filtering field is null.

        :meta private:
        """
        field = self._get_filtering_field()
        value_type = metadataTypeLookupReverse[field.valueType.value]
        return self.add_query_op("isnull", value_type())

    def is_not_null(self):
        """
        Check if the filtering field is not null.

        :meta private:
        """
        return self.is_null().add_query_op("not")

    def _get_filtering_field(self) -> MetadataFieldSchema:
        field_name = self.get_query().filter.column_filter
        if field_name is None:
            raise RuntimeError("The current query filter is not a field")
        for col in self.source.metadata_fields:
            if col.name == field_name:
                return col
        raise RuntimeError(f"Field {field_name} doesn't exist in the current uploaded metadata")

    def __and__(self, other: "Datasource"):
        return self.add_query_op("and", other)

    def __or__(self, other: "Datasource"):
        return self.add_query_op("or", other)

    # Prevent users from messing up their queries due to operator order
    # They always need to put the dataset query filters in parentheses, otherwise the binary and/or get executed before
    def __rand__(self, other):
        if type(other) is not Datasource:
            raise WrongOrderError(type(other))
        raise NotImplementedError

    def __ror__(self, other):
        if type(other) is not Datasource:
            raise WrongOrderError(type(other))
        raise NotImplementedError

    def add_query_op(
        self, op: str, other: Optional[Union[str, int, float, "Datasource", "QueryFilterTree"]] = None
    ) -> "Datasource":
        """
        Add a query operation to the current Datasource instance.

        Args:
            op (str): The operation to be performed in the query.
            other (Optional[Union[str, int, float, "Datasource", "DatasourceQuery"]], optional):
                The operand for the query operation. Defaults to None.

        Returns:
            Datasource: A new Datasource instance with the added query operation.

        :meta private:
        """
        new_ds = self.__deepcopy__()
        if type(other) is Datasource:
            other = other.get_query()
        new_ds._query.compose(op, other)
        return new_ds

    @staticmethod
    def _test_not_comparing_other_ds(other):
        if type(other) is Datasource:
            raise DatasetFieldComparisonError()


class MetadataContextManager:
    """
    Context manager for updating the metadata on a datasource.
    Batches the metadata changes, so they are being sent all at once.
    """

    def __init__(self, datasource: Datasource):
        self._datasource = datasource
        self._metadata_entries: List[DatapointMetadataUpdateEntry] = []
        self._multivalue_fields = datasource._get_multivalue_fields()

    def update_metadata(self, datapoints: Union[List[str], str], metadata: Dict[str, Any]):
        """
        Update metadata for the specified datapoints.

        .. note::
            If ``datapoints`` is a list, the same metadata is assigned to all the datapoints in the list.
            Call ``update_metadata()`` separately for each datapoint if you need to assign different metadata.

        Args:
            datapoints (Union[List[str], str]): A list of datapoints or a single datapoint path to update metadata for.
            metadata (Dict[str, Any]): A dictionary containing metadata key-value pairs to update.

        Example::

            with ds.metadata_context() as ctx:
                metadata = {
                    "episode": 5,
                    "has_baby_yoda": True,
                }

                # Attach metadata to a single specific file in the datasource.
                # The first argument is the filepath to attach metadata to, **relative to the root of the datasource**.
                ctx.update_metadata("images/005.jpg", metadata)

                # Attach metadata to several files at once:
                ctx.update_metadata(["images/006.jpg","images/007.jpg"], metadata)

        """
        if isinstance(datapoints, str):
            datapoints = [datapoints]

        field_value_types = {f.name: f.valueType for f in self._datasource.fields}

        for dp in datapoints:
            for k, v in metadata.items():
                if v is None:
                    continue
                if k in autogenerated_columns:
                    continue

                if type(v) is list:
                    if k not in self._multivalue_fields:
                        self._multivalue_fields.add(k)
                        # Promote all existing ones to multivalue
                        for e in self._metadata_entries:
                            if e.key == k:
                                e.allowMultiple = True
                    for sub_val in v:
                        value_type = field_value_types.get(k)
                        if value_type is None:
                            value_type = metadataTypeLookup[type(sub_val)]
                            field_value_types[k] = value_type
                        # Don't override bytes if they're not bytes - probably just undownloaded values
                        if value_type == MetadataFieldType.BLOB and type(sub_val) is not bytes:
                            continue

                        if type(v) is bytes:
                            sub_val = self.wrap_bytes(sub_val)
                        self._metadata_entries.append(
                            DatapointMetadataUpdateEntry(
                                url=dp,
                                key=k,
                                value=str(sub_val),
                                # todo: preliminary type check
                                valueType=value_type,
                                allowMultiple=k in self._multivalue_fields,
                            )
                        )

                else:
                    value_type = field_value_types.get(k)
                    if value_type is None:
                        value_type = metadataTypeLookup[type(v)]
                        field_value_types[k] = value_type
                    # Don't override bytes if they're not bytes - probably just undownloaded values
                    if value_type == MetadataFieldType.BLOB and type(v) is not bytes:
                        continue

                    if type(v) is bytes:
                        v = self.wrap_bytes(v)

                    self._metadata_entries.append(
                        DatapointMetadataUpdateEntry(
                            url=dp,
                            key=k,
                            value=str(v),
                            valueType=value_type,
                            # todo: preliminary type check
                            allowMultiple=k in self._multivalue_fields,
                        )
                    )

    @staticmethod
    def wrap_bytes(val: bytes) -> str:
        """
        Handles bytes values for uploading metadata
        The process is gzip -> base64

        :meta private:
        """
        compressed = gzip.compress(val)
        return base64.b64encode(compressed).decode("utf-8")

    def get_metadata_entries(self):
        return self._metadata_entries


@dataclass
class DatasourceQuery(DataClassJsonMixin):
    as_of: Optional[int] = field(default=None, metadata=config(exclude=exclude_if_none, letter_case=LetterCase.CAMEL))
    select: Optional[List[Dict]] = field(default=None, metadata=config(exclude=exclude_if_none))
    filter: "QueryFilterTree" = field(
        default=QueryFilterTree(),
        metadata=config(field_name="query", encoder=QueryFilterTree.serialize, decoder=QueryFilterTree.deserialize),
    )
    order_by: Optional[List] = field(
        default=None, metadata=config(exclude=exclude_if_none, letter_case=LetterCase.CAMEL)
    )

    def __deepcopy__(self, memodict={}):
        other = DatasourceQuery(
            as_of=self.as_of,
            filter=self.filter.__deepcopy__(),
        )
        if self.select is not None:
            other.select = self.select.copy()
        if self.order_by is not None:
            other.order_by = self.order_by.copy()
        return other

    def compose(
        self, op: str, other: Optional[Union[str, int, float, "DatasourceQuery", "QueryFilterTree", "Datasource"]]
    ):
        other_query: Optional["DatasourceQuery"] = None
        # Extract the filter tree for composition
        if isinstance(other, (DatasourceQuery, Datasource)):
            if type(other) is Datasource:
                other = other.get_query()
            other_query = other
            other = other.filter

        self.filter.compose(op, other)

        # If composition went successfully, carry over the as_of, select and order_by from the other query
        if other_query is not None:
            if other_query.select is not None:
                self.select = other_query.select
            if other_query.as_of is not None:
                self.as_of = other_query.as_of
            if other_query.order_by is not None:
                self.order_by = other_query.order_by


@dataclass
class DatasourceSerializedState(DataClassJsonMixin):
    """
    Serialized state of the datasource.
    This should be enough to recreate the exact copy of the datasource back (with additional requests)

    Also carries additional information that might be useful for the user:
        - if the state of the datasource at the point of saving differed from the dataset
        - the timestamp of saving
        - link to open the datasource on DagsHub
    """

    repo: str
    """Repository this datasource is on"""
    datasource_id: Union[str, int]
    """ID of the datasource"""
    datasource_name: str
    """Name of the datasource"""
    query: DatasourceQuery
    """Query at the time of saving"""
    dataset_id: Optional[Union[str, int]] = None
    """ID of the assigned dataset"""
    dataset_name: Optional[str] = None
    """Name of the assigned dataset"""
    timestamp: Optional[float] = None
    """Timestamp of serialization"""
    modified: Optional[bool] = None
    """Does the query differ from the query in the assigned dataset"""
    link: Optional[str] = None
    """URL to open this datasource on DagsHub"""


@dataclass
class DatasetState:
    """
    Information about the Dataset.

    Dataset is a Datasource with a Query applied on it.
    """

    dataset_id: Union[str, int]
    """
    ID of the dataset
    """
    dataset_name: str
    """
    Name of the dataset
    """
    datasource_id: Union[str, int]
    """
    ID of the datasource with which this dataset is associated
    """
    query: DatasourceQuery
    """
    Query of this dataset
    """

    @staticmethod
    def from_dataset_query(
        dataset_id: Union[str, int], dataset_name: str, datasource_id: Union[str, int], dataset_query: Union[Dict, str]
    ) -> "DatasetState":
        if type(dataset_query) is str:
            dataset_query = json.loads(dataset_query)
        query = DatasourceQuery.from_dict(dataset_query)
        res = DatasetState(dataset_id=dataset_id, dataset_name=dataset_name, datasource_id=datasource_id, query=query)
        return res

    @staticmethod
    def from_gql_dataset_result(dataset_result: "DatasetResult") -> "DatasetState":
        return DatasetState.from_dataset_query(
            dataset_id=dataset_result.id,
            dataset_name=dataset_result.name,
            datasource_id=dataset_result.datasource.id,
            dataset_query=dataset_result.datasetQuery,
        )
