# Copyright NVIDIA Corporation 2023
# SPDX-License-Identifier: Apache-2.0

import fnmatch
import io
import json
import re
import tarfile
from functools import partial
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Union

from ray._common.utils import env_bool
from ray.data._internal.util import iterate_with_retry
from ray.data.block import Block, BlockAccessor
from ray.data.datasource.file_based_datasource import FileBasedDatasource

ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR = (
    "RAY_DATA_WEBDATASET_ALLOW_UNSAFE_DESERIALIZATION"
)


if TYPE_CHECKING:
    import pyarrow


def _base_plus_ext(path: str):
    """Split off all file extensions.

    Returns base, allext.

    Args:
        path: path with extensions

    Returns:
        str: path with all extensions removed
    """
    match = re.match(r"^((?:.*/|)[^.]+)[.]([^/]*)$", path)
    if not match:
        return None, None
    return match.group(1), match.group(2)


def _valid_sample(sample: Dict[str, Any]):
    """Check whether a sample is valid.

    Args:
        sample: sample to be checked

    Returns:
        ``True`` if the sample is a non-empty dict without the ``__bad__`` flag.
    """
    return (
        sample is not None
        and isinstance(sample, dict)
        and len(list(sample.keys())) > 0
        and not sample.get("__bad__", False)
    )


def _apply_list(
    f: Union[Callable, List[Callable]], sample: Dict[str, Any], default: Callable = None
):
    """Apply a list of functions to a sample.

    Args:
        f: function or list of functions
        sample: sample to be modified
        default: default function to be applied to all keys.
            Defaults to None.

    Returns:
        modified sample
    """
    if f is None:
        return sample
    if not isinstance(f, list):
        f = [f]
    for g in f:
        if default is not None and not callable(g):
            g = partial(default, format=g)
        sample = g(sample)
    return sample


def _check_suffix(suffix: str, suffixes: Union[list, callable]):
    """Check whether a suffix is valid.

    Suffixes can be either None (=accept everything), a callable,
    or a list of patterns. If the pattern contains */? it is treated
    as a glob pattern, otherwise it is treated as a literal.

    Args:
        suffix: suffix to be checked
        suffixes: list of valid suffixes

    Returns:
        ``True`` if the suffix matches the allowed patterns.
    """
    if suffixes is None:
        return True
    if callable(suffixes):
        return suffixes(suffix)
    for pattern in suffixes:
        if "*" in pattern or "?" in pattern:
            if fnmatch.fnmatch("." + suffix, pattern):
                return True
        elif suffix == pattern or "." + suffix == pattern:
            return True
    return False


def _tar_file_iterator(
    fileobj: Any,
    fileselect: Optional[Union[bool, callable, list]] = None,
    filerename: Optional[Union[bool, callable, list]] = None,
    verbose_open: bool = False,
    meta: dict = None,
) -> Iterator[Dict[str, Any]]:
    """Iterate over tar file, yielding filename, content pairs for the given tar stream.

    Args:
        fileobj: file object
        fileselect: patterns or function selecting
            files to be selected
        filerename: patterns or function used to rename selected files
            before yielding them.
        verbose_open: if ``True``, print progress messages when starting
            and finishing iteration over the tar stream.
        meta: metadata to be added to each sample

    Yields:
        Dict[str, Any]: Dictionaries with ``fname`` and ``data`` keys for each
        selected file in the tar stream.
    """
    meta = meta or {}
    stream = tarfile.open(fileobj=fileobj, mode="r|*")
    if verbose_open:
        print(f"start {meta}")
    for tarinfo in stream:
        fname = tarinfo.name
        if not tarinfo.isreg() or fname is None:
            continue
        data = stream.extractfile(tarinfo).read()
        fname = _apply_list(filerename, fname)
        assert isinstance(fname, str)
        if not _check_suffix(fname, fileselect):
            continue
        result = dict(fname=fname, data=data)
        yield result
    if verbose_open:
        print(f"done {meta}")


