# noqa: isort:skip_file, yapf: disable
# Copyright (c) Alibaba, Inc. and its affiliates.
# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
"""ModelScope dataset loading orchestration.

This module provides :class:`DatasetsWrapperHF` and the
:func:`load_dataset_with_ctx` context manager that monkey-patch the
HuggingFace ``datasets`` library to work with the ModelScope Hub.

Sub-modules:
    _compat           – backward-compat shims for datasets>=4.0 script loading
    _module_factories – dataset module factory functions & data-file resolution
"""
import contextlib
import inspect
import os
import warnings
from dataclasses import fields
from pathlib import Path
from typing import Any, Dict, Iterable, List, Literal, Mapping, Optional, Sequence, Tuple, Union

from urllib.parse import urlencode

import requests
from datasets import (Dataset, DatasetBuilder, DatasetDict,
                      DownloadConfig, DownloadManager, DownloadMode, Features,
                      IterableDataset, IterableDatasetDict, Split,
                      VerificationMode, Version, config, data_files, LargeList,
                      Sequence as SequenceHf, SplitDict)

try:
    from datasets import List as DatasetList
except ImportError:
    DatasetList = None

from datasets.features import features
from datasets.features.features import _FEATURE_TYPES
from datasets.data_files import DataFilesDict, EmptyDatasetError
from datasets.exceptions import DataFilesNotFoundError, DatasetNotFoundError
from datasets.load import (
    CachedDatasetModuleFactory, DatasetModule,
    HubDatasetModuleFactoryWithParquetExport,
    PackagedDatasetModuleFactory,
    get_dataset_builder_class)
from datasets.packaged_modules import (_EXTENSION_TO_MODULE, _PACKAGED_DATASETS_MODULES)
from datasets.utils import file_utils
from datasets.utils.file_utils import (_raise_if_offline_mode_is_enabled,
                                       cached_path, relative_to_absolute_path)
from datasets.utils.info_utils import is_small_dataset
from datasets.utils.track import tracked_str

from huggingface_hub import hf_hub_url
from huggingface_hub.errors import OfflineModeIsEnabled
from huggingface_hub.hf_api import DatasetInfo as HfDatasetInfo
from huggingface_hub.hf_api import HfApi, RepoFile, RepoFolder
from huggingface_hub.hf_file_system import HfFileSystem

from modelscope import HubApi
from modelscope.hub.utils.utils import get_endpoint
from modelscope.msdatasets.utils.hf_file_utils import get_from_cache_ms
from modelscope.utils.config_ds import MS_DATASETS_CACHE
from modelscope.utils.constant import DEFAULT_DATASET_REVISION, REPO_TYPE_DATASET
from modelscope.utils.import_utils import has_attr_in_class
from modelscope.utils.file_utils import is_relative_path
from modelscope.utils.logger import get_logger

# -- Compat layer -----------------------------------------------------------
from modelscope.msdatasets.utils._compat import (
    _HAS_SCRIPT_LOADING,
    HubDatasetModuleFactoryWithScript,
    LocalDatasetModuleFactoryWithScript,
)

# -- Module factories --------------------------------------------------------
from modelscope.msdatasets.utils._module_factories import (
    _resolve_pattern,
    _download_repo_file,
    get_module_without_script,
    get_module_with_script,
    _compat_local_script_module,
    _compat_hub_script_module,
    _get_hub_api,
)

# Compatible with datasets 4.0+ (class name changed)
try:
    from datasets.load import (
        HubDatasetModuleFactory as HubDatasetModuleFactoryWithoutScript,
        LocalDatasetModuleFactory as LocalDatasetModuleFactoryWithoutScript)
except ImportError:
    from datasets.load import (
        HubDatasetModuleFactoryWithoutScript,
        LocalDatasetModuleFactoryWithoutScript)

logger = get_logger()


# ===================================================================
# Type definitions
# ===================================================================

ExpandDatasetProperty_T = Literal[
    'author',
    'cardData',
    'citation',
    'createdAt',
    'disabled',
    'description',
    'downloads',
    'downloadsAllTime',
    'gated',
    'lastModified',
    'likes',
    'paperswithcode_id',
    'private',
    'siblings',
    'sha',
    'tags',
]


# ===================================================================
# Feature patching (generate_from_dict_ms)
# ===================================================================

_NativeList = DatasetList if DatasetList is not None else SequenceHf


def generate_from_dict_ms(obj: Any):
    """Regenerate the nested feature object from a deserialized dict.

    This is a ModelScope-patched version of ``features.generate_from_dict``
    that handles backward compatibility for legacy ``Sequence`` types in
    datasets 4.0+ where ``Sequence`` is no longer a registered feature type.
    """
    if isinstance(obj, list):
        return [generate_from_dict_ms(value) for value in obj]
    if '_type' not in obj or isinstance(obj['_type'], dict):
        return {key: generate_from_dict_ms(value) for key, value in obj.items()}
    obj = dict(obj)
    _type = obj.pop('_type')

    if _type == 'Sequence':
        feature = obj.pop('feature')
        length = obj.get('length', -1)
        return SequenceHf(feature=generate_from_dict_ms(feature), length=length)

    class_type = _FEATURE_TYPES.get(_type, None) or globals().get(_type, None)

    if class_type is None:
        raise ValueError(f"Feature type '{_type}' not found. Available feature types: {list(_FEATURE_TYPES.keys())}")

    if class_type == LargeList:
        feature = obj.pop('feature')
        return LargeList(generate_from_dict_ms(feature), **obj)
    if _NativeList is not None and (class_type is _NativeList or issubclass(class_type, _NativeList)):
        feature = obj.pop('feature')
        return _NativeList(generate_from_dict_ms(feature), **obj)

    field_names = {f.name for f in fields(class_type)}
    return class_type(**{k: v for k, v in obj.items() if k in field_names})


