# SPDX-License-Identifier: Apache-2.0
# Copyright 2022 The HuggingFace Authors.

import logging
from typing import List, Optional, Union

from datasets import get_dataset_split_names
from datasets.builder import ManualDownloadError
from datasets.data_files import EmptyDatasetError as _EmptyDatasetError
from libcommon.constants import (
    PROCESSING_STEP_CONFIG_SPLIT_NAMES_FROM_INFO_VERSION,
    PROCESSING_STEP_CONFIG_SPLIT_NAMES_FROM_STREAMING_VERSION,
)
from libcommon.exceptions import (
    DatasetManualDownloadError,
    EmptyDatasetError,
    SplitNamesFromStreamingError,
)

from worker.dtos import CompleteJobResult, JobRunnerInfo, SplitItem, SplitsList
from worker.job_runners.config.config_job_runner import ConfigJobRunnerWithDatasetsCache


def compute_split_names_from_streaming_response(
    dataset: str,
    config: str,
    hf_token: Optional[str] = None,
) -> SplitsList:
    """
    Get the response of config-split-names-from-streaming for one specific dataset and config on huggingface.co.
    Dataset can be private or gated if you pass an acceptable token.

    It is assumed that the dataset exists and can be accessed using the token, and that the config exists in
    the dataset.

    This function relies on the streaming mode if the splits are not directly defined in the dataset config. See
    https://github.dev/huggingface/datasets/blob/e183a269067575db8765ee979bd8523d14a1adae/src/datasets/inspect.py#L389-L390

    The config-split-names-from-streaming response generated by this function does not include stats about the split,
    like the size or number of samples. See dataset-info or dataset-size for that.

    Args:
        dataset (`str`):
            A namespace (user or an organization) and a repo name separated
            by a `/`.
        config (`str`):
            A configuration name.
        hf_token (`str`, *optional*):
            An authentication token (See https://huggingface.co/settings/token)
    Returns:
        `SplitsList`: An object with the list of split names for the dataset and config.
    Raises the following errors:
        - [`libcommon.exceptions.DatasetManualDownloadError`]:
          If the dataset requires manual download.
        - [`libcommon.exceptions.EmptyDatasetError`]
          The dataset is empty.
        - [`libcommon.exceptions.SplitsNamesError`]
          If the list of splits could not be obtained using the datasets library.
    """
    logging.info(f"get split names for dataset={dataset}, config={config}")
    use_auth_token: Union[bool, str, None] = hf_token if hf_token is not None else False

    try:
        split_name_items: List[SplitItem] = [
            {"dataset": dataset, "config": config, "split": str(split)}
            for split in get_dataset_split_names(path=dataset, config_name=config, use_auth_token=use_auth_token)
        ]
    except ManualDownloadError as err:
        raise DatasetManualDownloadError(f"{dataset=} requires manual download.", cause=err) from err
    except _EmptyDatasetError as err:
        raise EmptyDatasetError("The dataset is empty.", cause=err) from err
    except Exception as err:
        raise SplitNamesFromStreamingError(
            f"Cannot get the split names for the config '{config}' of the dataset.",
            cause=err,
        ) from err
    return SplitsList(splits=split_name_items)


class ConfigSplitNamesFromStreamingJobRunner(ConfigJobRunnerWithDatasetsCache):
    @staticmethod
    def get_job_type() -> str:
        return "config-split-names-from-streaming"

    @staticmethod
    def get_job_runner_version() -> int:
        return PROCESSING_STEP_CONFIG_SPLIT_NAMES_FROM_STREAMING_VERSION

    @staticmethod
    def get_parallel_job_runner() -> JobRunnerInfo:
        return JobRunnerInfo(
            job_runner_version=PROCESSING_STEP_CONFIG_SPLIT_NAMES_FROM_INFO_VERSION,
            job_type="config-split-names-from-info",
        )

    def compute(self) -> CompleteJobResult:
        return CompleteJobResult(
            compute_split_names_from_streaming_response(
                dataset=self.dataset,
                config=self.config,
                hf_token=self.app_config.common.hf_token,
            )
        )