def _group_by_keys(
    data: List[Dict[str, Any]],
    keys: callable = _base_plus_ext,
    suffixes: Optional[Union[list, callable]] = None,
    meta: dict = None,
) -> Iterator[Dict[str, Any]]:
    """Return function over iterator that groups key, value pairs into samples.

    Args:
        data: iterator over key, value pairs
        keys: function that returns key, suffix for a given key
        suffixes: list of suffixes to be included in the sample
        meta: metadata to be added to each sample

    Yields:
        Dict[str, Any]: Grouped samples, where files sharing the same key prefix are
        combined into a single dictionary.
    """
    meta = meta or {}
    current_sample = None
    for filesample in data:
        assert isinstance(filesample, dict)
        fname, value = filesample["fname"], filesample["data"]
        prefix, suffix = keys(fname)
        if prefix is None:
            continue
        if current_sample is None or prefix != current_sample["__key__"]:
            if _valid_sample(current_sample):
                current_sample.update(meta)
                yield current_sample
            current_sample = dict(__key__=prefix)
            if "__url__" in filesample:
                current_sample["__url__"] = filesample["__url__"]
        if suffix in current_sample:
            raise ValueError(
                f"{fname}: duplicate file name in tar file "
                + f"{suffix} {current_sample.keys()}, tar is {meta['__url__']}"
            )
        if suffixes is None or _check_suffix(suffix, suffixes):
            current_sample[suffix] = value
    if _valid_sample(current_sample):
        current_sample.update(meta)
        yield current_sample


def _default_decoder(
    sample: Dict[str, Any],
    format: Optional[Union[bool, str]] = True,
    allow_unsafe: bool = False,
):
    """A default decoder for webdataset.

    This handles common file extensions: .txt, .cls, .cls2,
        .jpg, .png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl.
    These are the most common extensions used in webdataset.
    For other extensions, users can provide their own decoder.

    Args:
        sample: sample, modified in place
        format: optional image format hint (e.g. ``"PIL"`` to return PIL
            images instead of numpy arrays).
        allow_unsafe: if True, allow pickle/torch deserialization

    Returns:
        The sample with values decoded according to their key extension.
    """
    sample = dict(sample)
    for key, value in sample.items():
        extension = key.split(".")[-1]
        if key.startswith("__"):
            continue
        elif extension in ["txt", "text"]:
            sample[key] = value.decode("utf-8")
        elif extension in ["cls", "cls2"]:
            sample[key] = int(value.decode("utf-8"))
        elif extension in ["jpg", "png", "ppm", "pgm", "pbm", "pnm"]:
            import numpy as np
            import PIL.Image

            if format == "PIL":
                sample[key] = PIL.Image.open(io.BytesIO(value))
            else:
                sample[key] = np.asarray(PIL.Image.open(io.BytesIO(value)))
        elif extension == "json":
            sample[key] = json.loads(value)
        elif extension == "npy":
            import numpy as np

            sample[key] = np.load(io.BytesIO(value))
        elif extension == "mp":
            import msgpack

            sample[key] = msgpack.unpackb(value, raw=False)
        elif extension in ["pt", "pth"]:
            if not allow_unsafe:
                raise ValueError(
                    f"Refusing to load .{extension} member {key!r} from "
                    f"WebDataset with weights_only=False (arbitrary code "
                    f"execution risk). Provide a custom decoder or set "
                    f"{ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR}=1 "
                    f"for trusted sources."
                )
            import torch

            sample[key] = torch.load(io.BytesIO(value), weights_only=False)
        elif extension in ["pickle", "pkl"]:
            if not allow_unsafe:
                raise ValueError(
                    f"Refusing to unpickle WebDataset member {key!r} "
                    f"(arbitrary code execution risk). Provide a custom "
                    f"decoder or set "
                    f"{ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR}=1 "
                    f"for trusted sources."
                )
            import pickle

            sample[key] = pickle.loads(value)
    return sample


extension_to_format = {"jpg": "jpeg"}