# ===================================================================
# Download monkey-patch (_download_ms)
# ===================================================================

def _download_ms(self, url_or_filename: str, download_config: DownloadConfig) -> str:
    """ModelScope replacement for ``DownloadManager._download``.

    Rewrites relative paths and ``hf://`` URLs to ModelScope API endpoints.
    """
    url_or_filename = str(url_or_filename)
    if url_or_filename.startswith('hf://'):
        hf_path = url_or_filename[len('hf://'):]
        for _prefix in ('datasets/', 'models/'):
            if hf_path.startswith(_prefix):
                hf_path = hf_path[len(_prefix):]
                break
        if '@' in hf_path:
            at_idx = hf_path.index('@')
            after_at = hf_path[at_idx + 1:]
            slash_idx = after_at.find('/')
            if slash_idx == -1:
                revision = after_at
                file_path = ''
            else:
                revision = after_at[:slash_idx]
                file_path = after_at[slash_idx + 1:]
        else:
            parts = hf_path.split('/', 2)
            revision = DEFAULT_DATASET_REVISION
            file_path = parts[2] if len(parts) > 2 else ''
        params_str = urlencode({'Source': 'SDK', 'Revision': revision, 'FilePath': file_path})
        url_or_filename = self._base_path + params_str
    elif is_relative_path(url_or_filename):
        revision = DEFAULT_DATASET_REVISION
        params_str = urlencode({'Source': 'SDK', 'Revision': revision, 'FilePath': url_or_filename})
        url_or_filename = self._base_path + params_str

    out = cached_path(url_or_filename, download_config=download_config)
    out = tracked_str(out)
    out.set_origin(url_or_filename)
    return out


# ===================================================================
# HfApi monkey-patches (dataset_info, list_repo_tree, get_paths_info)
# ===================================================================

def _dataset_info(
    self,
    repo_id: str,
    *,
    revision: Optional[str] = None,
    timeout: Optional[float] = None,
    files_metadata: bool = False,
    token: Optional[Union[bool, str]] = None,
    expand: Optional[List[ExpandDatasetProperty_T]] = None,
) -> HfDatasetInfo:
    """ModelScope replacement for ``HfApi.dataset_info``."""
    repo_info_iter = self.list_repo_tree(
        repo_id=repo_id,
        path_in_repo='/',
        revision=revision,
        recursive=False,
        expand=expand,
        token=token,
        repo_type=REPO_TYPE_DATASET,
    )

    data_info = {
        'id': repo_id,
        'private': False,
        'author': repo_id.split('/')[0] if repo_id else None,
        'sha': revision,
        'lastModified': None,
        'gated': False,
        'disabled': False,
        'downloads': 0,
        'likes': 0,
        'tags': [],
        'cardData': [],
        'createdAt': None,
    }

    data_siblings = []
    for info_item in repo_info_iter:
        if isinstance(info_item, RepoFile):
            data_siblings.append(
                dict(
                    rfilename=info_item.rfilename,
                    blobId=info_item.blob_id,
                    size=info_item.size,
                )
            )
    data_info['siblings'] = data_siblings

    return HfDatasetInfo(**data_info)


# -- Repo tree cache ---------------------------------------------------------

_repo_tree_cache: Dict[tuple, List[Union[RepoFile, RepoFolder]]] = {}


def _derive_from_recursive_cache(
    repo_id: str,
    revision: str,
    path_in_repo: str,
    recursive: bool,
) -> Optional[List[Union[RepoFile, RepoFolder]]]:
    """Try to derive results from a cached recursive root listing."""
    root_key = (repo_id, revision, '/', True)
    root_cached = _repo_tree_cache.get(root_key)
    if root_cached is None:
        return None

    prefix = path_in_repo.strip('/') if path_in_repo and path_in_repo != '/' else ''
    results = []
    for item in root_cached:
        item_path = item.path
        if prefix:
            if not item_path.startswith(prefix + '/') and item_path != prefix:
                continue
            rel_path = item_path[len(prefix) + 1:] if item_path.startswith(prefix + '/') else ''
        else:
            rel_path = item_path
        if not recursive and '/' in rel_path:
            continue
        results.append(item)
    return results


