import contextlib
import csv
import functools
import hashlib
import io
import json
import logging
import os
import re
import shutil
import uuid
from collections import defaultdict
from dataclasses import asdict, dataclass
from keyword import iskeyword
from pathlib import Path
from typing import IO, Any, Callable, Dict, List, Optional, Union

import datasets
import jinja2
import pandas as pd
import requests
import yaml
from typeguard import typechecked

from gantry.api_client import APIClient
from gantry.dataset import constants
from gantry.exceptions import (
    DataSchemaMismatchError,
    DatasetCommitNotFoundError,
    DatasetDeletedException,
    DatasetHeadOutOfDateException,
    DatasetNotFoundError,
    GantryRequestException,
    NoTabularDataError,
)
from gantry.query.core.dataframe import GantryDataFrame
from gantry.utils import (
    download_file_from_url,
    get_files_checksum,
    list_all_files,
    parse_s3_path,
    read_lines_from_url,
    upload_file_to_url,
)

logger = logging.getLogger(__name__)

BATCH_SIZE = 5


def _ensure_not_deleted(func: Callable) -> Callable:
    """
    Decorator to ensure we do not perform an operation on a deleted dataset, and instead warn the
    user that the dataset is deleted.
    """

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        resp = self._api_client.request(
            "GET",
            f"/api/v1/datasets/{self.dataset_name}",
            raise_for_status=True,
        )
        dataset_info = resp["data"]

        if dataset_info["disabled"]:
            raise DatasetDeletedException("This dataset has been deleted!")

        return func(self, *args, **kwargs)

    return wrapper


def _ensure_dataset_exists_locally(func: Callable) -> Callable:
    """
    Decorator to ensure dataset exists locally, if not will warn the user to download the dataset.
    """

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        if not Path(self._get_file_path(constants.DATASET_MANIFEST_FILE)).exists():
            raise DatasetNotFoundError(
                f"Can't find dataset {self.dataset_name} locally, please run pull()"
                f"to download the dataset into your local workspace."
            )

        return func(self, *args, **kwargs)

    return wrapper


def _warn_when_local_changes(func: Callable) -> Callable:
    """
    Decorator to ensure there are no local changes in the dataset. If there are, will return None
    """

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        # only do the following checking when HEAD file exist(dataset exists on local host)
        if os.path.exists(self._get_file_path(constants.DATASET_HEAD_FILE)):
            localdiff = self.get_diff()
            if (
                localdiff[constants.DELETED_FILES]
                or localdiff[constants.MODIFIED_FILES]
                or localdiff[constants.NEW_FILES]
            ) and not kwargs.get("forced", False):
                logger.warning(
                    "Local changes detected! Please run stash() to stash your local changes, "
                    "or run this function again with forced=True"
                )
                return None

        return func(self, *args, **kwargs)

    return wrapper


@dataclass
class DatasetFileInfo:
    file_name: str
    url: str
    version_id: Optional[str] = None
    sha256: Optional[str] = None

    def to_jsonl(self):
        if not self.file_name or not self.url or not self.version_id or not self.sha256:
            raise ValueError(
                f"Failed to create file info jsonl: {json.dumps(asdict(self))}. \
                    Incomplete dataset file info line! "
            )

        return f"{json.dumps(asdict(self))}\n"