def _default_encoder(sample: Dict[str, Any], format: Optional[Union[str, bool]] = True):
    """A default encoder for webdataset.

    This handles common file extensions: .txt, .cls, .cls2, .jpg,
        .png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl
    These are the most common extensions used in webdataset.
    For other extensions, users can provide their own encoder.

    Args:
        sample: sample to encode.
        format: optional image format hint forwarded to the underlying
            image encoder.

    Returns:
        The sample with values encoded according to their key extension.
    """
    sample = dict(sample)
    for key, value in sample.items():
        extension = key.split(".")[-1]
        if key.startswith("__"):
            continue
        elif extension in ["txt"]:
            sample[key] = value.encode("utf-8")
        elif extension in ["cls", "cls2"]:
            sample[key] = str(value).encode("utf-8")
        elif extension in ["jpg", "jpeg", "png", "ppm", "pgm", "pbm", "pnm"]:
            import numpy as np
            import PIL.Image

            if isinstance(value, np.ndarray):
                value = PIL.Image.fromarray(value)
            assert isinstance(value, PIL.Image.Image)
            stream = io.BytesIO()
            value.save(
                stream, format=extension_to_format.get(extension.lower(), extension)
            )
            sample[key] = stream.getvalue()
        elif extension == "json":
            sample[key] = json.dumps(value).encode("utf-8")
        elif extension == "npy":
            import numpy as np

            stream = io.BytesIO()
            np.save(stream, value)
            sample[key] = stream.getvalue()
        elif extension == "mp":
            import msgpack

            sample[key] = msgpack.dumps(value)
        elif extension in ["pt", "pth"]:
            import torch

            stream = io.BytesIO()
            torch.save(value, stream)
            sample[key] = stream.getvalue()
        elif extension in ["pickle", "pkl"]:
            import pickle

            stream = io.BytesIO()
            pickle.dump(value, stream)
            sample[key] = stream.getvalue()
    return sample


def _make_iterable(block: BlockAccessor):
    """Make a block iterable.

    This is a placeholder for dealing with more complex blocks.

    Args:
        block: Ray Dataset block

    Returns:
        Iterable[Dict[str,Any]]: Iterable of samples
    """
    return block.iter_rows(public_row_format=False)


class WebDatasetDatasource(FileBasedDatasource):
    """A Datasource for WebDataset datasets (tar format with naming conventions)."""

    _FILE_EXTENSIONS = ["tar"]

    def __init__(
        self,
        paths: Union[str, List[str]],
        decoder: Optional[Union[bool, str, callable, list]] = True,
        fileselect: Optional[Union[bool, callable, list]] = None,
        filerename: Optional[Union[bool, callable, list]] = None,
        suffixes: Optional[Union[bool, callable, list]] = None,
        verbose_open: bool = False,
        expand_json: bool = False,
        **file_based_datasource_kwargs,
    ):
        super().__init__(paths, **file_based_datasource_kwargs)

        self.decoder = decoder
        self.fileselect = fileselect
        self.filerename = filerename
        self.suffixes = suffixes
        self.verbose_open = verbose_open
        self.expand_json = expand_json

        self._allow_unsafe_deserialization = env_bool(
            ALLOW_UNSAFE_DESERIALIZATION_ENV_VAR, False
        )

    def _read_stream(self, stream: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
        """Read and decode samples from a stream.

        Note that fileselect selects files during reading, while suffixes
        selects files during the grouping step.

        Args:
            stream: File descriptor to read from.
            path: Path to the data.

        Yields:
            Block: Single-row blocks (one per WebDataset sample).
        """

        import pandas as pd

        def get_tar_file_iterator():
            return _tar_file_iterator(
                stream,
                fileselect=self.fileselect,
                filerename=self.filerename,
                verbose_open=self.verbose_open,
            )

        # S3 can raise transient errors during iteration
        files = iterate_with_retry(
            get_tar_file_iterator,
            "iterate tar file",
            match=self._data_context.retried_io_errors,
        )

        samples = _group_by_keys(files, meta=dict(__url__=path), suffixes=self.suffixes)
        default_decoder = partial(
            _default_decoder, allow_unsafe=self._allow_unsafe_deserialization
        )
        for sample in samples:
            if self.decoder is not None:
                sample = _apply_list(self.decoder, sample, default=default_decoder)
            if self.expand_json:
                if isinstance(sample["json"], bytes):
                    parsed_json = json.loads(sample["json"].decode("utf-8"))
                elif isinstance(sample["json"], str):
                    parsed_json = json.loads(sample["json"])
                elif isinstance(sample["json"], dict):
                    parsed_json = sample["json"]
                else:
                    raise TypeError(
                        f"Unsupported data type" f" {type(sample['json'])} for sample"
                    )
                for k, v in parsed_json.items():
                    if k not in sample:
                        sample[k] = []
                    sample[k].append(v)
            yield pd.DataFrame(
                {
                    k: v if isinstance(v, list) and len(v) == 1 else [v]
                    for k, v in sample.items()
                }
            )