def _list_repo_tree(
    self,
    repo_id: str,
    path_in_repo: Optional[str] = None,
    *,
    recursive: bool = True,
    expand: bool = False,
    revision: Optional[str] = None,
    repo_type: Optional[str] = None,
    token: Optional[Union[bool, str]] = None,
) -> Iterable[Union[RepoFile, RepoFolder]]:
    """ModelScope replacement for ``HfApi.list_repo_tree``."""
    revision = revision or DEFAULT_DATASET_REVISION
    normalized_path = path_in_repo or '/'
    cache_key = (repo_id, revision, normalized_path, recursive)

    cached = _repo_tree_cache.get(cache_key)
    if cached is not None:
        yield from cached
        return

    derived = _derive_from_recursive_cache(repo_id, revision, normalized_path, recursive)
    if derived is not None:
        _repo_tree_cache[cache_key] = derived
        yield from derived
        return

    api = _get_hub_api()
    endpoint = api.get_endpoint_for_read(
        repo_id=repo_id, repo_type=REPO_TYPE_DATASET)

    _owner, _dataset_name = repo_id.split('/')
    dataset_hub_id, _ = api.get_dataset_id_and_type(
        dataset_name=_dataset_name, namespace=_owner, endpoint=endpoint)

    results: List[Union[RepoFile, RepoFolder]] = []
    page_number = 1
    page_size = 500
    max_pages = 10000
    while page_number <= max_pages:
        try:
            dataset_files = api.get_dataset_files(
                repo_id=repo_id,
                revision=revision,
                root_path=normalized_path,
                recursive=recursive,
                page_number=page_number,
                page_size=page_size,
                endpoint=endpoint,
                dataset_hub_id=dataset_hub_id,
            )
        except Exception as e:
            logger.error(f'Get dataset: {repo_id} file list failed, message: {e}')
            break

        if not dataset_files:
            break

        for file_info_d in dataset_files:
            path_info = {
                'type': 'directory' if file_info_d['Type'] == 'tree' else 'file',
                'path': file_info_d['Path'],
                'size': file_info_d['Size'],
                'oid': file_info_d['Sha256'],
            }
            item = RepoFile(**path_info) if path_info['type'] == 'file' else RepoFolder(**path_info)
            results.append(item)
            yield item

        if len(dataset_files) < page_size:
            break
        page_number += 1

    _repo_tree_cache[cache_key] = results


def _get_paths_info(
    self,
    repo_id: str,
    paths: Union[List[str], str],
    *,
    expand: bool = False,
    revision: Optional[str] = None,
    repo_type: Optional[str] = None,
    token: Optional[Union[bool, str]] = None,
) -> List[Union[RepoFile, RepoFolder]]:
    """ModelScope replacement for ``HfApi.get_paths_info``."""
    revision = revision or DEFAULT_DATASET_REVISION
    if isinstance(paths, str):
        paths = [paths]
    paths_set = set(paths)

    # Check all available caches for matching paths
    for cache_key, cached_items in _repo_tree_cache.items():
        if cache_key[0] != repo_id or cache_key[1] != revision:
            continue
        matched = [item for item in cached_items if item.path in paths_set]
        if matched:
            return matched
        # Recursive root cache is authoritative – if paths not found, they don't exist
        if cache_key == (repo_id, revision, '/', True):
            return []

    repo_info_iter = self.list_repo_tree(
        repo_id=repo_id,
        recursive=False,
        expand=expand,
        revision=revision,
        repo_type=repo_type,
        token=token,
    )

    return [item for item in repo_info_iter if item.path in paths_set]


# ===================================================================
# HfFileSystem patches (_hf_fs_open, _hf_fs_init)
# ===================================================================

_hf_fs_open_original = None
_hf_fs_init_original = None


def _hf_fs_init_with_cookie(self, *args, endpoint=None, token=None, **kwargs):
    """HfFileSystem.__init__ wrapper that injects ModelScope cookie auth.

    ModelScope's /resolve/ endpoint authenticates via `m_session_id` cookie
    rather than the `Authorization: Bearer` header used by HuggingFace Hub.
    This wrapper ensures the cookie is included in all subsequent HTTP
    requests made by the HfFileSystem instance.
    """
    _hf_fs_init_original(self, *args, endpoint=endpoint, token=token, **kwargs)
    if token and isinstance(token, str):
        if not hasattr(self._api, 'headers') or self._api.headers is None:
            self._api.headers = {}
        self._api.headers['Cookie'] = f'm_session_id={token}'


def _hf_fs_open(self, path, mode='rb', **kwargs):
    """Wrapper for HfFileSystem._open that fixes size=0 from ModelScope API.

    The ModelScope tree API may report Size=0 for files. When HfFileSystem
    caches this, AbstractBufferedFile treats the file as empty (0 bytes).
    This wrapper detects size=0 for files opened in read mode and resolves
    the actual size via a HEAD request before creating the file object.
    """
    if mode == 'rb' and 'size' not in kwargs:
        try:
            resolved = self.resolve_path(path)
            resolved_name = resolved.unresolve()
            parent = self._parent(resolved_name)
            cached_size = None
            if parent in self.dircache:
                for entry in self.dircache[parent]:
                    if entry['name'] == resolved_name and entry.get('type') == 'file':
                        cached_size = entry.get('size', -1)
                        break
            if cached_size == 0:
                url = hf_hub_url(
                    repo_id=resolved.repo_id,
                    revision=resolved.revision,
                    filename=resolved.path_in_repo,
                    repo_type=resolved.repo_type,
                    endpoint=self.endpoint,
                )
                headers = self._api._build_hf_headers()
                resp = requests.head(url, headers=headers, allow_redirects=True, timeout=30)
                if resp.status_code == 200:
                    cl = resp.headers.get('Content-Length')
                    if cl:
                        actual_size = int(cl)
                        kwargs['size'] = actual_size
                        for entry in self.dircache.get(parent, []):
                            if entry['name'] == resolved_name:
                                entry['size'] = actual_size
                                break
        except Exception:
            pass
    return _hf_fs_open_original(self, path, mode=mode, **kwargs)