class GantryDataset:
    """
    A class representing a lightweight container around your data that
    provides simple version control semantics for helping to make ML
    pipelines reproducible.
    """

    def __init__(
        self,
        api_client: APIClient,
        dataset_name: str,
        dataset_id: uuid.UUID,
        bucket_name: str,
        aws_region: str,
        dataset_s3_prefix: str,
        workspace: str,
    ):
        """
        :meta private:
        """
        self.dataset_name = dataset_name
        self.workspace = workspace
        self._api_client = api_client
        self._dataset_id = dataset_id
        self._bucket_name = bucket_name
        self._dataset_s3_prefix = dataset_s3_prefix
        self._aws_region = aws_region

    def push_gantry_dataframe(self, gantrydf: GantryDataFrame, filename: str = ""):
        """
        This method will take a GantryDataFrame as input, and save its contents to a new file in
        a Dataset. [TODO: add ability to specify custom version message]

        Example usage:

            >>> dataset = gdataset.get_dataset(dataset_name)
            >>> os.listdir(workspace + dataset_name)
            ['.dataset_metadata', 'dataset_config.yaml', 'README.md', 'tabular_manifests']
            >>> os.listdir(workspace + dataset_name + '/tabular_manifests')
            []
            >>> gantrydf = gquery.query(
            ...   application=application,
            ...   start_time=start,
            ...   end_time=end,
            ...   version="1.1",
            ... )
            >>>
            >>> _ = dataset.push_gantry_dataframe(gantrydf)
            >>>
            >>> _ = dataset.pull()
            >>> os.listdir(workspace + dataset_name + '/tabular_manifests')
            ['GantryDataAutoSave__4390b343-b802-4add-ab51-f86e53979c73_2023-02-13_12:09_to_2023-02-13_17:09_rows_0_69.csv']

        Args:
            gantrydf (GantryDataFrame): the GantryDataFrame that needs to be added to the dataset.
            filename (str): the name to be given to the file with data from the GantryDataFrame.
        Returns:
            Latest commit info.
        """
        start_time = gantrydf.query_info.start_time
        end_time = gantrydf.query_info.end_time
        final_filename = filename
        if filename == "":
            final_filename = f"GantryDataAutoSave__{uuid.uuid4()}_{start_time.strftime('%Y-%m-%d_%H:%M')}_to_{end_time.strftime('%Y-%m-%d_%H:%M')}"  # noqa: E501
        app_node_id = gantrydf.query_info.application_node_id
        filters = gantrydf.filters
        tags = gantrydf.tags
        resp = self._api_client.request(
            "POST",
            f"/api/v1/datasets/{self._dataset_id}/raw_data",
            json={
                "model_id": app_node_id,
                "raw_data_query_parameters": {
                    "filters": filters,
                    "tags": tags,
                    "include_join_id": False,
                    "hidden_columns": None,
                    "start_time": start_time,
                    "end_time": end_time,
                },
                "raw_data_csv_prefix": final_filename,
            },
            raise_for_status=True,
        )
        commit_info = resp["data"]
        return self._prune_commit_info(commit_info)

    def list_versions(self) -> List[Dict[str, Any]]:
        """
        List dataset version history

        Returns:
            ``List[Dict[str, Any]]``: dataset versions from latest to earliest.
        """
        response = self._api_client.request(
            "GET", f"/api/v1/datasets/{self._dataset_id}/commits", raise_for_status=True
        )
        return [self._prune_commit_info(commit_info) for commit_info in response["data"]]

    @_ensure_dataset_exists_locally
    def get_diff(self) -> Dict[str, List[str]]:
        """
        Get lists of new, modified, and deleted files in your local
        working directory.

        Returns:
            ``dict``: a dictionary representing the diff, which looks like this:

            .. code-block:: json

                {
                    "new_files": List[str],
                    "modified_files": List[str],
                    "deleted_files": List[str],
                }
        """

        diff = self._get_diff()
        return {
            constants.NEW_FILES: [f for f, _ in diff[constants.NEW_FILES]],
            constants.MODIFIED_FILES: [f for f, _ in diff[constants.MODIFIED_FILES]],
            constants.DELETED_FILES: [f for f, _ in diff[constants.DELETED_FILES]],
        }

    @typechecked
    @_ensure_dataset_exists_locally
    @_ensure_not_deleted
    def push_version(self, message: str) -> Dict[str, Any]:
        """
        Create a new version of this dataset that snapshots any new, modified, or deleted
        files that exist in the local ``workspace``, a directory where a local copy of your
        dataset lives, and pushes it to Gantry.

        Args:
            message (str): a non-empty string representing a message to associate with
                the created version.
        Returns:
            ``Dict[str, Any]``: a dictionary with metadata representing the version created.
        """
        commit_id = uuid.uuid4()
        diff = self._get_diff_for_commit(commit_id)

        if (
            not diff[constants.NEW_FILES]
            and not diff[constants.MODIFIED_FILES]
            and not diff[constants.DELETED_FILES]
        ):
            logger.warning("No local changes to add!")
            return self._get_current_commit()

        # Check that HEAD is up to date
        resp = self._api_client.request(
            "GET", f"/api/v1/datasets/{self._dataset_id}/commits", raise_for_status=True
        )
        latest_commit_id = resp["data"][0]["id"]

        with open(self._get_file_path(constants.DATASET_HEAD_FILE), "r") as f:
            current_commit = json.load(f)

            if current_commit["id"] != latest_commit_id:
                raise DatasetHeadOutOfDateException

        diff[constants.NEW_FILES] = self._upload_files(diff[constants.NEW_FILES])
        diff[constants.MODIFIED_FILES] = self._upload_files(diff[constants.MODIFIED_FILES])

        self._backup_file(constants.DATASET_MANIFEST_FILE)
        self._update_manifest_file(diff)

        result = self._upload_files(
            [
                DatasetFileInfo(
                    file_name=constants.DATASET_MANIFEST_FILE,
                    url=self._get_s3_url(constants.DATASET_MANIFEST_FILE),
                )
            ]
        )
        version_id = result[0].version_id
        try:
            resp = self._api_client.request(
                "POST",
                f"/api/v1/datasets/{self._dataset_id}/commits",
                json={
                    "message": message,
                    "metadata_s3_file_version": version_id,
                    "parent_commit_id": self._get_current_commit()["id"],
                    "commit_id": commit_id,
                },
                raise_for_status=True,
            )
            os.remove(
                self._get_file_path(f"{constants.DATASET_MANIFEST_FILE}{constants.BACKUP_SUFFIX}")
            )
        except GantryRequestException as e:
            self._recover_file(constants.DATASET_MANIFEST_FILE)
            if (
                e.status_code == 400
                and 'Details: {"error":"Parent version out of date!","response":"error"}'
                in e.args[0]
            ):
                raise DatasetHeadOutOfDateException
            raise e

        commit_info = resp["data"]
        self._update_head(commit_info)

        return self._prune_commit_info(commit_info)

    @typechecked
    @_ensure_not_deleted
    def push_file(
        self,
        file: IO,
        filename: Optional[str] = None,
        message: Optional[str] = None,
        parent_version_id: Optional[Union[uuid.UUID, str]] = None,
    ) -> Dict[str, Any]:
        """
        Create a new version of this dataset that contains a single new file passed as the
        ``file`` parameter for this method, and returns information about the newly created
        version.

        Args:
            file (IO): the file to be uploaded
            filename (str, Optional): the name of the file to be uploaded. Defaults to the
                name of the file passed in the ``file`` parameter.
            message (str, Optional): the version message that will be associated with the
                upload of the file. Defaults to a generic message if not set
            parent_version_id (uuid.UUID, Optional): If specified, SDK will check whether you are
                making changes on top of the latest version, if not the operation will fail.

        Returns:
            ``Dict[str, Any]``: A dictionary with metadata representing the version created.
        """
        if filename:
            name = filename
        else:
            name = os.path.basename(file.name)

        return self._add_single_file(file, name, message, parent_version_id)

    @typechecked
    @_ensure_not_deleted
    def push_tabular_file(
        self,
        file: IO,
        filename: Optional[str] = None,
        message: Optional[str] = None,
        parent_version_id: Optional[Union[uuid.UUID, str]] = None,
    ) -> Dict[str, Any]:
        """
        Create a new version of this dataset that contains a single new CSV file passed as the
        ``file`` parameter for this method, and returns information about the newly created
        version. The filename must be postfixed with ``.csv``.

        Args:
            file (IO): the csv file to be uploaded
            filename (str, Optional): the name of the file to be uploaded. Defaults to the local
                file name of the file. Name must end with ``.csv``.
            message (str, optional): the version message that will be associated with the
                upload of the file. Defaults to a generic message if not set
            parent_version_id (uuid.UUID, optional): If specified, SDK will check whether you are
                making changes on top of the latest version, if not the operation will fail.

        Returns:
            ``Dict[str, Any]``: A dictionary with metadata representing the version created.
        """
        if filename:
            name = filename
        else:
            name = os.path.basename(file.name)

        if name[-4:] != constants.CSV_SUFFIX:
            raise ValueError(f"filename must end with '.csv'! Got filename: {name}")

        return self._add_single_file(file, name, message, parent_version_id)

    @typechecked
    @_ensure_not_deleted
    def push_dataframe(
        self,
        dataframe: pd.DataFrame,
        csvname: Optional[str] = None,
        message: Optional[str] = None,
        parent_version_id: Optional[Union[uuid.UUID, str]] = None,
    ) -> Dict[str, Any]:
        """
        Create a new version of this dataset that contains a single new CSV created by calling
        ``to_csv(tempfile, index=False)`` on the ``pandas.DataFrame`` passed as the ``dataframe``
        parameter to this function. Optionally override the CSV name which will default to a
        random hash.

        Args:
            dataframe (pd.DataFrame): the dataframe to be uploaded
            csvname (str, Optional): the name for the csv that will be uploaded. Defaults to
                dataframe-<random-hash>.csv
            message (str, Optional): the version message that will be associated with the
                upload of the file. Defaults to a generic message if not set
            parent_version_id (uuid.UUID, Optional): If specified, SDK will check whether you are
                making changes on top of the latest version, if not the operation will fail.

        Returns:
            ``Dict[str, Any]``: A dictionary with metadata representing the version created.
        """
        if csvname:
            if csvname[-4:] == constants.CSV_SUFFIX:
                name = csvname
            else:
                name = f"{csvname}{constants.CSV_SUFFIX}"
        else:
            hash = hashlib.md5(pd.util.hash_pandas_object(dataframe).values).hexdigest()
            name = f"dataframe-{hash}{constants.CSV_SUFFIX}"

        file = io.StringIO()
        dataframe.to_csv(file, index=False)
        file.seek(0)

        return self._add_single_file(file, name, message, parent_version_id)

    @typechecked
    @_ensure_not_deleted
    @_warn_when_local_changes
    def rollback(
        self, version_id: Union[str, uuid.UUID], forced: bool = False
    ) -> Optional[Dict[str, Any]]:
        """
        Roll back the dataset to a previous version based on version ID, then pull the target
        version to the local directory specified in ``workspace``.

        Args:
            version_id (str): ID of the version to rollback to.
            forced (bool): whether to discard local changes when rolling back.

        Returns:
            Optional[Dict[str, Any]]: A dictionary with metadata representing the versioned
            rolled back to if successful, ``None`` otherwise.
        """
        target_commit = self._get_commit(version_id)

        resp = self._api_client.request(
            "POST",
            f"/api/v1/datasets/{self._dataset_id}/commits",
            json={
                "message": f"Rollback dataset to version: {target_commit['id']}",
                "metadata_s3_file_version": target_commit["metadata_s3_file_version"],
                "parent_commit_id": self._get_current_commit()["id"],
            },
            raise_for_status=True,
        )

        commit_info = resp["data"]
        self.pull(commit_info["id"], forced=True)

        return self._prune_commit_info(commit_info)

    @typechecked
    @_ensure_not_deleted
    @_warn_when_local_changes
    def pull(
        self, version_id: Optional[Union[str, uuid.UUID]] = None, forced: bool = False
    ) -> Optional[Dict[str, Any]]:
        """
        Sync local dataset folder based on a version id. If version_id was not provided, will sync
        based on the latest version. Will not allow a pull if there are local changes, to avoid
        accidentally losing changes

        Args:
            version_id (str, optional): target version ID, defaults to ``None`` which will pull the
                latest version.
            forced (bool): whether to ignore the presence of any local changes

        Returns:
            Optional[Dict[str, Any]]: A dictionary with metadata representing the versioned rolled
            back to if successful, ``None`` otherwise.
        """
        target_commit = (
            self._get_commit(commit_id=version_id) if version_id else self._get_latest_commit()
        )

        metadata_version_id = target_commit[constants.METADATA_S3_FILE_VERSION]

        self._download_files(
            [
                DatasetFileInfo(
                    file_name=constants.DATASET_MANIFEST_FILE,
                    url=self._get_s3_url(constants.DATASET_MANIFEST_FILE),
                    version_id=metadata_version_id,
                )
            ]
        )

        diff = self._get_diff_for_pull()
        self._download_files(diff[constants.MODIFIED_FILES])  # Overwrite modified files
        self._download_files(diff[constants.DELETED_FILES])  # redownload deleted files

        # delete new added files
        for f in diff[constants.NEW_FILES]:
            os.remove(self._get_file_path(f.file_name))

        self._update_head(target_commit)
        if not os.path.exists(self._get_file_path(constants.STASH_FOLDER)):
            os.mkdir(self._get_file_path(constants.STASH_FOLDER))

        # Create the tabular_manifests and artifact folder
        (Path(self.workspace) / self.dataset_name / constants.TABULAR_MANIFESTS).resolve().mkdir(
            parents=True, exist_ok=True
        )

        (Path(self.workspace) / self.dataset_name / constants.ARTIFACTS).resolve().mkdir(
            parents=True, exist_ok=True
        )

        return self._prune_commit_info(target_commit)

    def stash(self) -> None:
        """
        Stashes local changes for later use.

        Returns:
            ``None``
        """
        # deletes existing stash
        if os.path.exists(self._get_file_path(constants.DATASET_STASH_FILE)):
            with open(self._get_file_path(constants.DATASET_STASH_FILE)) as stashfile:
                old_stash = json.load(stashfile)
            for file in old_stash[constants.MODIFIED_FILES] + old_stash[constants.NEW_FILES]:
                with contextlib.suppress(FileNotFoundError):  # file not exist is OK
                    os.remove(self._get_file_path(f"{constants.STASH_FOLDER}/{file}"))

        # creates new stash
        local_diff = self.get_diff()
        with open(self._get_file_path(constants.DATASET_STASH_FILE), "w+") as stashfile:
            json.dump(local_diff, stashfile)

        # stash files
        for file in local_diff[constants.MODIFIED_FILES]:
            self._stash_file(file)
        for file in local_diff[constants.NEW_FILES]:
            self._stash_file(file)

        # re-sync to current commit
        diff = self._get_diff_for_pull()
        logger.warning(diff)
        self._download_files(diff[constants.MODIFIED_FILES])  # overwrite modified files
        self._download_files(diff[constants.DELETED_FILES])  # redownload deleted files

        # delete new added files
        for file in diff[constants.NEW_FILES]:
            os.remove(self._get_file_path(file.file_name))

    def restore(self) -> None:
        """
        Restore changes from local stash.

        Returns:
            ``None``
        """
        if not os.path.exists(self._get_file_path(constants.DATASET_STASH_FILE)):
            logger.warning("No stashed files were found!")
            return

        diff = {}
        with open(self._get_file_path(constants.DATASET_STASH_FILE)) as stashfile:
            diff = json.load(stashfile)
        os.remove(self._get_file_path(constants.DATASET_STASH_FILE))
        if not diff:
            return

        for file in diff[constants.MODIFIED_FILES]:
            self._restore_stash_file(file)
        for file in diff[constants.NEW_FILES]:
            self._restore_stash_file(file)
        for file in diff[constants.DELETED_FILES]:
            os.remove(self._get_file_path(file))

    def delete(self):
        """
        Soft delete this dataset.
        """
        self._api_client.request(
            "DELETE",
            f"/api/v1/datasets/{self._dataset_id}",
            raise_for_status=True,
        )
        logger.info(f"Dataset {self.dataset_name} has been deleted")
        return

    @typechecked
    def _get_commit(self, commit_id: Union[str, uuid.UUID]) -> Dict[str, Any]:
        """
        Get commit details

        Args:
            commit_id (str): commit id
        Returns:
            Commit information
        """
        try:
            resp = self._api_client.request(
                "GET",
                f"/api/v1/datasets/{self._dataset_id}/commits/{commit_id}",
                raise_for_status=True,
            )
        except GantryRequestException as e:
            if e.status_code == 404:
                raise DatasetCommitNotFoundError(
                    f"Could not find dataset commit with id {commit_id}"
                ) from e
            raise
        return resp["data"]

    def _get_latest_commit(self) -> Dict[str, Any]:
        """
        Get latest commit details

        Returns:
            Commit information
        """
        resp = self._api_client.request(
            "GET", f"/api/v1/datasets/{self._dataset_id}/commits", raise_for_status=True
        )
        return resp["data"][0]  # return latest commit

    def _update_manifest_file(self, diff: Dict[str, List[DatasetFileInfo]]):
        new_files = {f.file_name: f for f in diff[constants.NEW_FILES]}
        modified_files = {f.file_name: f for f in diff[constants.MODIFIED_FILES]}
        deleted_files = {f.file_name: f for f in diff[constants.DELETED_FILES]}

        with open(
            self._get_file_path(f"{constants.DATASET_MANIFEST_FILE}{constants.NEW_SUFFIX}"), "w"
        ) as new_gantry_manifest:
            with open(
                self._get_file_path(constants.DATASET_MANIFEST_FILE), "r"
            ) as cur_gantry_manifest:
                for line in cur_gantry_manifest:
                    cur_file_info = json.loads(line)
                    cur_file_name = cur_file_info[constants.FILE_NAME]

                    if cur_file_name in modified_files:  # update modified files
                        new_gantry_manifest.write(modified_files[cur_file_name].to_jsonl())
                    elif cur_file_name in deleted_files:  # remove deleted files
                        continue
                    else:  # keep unchanged files
                        new_gantry_manifest.write(line)

            for _, file_info in new_files.items():
                new_gantry_manifest.write(file_info.to_jsonl())  # add new added files

        os.replace(
            self._get_file_path(f"{constants.DATASET_MANIFEST_FILE}{constants.NEW_SUFFIX}"),
            self._get_file_path(constants.DATASET_MANIFEST_FILE),
        )

    def _get_diff_for_commit(self, commit_id: uuid.UUID):
        repo_diff = self._get_diff()
        commit_diff = {}

        # generate new url for new files
        commit_diff[constants.NEW_FILES] = [
            DatasetFileInfo(
                file_name=file_name,
                url=self._get_s3_url(file_name, commit_id),
                sha256=checksum,
            )
            for file_name, checksum in repo_diff[constants.NEW_FILES]
        ]
        # generate new url for modified files
        commit_diff[constants.MODIFIED_FILES] = [
            DatasetFileInfo(
                file_name=file_name,
                url=self._get_s3_url(file_name, commit_id),
                sha256=checksum,
            )
            for file_name, checksum in repo_diff[constants.MODIFIED_FILES]
        ]
        # don't need url for deleted files
        commit_diff[constants.DELETED_FILES] = [
            DatasetFileInfo(
                file_name=file_name,
                url="",
                sha256=checksum,
            )
            for file_name, checksum in repo_diff[constants.DELETED_FILES]
        ]

        return commit_diff

    def _get_diff_for_pull(self):
        repo_diff = self._get_diff()
        sync_diff = {}

        commit_snapshot = {}
        with open(self._get_file_path(constants.DATASET_MANIFEST_FILE)) as cur_gantry_manifest:
            for line in cur_gantry_manifest:
                file_info = json.loads(line)
                commit_snapshot[file_info[constants.FILE_NAME]] = file_info

        # don't need url for new files and use checksum for empty string
        sync_diff[constants.NEW_FILES] = [
            DatasetFileInfo(
                file_name=file_name,
                url="",
                sha256=constants.EMPTY_STR_SHA256,
            )
            for file_name, _ in repo_diff[constants.NEW_FILES]
        ]
        # retrieve url for modified files
        sync_diff[constants.MODIFIED_FILES] = [
            DatasetFileInfo(
                file_name=file_name,
                url=commit_snapshot[file_name][constants.URL],
                sha256=commit_snapshot[file_name][constants.SHA256],
                version_id=commit_snapshot[file_name][constants.VERSION_ID],
            )
            for file_name, _ in repo_diff[constants.MODIFIED_FILES]
        ]
        # retrieve url for deleted files
        sync_diff[constants.DELETED_FILES] = [
            DatasetFileInfo(
                file_name=file_name,
                url=commit_snapshot[file_name][constants.URL],
                sha256=commit_snapshot[file_name][constants.SHA256],
                version_id=commit_snapshot[file_name][constants.VERSION_ID],
            )
            for file_name, _ in repo_diff[constants.DELETED_FILES]
        ]

        return sync_diff

    def _get_diff(self):
        """
        Generate local diff.

        Args:
            return_local_info (bool): if true return local file info for modified files else
            return file info in dataset repo

        Returns:
            {
                "new_files": List[Tuple(str,str)],
                "modified_files": List[Tuple(str,str)],
                "deleted_files": List[Tuple(str,str)],
            }
        """
        repo_diff = defaultdict(list)

        commit_snapshot = {}
        with open(self._get_file_path(constants.DATASET_MANIFEST_FILE)) as cur_gantry_manifest:
            for line in cur_gantry_manifest:
                file_info = json.loads(line)
                commit_snapshot[file_info[constants.FILE_NAME]] = file_info

        local_files = get_files_checksum(Path(self.workspace) / self.dataset_name)

        # TODO:// Simplify the following logic to make it easier to maintain
        for file_name, checksum in local_files.items():
            if file_name.startswith(constants.GANTRY_FOLDER):  # skip .dataset_metadata folder
                continue
            elif file_name.startswith(constants.STASH_FOLDER):  # skip stash folder
                continue
            elif file_name not in commit_snapshot:  # new added file
                repo_diff[constants.NEW_FILES].append((file_name, checksum))
            elif checksum != commit_snapshot[file_name].get(
                constants.SHA256
            ):  # if file_name in current_snapshot --> check if it has been modified
                repo_diff[constants.MODIFIED_FILES].append((file_name, checksum))

        for file_name in commit_snapshot:
            if file_name not in local_files:
                repo_diff[constants.DELETED_FILES].append((file_name, constants.EMPTY_STR_SHA256))

        return repo_diff

    def _add_single_file(
        self,
        file: IO,
        filename: str,
        commit_msg: Optional[str],
        parent_commit_id: Optional[Union[uuid.UUID, str]],
    ):
        if parent_commit_id is not None:
            if str(parent_commit_id) != self._get_latest_commit()["id"]:
                raise DatasetHeadOutOfDateException()

        if commit_msg is None:
            commit_msg = f"Uploaded {filename} to dataset"

        data = {
            "commit_msg": commit_msg,
        }
        files = {
            "file": (filename, file),
        }

        resp = self._api_client.request(
            "POST",
            f"/api/v1/datasets/{self._dataset_id}/file",
            files=files,
            data=data,
            raise_for_status=True,
        )

        commit_info = resp["data"]
        return self._prune_commit_info(commit_info)

    def _get_file_path(self, file_name: str) -> str:
        return os.path.join(self.workspace, self.dataset_name, file_name)

    def _get_s3_url(self, file_name: str, commit_id: Optional[uuid.UUID] = None) -> str:
        if commit_id:
            return (
                f"s3://{self._bucket_name}/{self._dataset_s3_prefix}/"
                f"{self.dataset_name}/{commit_id}/{file_name}"
            )
        else:
            return (
                f"s3://{self._bucket_name}/{self._dataset_s3_prefix}/"
                f"{self.dataset_name}/{file_name}"
            )

    def _get_obj_key(self, s3_url: str) -> str:
        _, obj_key = parse_s3_path(s3_url)
        return obj_key

    def _update_head(self, commit_info: Dict[str, Any]):
        with open(self._get_file_path(constants.DATASET_HEAD_FILE), "w") as f:
            json.dump(commit_info, f)

    def _get_current_commit(self) -> Dict[str, Any]:
        with open(self._get_file_path(constants.DATASET_HEAD_FILE), "r") as f:
            return json.load(f)

    def _backup_file(self, file_path: str):
        # make a copy of the original file
        shutil.copyfile(
            self._get_file_path(file_path),
            self._get_file_path(f"{file_path}{constants.BACKUP_SUFFIX}"),
        )

    def _recover_file(self, file_path: str):
        # recover the file from a backup copy
        os.replace(
            self._get_file_path(f"{file_path}{constants.BACKUP_SUFFIX}"),
            self._get_file_path(file_path),
        )

    def _stash_file(self, file_path: str):
        # move file to stash folder
        dest_path = self._get_file_path(f"{constants.STASH_FOLDER}/{file_path}")
        os.makedirs(Path(dest_path).parent, exist_ok=True)
        os.replace(
            self._get_file_path(file_path),
            dest_path,
        )

    def _restore_stash_file(self, file_path: str):
        # recover file from stash folder
        dest_path = self._get_file_path(file_path)
        os.makedirs(Path(dest_path).parent, exist_ok=True)
        os.replace(
            self._get_file_path(f"{constants.STASH_FOLDER}/{file_path}"),
            dest_path,
        )

    def _download_files(self, file_list: List[DatasetFileInfo]):
        """
        download file to the specific path

        Args:
            file_list (List[DatasetFileInfo]):
        """
        for i in range(0, len(file_list), BATCH_SIZE):
            batch = file_list[i : i + BATCH_SIZE]
            presigned_urls = self._generate_get_presign_url(batch=batch)

            for item in batch:
                local_path = self._get_file_path(item.file_name)
                presigned_url = presigned_urls[self._get_obj_key(item.url)]
                download_file_from_url(presigned_url=presigned_url, local_path=Path(local_path))

    def _generate_get_presign_url(self, batch: List[DatasetFileInfo]):
        resp = self._api_client.request(
            "POST",
            f"/api/v1/datasets/{self._dataset_id}/presign/getobject",
            json={
                "obj_infos": [
                    {
                        constants.OBJ_KEY: self._get_obj_key(item.url),
                        constants.VERSION_ID: item.version_id,
                    }
                    for item in batch
                ],
            },
            raise_for_status=True,
        )
        return resp["data"]

    def _upload_files(self, file_list: List[DatasetFileInfo]):
        """
        download file to the specific path

        Args:
            file_list (List[DatasetFileInfo]):
        Returns:
            List[DatasetFileInfo]
        """
        for i in range(0, len(file_list), BATCH_SIZE):
            batch = file_list[i : i + BATCH_SIZE]
            resp = self._api_client.request(
                "POST",
                f"/api/v1/datasets/{self._dataset_id}/presign/putobject",
                json={"obj_keys": [self._get_obj_key(f.url) for f in batch]},
                raise_for_status=True,
            )
            presigned_urls = resp["data"]

            for item in batch:
                local_path = self._get_file_path(item.file_name)
                presigned_url = presigned_urls[self._get_obj_key(item.url)]
                version_id = upload_file_to_url(
                    presigned_url=presigned_url, local_path=Path(local_path)
                )
                item.version_id = version_id
                logger.info(f"successfully upload: {item.file_name}")

        return file_list

    def _extract_feature_type(self) -> Dict[str, Any]:
        """
        Read feature type from dataset_config and convert them to huggingface feature type
        Returns:
            A key value pair of Dict where key is the column name and value is huggingface
            feature type
        """
        config = yaml.safe_load(
            (Path(self.workspace) / self.dataset_name / constants.DATASET_CONFIG_FILE).open(
                mode="r"
            )
        )
        return self._get_feature_map_from_config(config=config)

    def pull_huggingface_dataset_by_version(
        self,
        version_id: Optional[str] = None,
        hf_dataset_name: Optional[str] = None,
        streaming=False,
    ):
        """
        This function will pull the dataset tabular files from gantry server and convert it
        to huggingface dataset. This operation will not change your local dataset folder, it
        will only load dataset tabular files exist in the version you specified into a
        huggingface dataset object.

        Args:
            version_id (Optional[str], optional): dataset version id to load. Defaults to latest.
            hf_dataset_name (Optional[str], optional): param to overwrite the huggingface
                                                dataset name. Note this param must be a valid
                                                python identifier.
            streaming (bool, default ``False``): If set to True, don't download the data files.
                                                 Instead, it streams the data progressively while
                                                 iterating on the dataset. An IterableDataset or
                                                 IterableDatasetDict is returned instead
                                                 in this case.
        """
        if not hf_dataset_name:
            hf_dataset_name = self._clean_variable_name(self.dataset_name)

        if not GantryDataset._is_valid_python_identifier(hf_dataset_name=hf_dataset_name):
            return

        manifest_presign_url = self._get_manifest_presign_url(version_id=version_id)
        tabular_files = self._get_tabular_files(manifest_presign_url=manifest_presign_url)
        features = self._extract_feature_schema_from_s3(manifest_presign_url=manifest_presign_url)
        if not features:
            raise DataSchemaMismatchError()
        training_files = []
        for i in range(0, len(tabular_files), BATCH_SIZE):
            batch = tabular_files[i : i + BATCH_SIZE]
            presigned_urls = self._generate_get_presign_url(batch=batch)
            training_files.extend(presigned_urls.values())

        if not training_files:
            raise NoTabularDataError()

        env = jinja2.Environment(
            loader=jinja2.FileSystemLoader(
                os.path.join(os.path.dirname(os.path.realpath(__file__)), "templates")
            )
        )

        template = env.get_template("load_script_template.py.jinja")

        content = template.render(
            training_file_list=training_files,
            dataset_version="1.0.0",  # Format should be x.y.z with {{x,y,z}} being digits
            dataset_name=hf_dataset_name,
            features=features,
        )
        hf_loading_script = (
            Path(self.workspace) / self.dataset_name / constants.HF_FOLDER / f"{hf_dataset_name}.py"
        )
        hf_loading_script.parent.mkdir(parents=True, exist_ok=True)
        with hf_loading_script.open("w") as f:
            f.write(content)

        try:
            hf_dataset = datasets.load_dataset(
                str(hf_loading_script), streaming=streaming, split="train"
            )
        except datasets.builder.DatasetGenerationError as e:
            raise DataSchemaMismatchError from e

        return hf_dataset

    def _get_manifest_presign_url(self, version_id: Optional[str] = None) -> str:
        target_commit = (
            self._get_commit(commit_id=version_id) if version_id else self._get_latest_commit()
        )

        metadata_version_id = target_commit[constants.METADATA_S3_FILE_VERSION]
        manifest_s3_url = self._get_s3_url(constants.DATASET_MANIFEST_FILE)

        return self._generate_get_presign_url(
            [
                DatasetFileInfo(
                    file_name=constants.DATASET_MANIFEST_FILE,
                    url=manifest_s3_url,
                    version_id=metadata_version_id,
                )
            ]
        )[self._get_obj_key(manifest_s3_url)]

    def _get_tabular_files(self, manifest_presign_url: str) -> List[DatasetFileInfo]:
        """
        Get tabular file list based on the version id

        Args:
            version_id (Optional[str], optional): dataset version. Defaults to None.

        Returns:
            List[DatasetFileInfo]: tabular file list
        """
        tabular_files = []
        for line in read_lines_from_url(manifest_presign_url):
            file_info = json.loads(line)
            if file_info[constants.FILE_NAME].startswith(constants.TABULAR_MANIFESTS) and file_info[
                constants.FILE_NAME
            ].endswith(constants.CSV_SUFFIX):
                tabular_files.append(
                    DatasetFileInfo(
                        file_name=file_info[constants.FILE_NAME],
                        url=file_info[constants.URL],
                        version_id=file_info[constants.VERSION_ID],
                    )
                )
        return tabular_files

    def _extract_feature_schema_from_s3(self, manifest_presign_url: str):
        for line in read_lines_from_url(manifest_presign_url):
            file_info = json.loads(line)
            if file_info[constants.FILE_NAME] == constants.DATASET_CONFIG_FILE:
                # read config file
                # generate the feature map
                config_presign_url = self._generate_get_presign_url(
                    [
                        DatasetFileInfo(
                            file_name=file_info[constants.FILE_NAME],
                            url=file_info[constants.URL],
                            version_id=file_info[constants.VERSION_ID],
                        )
                    ]
                )[self._get_obj_key(file_info[constants.URL])]

                # generate presign url for config file
                response = requests.get(config_presign_url)
                config = yaml.safe_load(response.text)

                return GantryDataset._get_feature_map_from_config(config=config)

    @staticmethod
    def _get_feature_map_from_config(config):
        columns = dict()
        if constants.DATASET_FEATURES_KEY in config:
            for feature_name, feature_dtype in config[constants.DATASET_FEATURES_KEY].items():
                columns[feature_name] = constants.GANTRY_2_HF_DTYPE[feature_dtype]

        if constants.DATASET_FEEDBACK_KEY in config:
            for feature_name, feature_dtype in config[constants.DATASET_FEEDBACK_KEY].items():
                columns[feature_name] = constants.GANTRY_2_HF_DTYPE[feature_dtype]

        return columns

    def get_huggingface_dataset(
        self, hf_dataset_name: Optional[str] = None, streaming: bool = False
    ):
        """
        Create a huggingface dataset based on the local dataset files

        This method will create a huggingface load script based on the dataset config and
        then load the dataset into huggingface pyarrow dataset.

        Args:
            versoion_id (Optional[str], optional): dataset version id to load. Defaults to latest.
            hf_dataset_name (Optional[str], optional): param to overwrite the huggingface
                                                dataset name. Note this param must be a valid
                                                python identifier.
            streaming (bool, default ``False``): If set to True, don't download the data files.
                                                 Instead, it streams the data progressively while
                                                 iterating on the dataset. An IterableDataset or
                                                 IterableDatasetDict is returned instead
                                                 in this case.
        """

        if not hf_dataset_name:
            hf_dataset_name = GantryDataset._clean_variable_name(self.dataset_name)

        if not GantryDataset._is_valid_python_identifier(hf_dataset_name=hf_dataset_name):
            return

        tabular_manifests_path = (
            Path(self.workspace) / self.dataset_name / constants.TABULAR_MANIFESTS
        ).resolve()

        if not os.path.exists(tabular_manifests_path) or not any(
            Path(tabular_manifests_path).iterdir()
        ):
            raise NoTabularDataError()

        env = jinja2.Environment(
            loader=jinja2.FileSystemLoader(
                os.path.join(os.path.dirname(os.path.realpath(__file__)), "templates")
            )
        )

        template = env.get_template("load_script_template.py.jinja")

        training_files = [str(p) for p in list_all_files(tabular_manifests_path.resolve())]
        features = self._extract_feature_type()

        # verify the csvs agree with the schema
        for csv_path in training_files:
            with open(csv_path) as data_csv:
                cols = next(csv.reader(data_csv))

                if set(cols) != set(features):
                    raise DataSchemaMismatchError()

        content = template.render(
            training_file_list=training_files,
            dataset_version="1.0.0",  # Format should be x.y.z with {{x,y,z}} being digits
            dataset_name=hf_dataset_name,
            features=features,
        )
        hf_loading_script = (
            Path(self.workspace) / self.dataset_name / constants.HF_FOLDER / f"{hf_dataset_name}.py"
        )
        os.makedirs(hf_loading_script.parent, exist_ok=True)
        with hf_loading_script.open("w") as f:
            f.write(content)

        try:
            hf_dataset = datasets.load_dataset(
                str(hf_loading_script), streaming=streaming, split="train"
            )
        except datasets.builder.DatasetGenerationError as e:
            raise DataSchemaMismatchError from e

        return hf_dataset

    @staticmethod
    def _is_valid_python_identifier(hf_dataset_name):
        if (
            (not hf_dataset_name)
            or (not hf_dataset_name.isidentifier())
            or iskeyword(hf_dataset_name)
        ):
            logger.error(
                "Can't get a valid hugging face dataset name, the name has to be a valid python"
                + "identifier. Please use hf_dataset_name param to set a valid name and try again."
            )
            return False
        else:
            return True

    def _prune_commit_info(self, commit_info):
        return {
            "version_id": commit_info["id"],
            "dataset": self.dataset_name,
            "message": commit_info["message"],
            # TODO:// uncomment the following line once we add sdk support for user version
            # "version_name": commit_info["user_version"],
            "created_at": commit_info["created_at"],
            "created_by": commit_info["created_by"],
            "is_latest_version": commit_info["is_latest_commit"],
        }

    @staticmethod
    def _clean_variable_name(ds_name):
        # Remove invalid characters
        ds_name = re.sub("[^0-9a-zA-Z_]", "", ds_name)
        # Remove leading characters until we find a letter or underscore
        ds_name = re.sub("^[^a-zA-Z_]+", "", ds_name)
        return ds_name