class _DryRunDownloadManager:
    """Minimal download-manager stub for split discovery without I/O.

    Returns placeholder paths for all download/extract calls, allowing
    _split_generators() to execute and return SplitGenerator objects
    (which carry split names) without triggering actual network or disk I/O.
    """

    _PLACEHOLDER = os.devnull

    def download(self, url_or_urls):
        return self._map(url_or_urls)

    def download_and_extract(self, url_or_urls):
        return self._map(url_or_urls)

    def extract(self, path_or_paths):
        return self._map(path_or_paths)

    def download_custom(self, url_or_urls, custom_download):
        return self._map(url_or_urls)

    def iter_archive(self, path):
        return iter([])

    def iter_files(self, paths):
        return iter([])

    @property
    def manual_dir(self):
        return None

    @property
    def is_streaming(self):
        return False

    def _map(self, input_):
        if isinstance(input_, dict):
            return {k: self._PLACEHOLDER for k in input_}
        if isinstance(input_, (list, tuple, set)):
            return type(input_)(self._PLACEHOLDER for _ in input_)
        return self._PLACEHOLDER


def _discover_splits_from_builder(builder_instance):
    """Discover available split names by dry-running _split_generators().

    For script-based datasets that lack README split metadata, this calls
    the builder's _split_generators() with a stub download manager to
    extract split names without performing any actual downloads.

    Returns:
        A set of split name strings, or an empty set if discovery fails.
    """
    if not hasattr(builder_instance, '_split_generators'):
        return set()
    try:
        generators = builder_instance._split_generators(_DryRunDownloadManager())
        return {str(sg.name) for sg in generators}
    except Exception as e:
        logger.debug(f'Failed to discover splits from builder: {e}')
        return set()


def _validate_split_exists(builder_instance, split):
    """Fail-fast check: raise ValueError before downloading if the
    requested split does not exist in the dataset metadata.

    Args:
        builder_instance: The DatasetBuilder instance with info/config.
        split: The user-requested split specification (may be None).

    Raises:
        ValueError: If any requested split name is not found among
            the available splits declared in the dataset metadata.
    """
    if split is None:
        return

    from modelscope.msdatasets.utils._module_factories import _extract_split_names
    split_names = _extract_split_names(split)
    if not split_names:
        return

    # Source 1: info.splits (original metadata from README)
    available = set()
    info = getattr(builder_instance, 'info', None)
    if info is not None and info.splits:
        available = set(info.splits.keys())

    # Source 2: config.data_files keys
    if not available:
        config = getattr(builder_instance, 'config', None)
        data_files = getattr(config, 'data_files', None)
        if isinstance(data_files, dict):
            available = set(data_files.keys())

    # Source 3: dry-run _split_generators() for script-based datasets
    if not available:
        available = _discover_splits_from_builder(builder_instance)

    if not available:
        logger.debug(
            'Cannot determine available splits from dataset metadata; '
            'split validation skipped. Invalid splits will be caught downstream.'
        )
        return

    missing = split_names - available
    if missing:
        raise ValueError(
            f'Split {sorted(missing)} not found in dataset. '
            f'Available splits: {sorted(available)}'
        )


def _align_builder_splits_with_data_files(builder_instance, split):
    """Align builder.info.splits with the actually requested split(s).

    When data_files have been filtered to a subset of splits (see
    _filter_data_files_by_split in _module_factories.py), the builder's
    info.splits metadata may still list all original splits from the
    README.  download_and_prepare() calls verify_splits() which would
    then raise ExpectedMoreSplitsError.  This helper prunes info.splits
    to only contain the splits that will actually be generated.
    """
    if split is None:
        return
    info = getattr(builder_instance, 'info', None)
    if info is None or info.splits is None:
        return

    from modelscope.msdatasets.utils._module_factories import _extract_split_names
    split_names = _extract_split_names(split)
    if not split_names:
        return

    existing_keys = set(info.splits.keys())
    if split_names >= existing_keys:
        return  # All splits requested, no filtering needed

    filtered = {k: v for k, v in info.splits.items() if k in split_names}
    if not filtered:
        return  # Safety: don't empty out splits

    info.splits = SplitDict(filtered, dataset_name=info.splits.dataset_name)


# ===================================================================
# DatasetsWrapperHF
# ===================================================================

class DatasetsWrapperHF:

    @staticmethod
    def load_dataset(
        path: str,
        name: Optional[str] = None,
        data_dir: Optional[str] = None,
        data_files: Optional[Union[str, Sequence[str],
                                   Mapping[str, Union[str,
                                                      Sequence[str]]]]] = None,
        split: Optional[Union[str, Split]] = None,
        cache_dir: Optional[str] = None,
        features: Optional[Features] = None,
        download_config: Optional[DownloadConfig] = None,
        download_mode: Optional[Union[DownloadMode, str]] = None,
        verification_mode: Optional[Union[VerificationMode, str]] = None,
        keep_in_memory: Optional[bool] = None,
        save_infos: bool = False,
        revision: Optional[Union[str, Version]] = None,
        token: Optional[Union[bool, str]] = None,
        use_auth_token='deprecated',
        task='deprecated',
        streaming: bool = False,
        num_proc: Optional[int] = None,
        storage_options: Optional[Dict] = None,
        trust_remote_code: bool = False,
        dataset_info_only: Optional[bool] = False,
        **config_kwargs,
    ) -> Union[DatasetDict, Dataset, IterableDatasetDict, IterableDataset,
               dict]:

        if use_auth_token != 'deprecated':
            warnings.warn(
                "'use_auth_token' was deprecated in favor of 'token' in version 2.14.0 and will be removed in 3.0.0.\n"
                "You can remove this warning by passing 'token=<use_auth_token>' instead.",
                FutureWarning,
            )
            token = use_auth_token
        if task != 'deprecated':
            warnings.warn(
                "'task' was deprecated in version 2.13.0 and will be removed in 3.0.0.\n",
                FutureWarning,
            )
        else:
            task = None
        if data_files is not None and not data_files:
            raise ValueError(
                f"Empty 'data_files': '{data_files}'. It should be either non-empty or None (default)."
            )
        if Path(path, config.DATASET_STATE_JSON_FILENAME).exists():
            raise ValueError(
                'You are trying to load a dataset that was saved using `save_to_disk`. '
                'Please use `load_from_disk` instead.')

        if streaming and num_proc is not None:
            raise NotImplementedError(
                'Loading a streaming dataset in parallel with `num_proc` is not implemented. '
                'To parallelize streaming, you can wrap the dataset with a PyTorch DataLoader '
                'using `num_workers` > 1 instead.')

        download_mode = DownloadMode(download_mode
                                     or DownloadMode.REUSE_DATASET_IF_EXISTS)
        verification_mode = VerificationMode((
            verification_mode or VerificationMode.BASIC_CHECKS
        ) if not save_infos else VerificationMode.ALL_CHECKS)

        if trust_remote_code:
            logger.warning(f'Use trust_remote_code=True. Will invoke codes from {path}. Please make sure '
                           'that you can trust the external codes.')

        builder_instance = DatasetsWrapperHF.load_dataset_builder(
            path=path,
            name=name,
            data_dir=data_dir,
            data_files=data_files,
            cache_dir=cache_dir,
            features=features,
            download_config=download_config,
            download_mode=download_mode,
            revision=revision,
            token=token,
            storage_options=storage_options,
            trust_remote_code=trust_remote_code,
            _require_default_config_name=name is None,
            split=split,
            **config_kwargs,
        )

        if dataset_info_only:
            ret_dict = {}

            # Case 1: Local .py script file
            if isinstance(path, str) and path.endswith('.py') and os.path.exists(path):
                from datasets import get_dataset_config_names
                subset_list = get_dataset_config_names(path)
                ret_dict = {_subset: [] for _subset in subset_list}
                return ret_dict

            if builder_instance is None:
                logger.error(f'No builder instance created for {path} dataset.')
                return ret_dict

            # Case 2: Try builder_configs with data_files (packaged datasets)
            _tmp_builder_configs = getattr(builder_instance, 'builder_configs', None)
            if _tmp_builder_configs:
                if hasattr(_tmp_builder_configs, 'items'):
                    configs_iter = _tmp_builder_configs.items()
                else:
                    configs_iter = [(getattr(c, 'name', 'default'), c) for c in _tmp_builder_configs]
                for tmp_config_name, tmp_builder_config in configs_iter:
                    tmp_config_name = str(tmp_config_name)
                    if hasattr(tmp_builder_config, 'data_files') and tmp_builder_config.data_files is not None:
                        ret_dict[tmp_config_name] = [str(item) for item in list(tmp_builder_config.data_files.keys())]

            # Case 3: Fallback for script datasets — use info.splits or dry-run discovery
            if not ret_dict or all(not v for v in ret_dict.values()):
                config_name = getattr(builder_instance, 'config_name', 'default') or 'default'
                splits = []

                # Try info.splits (from README metadata)
                info = getattr(builder_instance, 'info', None)
                if info is not None and info.splits:
                    splits = sorted(info.splits.keys())

                # Fallback: dry-run _split_generators()
                if not splits:
                    discovered = _discover_splits_from_builder(builder_instance)
                    if discovered:
                        splits = sorted(discovered)

                if splits:
                    ret_dict = {config_name: splits}
                elif not ret_dict:
                    ret_dict = {config_name: []}

            return ret_dict

        _validate_split_exists(builder_instance, split)

        if streaming:
            return builder_instance.as_streaming_dataset(split=split)

        _align_builder_splits_with_data_files(builder_instance, split)

        builder_instance.download_and_prepare(
            download_config=download_config,
            download_mode=download_mode,
            verification_mode=verification_mode,
            num_proc=num_proc,
            storage_options=storage_options,
        )

        keep_in_memory = (
            keep_in_memory if keep_in_memory is not None else is_small_dataset(
                builder_instance.info.dataset_size))
        ds = builder_instance.as_dataset(
            split=split,
            verification_mode=verification_mode,
            in_memory=keep_in_memory)
        if task is not None:
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', FutureWarning)
                ds = ds.prepare_for_task(task)
        if save_infos:
            builder_instance._save_infos()

        try:
            api = _get_hub_api()
            if is_relative_path(path) and path.count('/') == 1:
                _namespace, _dataset_name = path.split('/')
                endpoint = api.get_endpoint_for_read(
                    repo_id=path, repo_type=REPO_TYPE_DATASET)
                api.dataset_download_statistics(dataset_name=_dataset_name, namespace=_namespace, endpoint=endpoint)
        except Exception as e:
            logger.warning(f'Could not record download statistics: {e}')

        return ds

    @staticmethod
    def load_dataset_builder(
        path: str,
        name: Optional[str] = None,
        data_dir: Optional[str] = None,
        data_files: Optional[Union[str, Sequence[str],
                                   Mapping[str, Union[str,
                                                      Sequence[str]]]]] = None,
        cache_dir: Optional[str] = None,
        features: Optional[Features] = None,
        download_config: Optional[DownloadConfig] = None,
        download_mode: Optional[Union[DownloadMode, str]] = None,
        revision: Optional[Union[str, Version]] = None,
        token: Optional[Union[bool, str]] = None,
        use_auth_token='deprecated',
        storage_options: Optional[Dict] = None,
        trust_remote_code: Optional[bool] = None,
        _require_default_config_name=True,
        split: Optional[Union[str, Split]] = None,
        **config_kwargs,
    ) -> DatasetBuilder:

        if use_auth_token != 'deprecated':
            warnings.warn(
                "'use_auth_token' was deprecated in favor of 'token' in version 2.14.0 and will be removed in 3.0.0.\n"
                "You can remove this warning by passing 'token=<use_auth_token>' instead.",
                FutureWarning,
            )
            token = use_auth_token
        download_mode = DownloadMode(download_mode
                                     or DownloadMode.REUSE_DATASET_IF_EXISTS)
        if token is not None:
            download_config = download_config.copy(
            ) if download_config else DownloadConfig()
            download_config.token = token
        if storage_options is not None:
            download_config = download_config.copy(
            ) if download_config else DownloadConfig()
            download_config.storage_options.update(storage_options)
        if split is not None:
            download_config = download_config.copy(
            ) if download_config else DownloadConfig()
            if download_config.storage_options is None:
                download_config.storage_options = {}
            download_config.storage_options['split'] = split

        dataset_module = DatasetsWrapperHF.dataset_module_factory(
            path,
            revision=revision,
            download_config=download_config,
            download_mode=download_mode,
            data_dir=data_dir,
            data_files=data_files,
            cache_dir=cache_dir,
            trust_remote_code=trust_remote_code,
            _require_default_config_name=_require_default_config_name,
            _require_custom_configs=bool(config_kwargs),
            name=name,
        )
        builder_kwargs = dataset_module.builder_kwargs
        data_dir = builder_kwargs.pop('data_dir', data_dir)
        data_files = builder_kwargs.pop('data_files', data_files)
        config_name = builder_kwargs.pop(
            'config_name', name
            or dataset_module.builder_configs_parameters.default_config_name)
        dataset_name = builder_kwargs.pop('dataset_name', None)
        info = dataset_module.dataset_infos.get(
            config_name) if dataset_module.dataset_infos else None

        if (path in _PACKAGED_DATASETS_MODULES and data_files is None
                and dataset_module.builder_configs_parameters.
                builder_configs[0].data_files is None):
            error_msg = f'Please specify the data files or data directory to load for the {path} dataset builder.'
            example_extensions = [
                extension for extension in _EXTENSION_TO_MODULE
                if _EXTENSION_TO_MODULE[extension] == path
            ]
            if example_extensions:
                error_msg += f'\nFor example `data_files={{"train": "path/to/data/train/*.{example_extensions[0]}"}}`'
            raise ValueError(error_msg)

        builder_cls = get_dataset_builder_class(
            dataset_module, dataset_name=dataset_name)

        _config_cls = builder_cls.BUILDER_CONFIG_CLASS
        if hasattr(_config_cls, '__dataclass_fields__'):
            _valid_fields = set(_config_cls.__dataclass_fields__.keys())
            # Also preserve parameters accepted by the builder's
            # __init__ (e.g. writer_batch_size, base_path, repo_id)
            # so they are not inadvertently stripped.
            try:
                _init_params = set(
                    inspect.signature(builder_cls.__init__).parameters.keys()
                )
            except (ValueError, TypeError):
                _init_params = set()
            _valid_fields = _valid_fields | _init_params
            config_kwargs = {
                k: v for k, v in config_kwargs.items()
                if k in _valid_fields
            }

        builder_instance: DatasetBuilder = builder_cls(
            cache_dir=cache_dir,
            dataset_name=dataset_name,
            config_name=config_name,
            data_dir=data_dir,
            data_files=data_files,
            hash=dataset_module.hash,
            info=info,
            features=features,
            token=token,
            storage_options=storage_options,
            **builder_kwargs,
            **config_kwargs,
        )
        builder_instance._use_legacy_cache_dir_if_possible(dataset_module)

        return builder_instance

    @staticmethod
    def dataset_module_factory(
        path: str,
        revision: Optional[Union[str, Version]] = None,
        download_config: Optional[DownloadConfig] = None,
        download_mode: Optional[Union[DownloadMode, str]] = None,
        dynamic_modules_path: Optional[str] = None,
        data_dir: Optional[str] = None,
        data_files: Optional[Union[Dict, List, str, DataFilesDict]] = None,
        cache_dir: Optional[str] = None,
        trust_remote_code: Optional[bool] = None,
        _require_default_config_name=True,
        _require_custom_configs=False,
        **download_kwargs,
    ) -> DatasetModule:

        subset_name: str = download_kwargs.pop('name', None)
        revision = revision or DEFAULT_DATASET_REVISION
        if download_config is None:
            download_config = DownloadConfig(**download_kwargs)
        download_config.storage_options.update({'name': subset_name})
        download_config.storage_options.update({'revision': revision})

        if download_config and download_config.cache_dir is None:
            download_config.cache_dir = MS_DATASETS_CACHE

        download_mode = DownloadMode(download_mode
                                     or DownloadMode.REUSE_DATASET_IF_EXISTS)
        download_config.extract_compressed_file = True
        download_config.force_extract = True
        download_config.force_download = download_mode == DownloadMode.FORCE_REDOWNLOAD

        filename = list(
            filter(lambda x: x,
                   path.replace(os.sep, '/').split('/')))[-1]
        if not filename.endswith('.py'):
            filename = filename + '.py'
        combined_path = os.path.join(path, filename)

        # Try packaged
        if path in _PACKAGED_DATASETS_MODULES:
            return PackagedDatasetModuleFactory(
                path,
                data_dir=data_dir,
                data_files=data_files,
                download_config=download_config,
                download_mode=download_mode,
            ).get_module()
        # Try locally with script
        elif path.endswith(filename):
            if os.path.isfile(path):
                if _HAS_SCRIPT_LOADING:
                    return LocalDatasetModuleFactoryWithScript(
                        path,
                        download_mode=download_mode,
                        dynamic_modules_path=dynamic_modules_path,
                        trust_remote_code=trust_remote_code,
                    ).get_module()
                return _compat_local_script_module(
                    path,
                    download_mode=download_mode,
                    dynamic_modules_path=dynamic_modules_path,
                    trust_remote_code=trust_remote_code,
                )
            else:
                raise FileNotFoundError(
                    f"Couldn't find a dataset script at {relative_to_absolute_path(path)}"
                )
        elif os.path.isfile(combined_path):
            if _HAS_SCRIPT_LOADING:
                return LocalDatasetModuleFactoryWithScript(
                    combined_path,
                    download_mode=download_mode,
                    dynamic_modules_path=dynamic_modules_path,
                    trust_remote_code=trust_remote_code,
                ).get_module()
            return _compat_local_script_module(
                combined_path,
                download_mode=download_mode,
                dynamic_modules_path=dynamic_modules_path,
                trust_remote_code=trust_remote_code,
            )
        elif os.path.isdir(path):
            return LocalDatasetModuleFactoryWithoutScript(
                path,
                data_dir=data_dir,
                data_files=data_files,
                download_mode=download_mode).get_module()
        # Try remotely
        elif is_relative_path(path) and path.count('/') == 1:
            try:
                _raise_if_offline_mode_is_enabled()

                try:
                    dataset_info = HfApi().dataset_info(
                        repo_id=path,
                        revision=revision,
                        token=download_config.token,
                        timeout=100.0,
                    )
                except Exception as e:  # noqa: broad exception from hf_hub
                    if isinstance(
                            e,
                        (  # noqa: E131
                            OfflineModeIsEnabled,
                            requests.exceptions.ConnectTimeout,
                            requests.exceptions.ConnectionError,
                        ),
                    ):
                        raise ConnectionError(
                            f"Couldn't reach '{path}' on the Hub ({type(e).__name__})"
                        )
                    elif '404' in str(e):
                        msg = f"Dataset '{path}' doesn't exist on the Hub"
                        raise DatasetNotFoundError(
                            msg
                            + f" at revision '{revision}'" if revision else msg
                        )
                    elif '401' in str(e):
                        msg = f"Dataset '{path}' doesn't exist on the Hub"
                        msg = msg + f" at revision '{revision}'" if revision else msg
                        raise DatasetNotFoundError(
                            msg + '. If the repo is private or gated, '
                            'make sure to log in with `huggingface-cli login`.'
                        )
                    else:
                        raise e

                dataset_readme_path = _download_repo_file(
                    repo_id=path,
                    path_in_repo='README.md',
                    download_config=download_config,
                    revision=revision,
                )
                commit_hash = os.path.basename(os.path.dirname(dataset_readme_path))

                if filename in [
                        sibling.rfilename for sibling in dataset_info.siblings
                ]:
                    can_load_config_from_parquet_export = False
                    if config.USE_PARQUET_EXPORT and can_load_config_from_parquet_export:
                        try:
                            if has_attr_in_class(HubDatasetModuleFactoryWithParquetExport, 'revision'):
                                return HubDatasetModuleFactoryWithParquetExport(
                                    path,
                                    revision=revision,
                                    download_config=download_config).get_module()

                            return HubDatasetModuleFactoryWithParquetExport(
                                path,
                                commit_hash=commit_hash,
                                download_config=download_config).get_module()
                        except Exception as e:
                            logger.error(e)

                    if _HAS_SCRIPT_LOADING:
                        if has_attr_in_class(HubDatasetModuleFactoryWithScript, 'revision'):
                            return HubDatasetModuleFactoryWithScript(
                                path,
                                revision=revision,
                                download_config=download_config,
                                download_mode=download_mode,
                                dynamic_modules_path=dynamic_modules_path,
                                trust_remote_code=trust_remote_code,
                            ).get_module()

                        return HubDatasetModuleFactoryWithScript(
                            path,
                            commit_hash=commit_hash,
                            download_config=download_config,
                            download_mode=download_mode,
                            dynamic_modules_path=dynamic_modules_path,
                            trust_remote_code=trust_remote_code,
                        ).get_module()

                    return _compat_hub_script_module(
                        path,
                        revision=revision,
                        download_config=download_config,
                        download_mode=download_mode,
                        dynamic_modules_path=dynamic_modules_path,
                        trust_remote_code=trust_remote_code,
                    )
                else:
                    if has_attr_in_class(HubDatasetModuleFactoryWithoutScript, 'revision'):
                        return HubDatasetModuleFactoryWithoutScript(
                            path,
                            revision=revision,
                            data_dir=data_dir,
                            data_files=data_files,
                            download_config=download_config,
                            download_mode=download_mode,
                        ).get_module()

                    return HubDatasetModuleFactoryWithoutScript(
                        path,
                        commit_hash=commit_hash,
                        data_dir=data_dir,
                        data_files=data_files,
                        download_config=download_config,
                        download_mode=download_mode,
                    ).get_module()
            except Exception as e1:
                logger.error(f'>> Error loading {path}: {e1}')

                try:
                    _cached_factory_kwargs = {'cache_dir': cache_dir}
                    if _HAS_SCRIPT_LOADING:
                        _cached_factory_kwargs['dynamic_modules_path'] = dynamic_modules_path
                    return CachedDatasetModuleFactory(
                        path, **_cached_factory_kwargs).get_module()
                except Exception:
                    if isinstance(e1, OfflineModeIsEnabled):
                        raise ConnectionError(
                            f"Couldn't reach the Hugging Face Hub for dataset '{path}': {e1}"
                        ) from None
                    if isinstance(e1,
                                  (DataFilesNotFoundError,
                                   DatasetNotFoundError, EmptyDatasetError)):
                        raise e1 from None
                    if isinstance(e1, FileNotFoundError):
                        raise FileNotFoundError(
                            f"Couldn't find a dataset script at {relative_to_absolute_path(combined_path)} or "
                            f'any data file in the same directory. '
                            f"Couldn't find '{path}' on the Hugging Face Hub either: {type(e1).__name__}: {e1}"
                        ) from None
                    raise e1 from None
        else:
            raise FileNotFoundError(
                f"Couldn't find a dataset script at {relative_to_absolute_path(combined_path)} or "
                f'any data file in the same directory.')


# ===================================================================
# Context manager – load_dataset_with_ctx
# ===================================================================

@contextlib.contextmanager
def load_dataset_with_ctx(*args, **kwargs):
    """Context manager that monkey-patches ``datasets`` to use ModelScope.

    All monkey-patches are applied on entry and restored on exit (for
    non-streaming mode) or kept alive (for streaming mode, where lazy
    iteration needs the patches to remain active).
    """
    global _hf_fs_open_original, _hf_fs_init_original

    # Save originals
    hf_endpoint_origin = config.HF_ENDPOINT
    get_from_cache_origin = file_utils.get_from_cache
    _download_origin = DownloadManager._download if hasattr(DownloadManager, '_download') \
        else DownloadManager._download_single
    dataset_info_origin = HfApi.dataset_info
    list_repo_tree_origin = HfApi.list_repo_tree
    get_paths_info_origin = HfApi.get_paths_info
    resolve_pattern_origin = data_files.resolve_pattern
    get_module_without_script_origin = HubDatasetModuleFactoryWithoutScript.get_module
    get_module_with_script_origin = (
        HubDatasetModuleFactoryWithScript.get_module if _HAS_SCRIPT_LOADING else None)
    generate_from_dict_origin = features.generate_from_dict
    hf_fs_open_origin = HfFileSystem._open
    hf_fs_init_origin = HfFileSystem.__init__
    hf_fs_open_was_patched = hf_fs_open_origin is _hf_fs_open
    hf_fs_init_was_patched = hf_fs_init_origin is _hf_fs_init_with_cookie

    # Apply patches
    config.HF_ENDPOINT = get_endpoint()
    file_utils.get_from_cache = get_from_cache_ms
    if hasattr(DownloadManager, '_download'):
        DownloadManager._download = _download_ms
    else:
        DownloadManager._download_single = _download_ms
    HfApi.dataset_info = _dataset_info
    HfApi.list_repo_tree = _list_repo_tree
    HfApi.get_paths_info = _get_paths_info
    data_files.resolve_pattern = _resolve_pattern
    HubDatasetModuleFactoryWithoutScript.get_module = get_module_without_script
    if _HAS_SCRIPT_LOADING:
        HubDatasetModuleFactoryWithScript.get_module = get_module_with_script
    features.generate_from_dict = generate_from_dict_ms
    if not hf_fs_open_was_patched:
        _hf_fs_open_original = hf_fs_open_origin
        HfFileSystem._open = _hf_fs_open
    if not hf_fs_init_was_patched:
        _hf_fs_init_original = hf_fs_init_origin
        HfFileSystem.__init__ = _hf_fs_init_with_cookie

    streaming = kwargs.get('streaming', False)

    _streaming_dataset_returned = False

    try:
        dataset_res = DatasetsWrapperHF.load_dataset(*args, **kwargs)
        _streaming_dataset_returned = streaming
        yield dataset_res
    finally:
        _repo_tree_cache.clear()
        HubApi._dataset_id_type_cache.clear()

        should_restore = not _streaming_dataset_returned
        if should_restore:
            if not hf_fs_open_was_patched:
                HfFileSystem._open = hf_fs_open_origin
                _hf_fs_open_original = None
            if not hf_fs_init_was_patched:
                HfFileSystem.__init__ = hf_fs_init_origin
                _hf_fs_init_original = None

            config.HF_ENDPOINT = hf_endpoint_origin
            file_utils.get_from_cache = get_from_cache_origin
            features.generate_from_dict = generate_from_dict_origin

            if hasattr(DownloadManager, '_download'):
                DownloadManager._download = _download_origin
            else:
                DownloadManager._download_single = _download_origin

            HfApi.dataset_info = dataset_info_origin
            HfApi.list_repo_tree = list_repo_tree_origin
            HfApi.get_paths_info = get_paths_info_origin
            data_files.resolve_pattern = resolve_pattern_origin
            HubDatasetModuleFactoryWithoutScript.get_module = get_module_without_script_origin
            if _HAS_SCRIPT_LOADING:
                HubDatasetModuleFactoryWithScript.get_module = get_module_with_script_origin
