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

import glob
import io
import os
import pickle
import tarfile

import pytest
import webdataset as wds

import ray
from ray.tests.conftest import *  # noqa


class TarWriter:
    def __init__(self, path):
        self.path = path
        self.tar = tarfile.open(path, "w")

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.tar.close()

    def write(self, name, data):
        f = self.tar.tarinfo()
        f.name = name
        f.size = len(data)
        self.tar.addfile(f, io.BytesIO(data))


def test_webdataset_read(ray_start_2_cpus, tmp_path):
    path = os.path.join(tmp_path, "bar_000000.tar")
    with TarWriter(path) as tf:
        for i in range(100):
            tf.write(f"{i}.a", str(i).encode("utf-8"))
            tf.write(f"{i}.b", str(i**2).encode("utf-8"))
    assert os.path.exists(path)
    assert len(glob.glob(f"{tmp_path}/*.tar")) == 1
    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
    samples = ds.take(100)
    assert len(samples) == 100
    for i, sample in enumerate(samples):
        assert isinstance(sample, dict), sample
        assert sample["__key__"] == str(i)
        assert sample["a"].decode("utf-8") == str(i)
        assert sample["b"].decode("utf-8") == str(i**2)


@pytest.fixture
def allow_unsafe_deserialization(monkeypatch):
    monkeypatch.setenv("RAY_DATA_WEBDATASET_ALLOW_UNSAFE_DESERIALIZATION", "1")


def test_webdataset_expand_json(
    ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
):
    import numpy as np
    import torch

    image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
    gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
    dstruct = dict(a=[1, 2], b=dict(c=2), d="hello")
    ttensor = torch.tensor([1, 2, 3]).numpy()

    sample = {
        "__key__": "foo",
        "jpg": image,
        "gray.png": gray,
        "mp": dstruct,
        "json": dstruct,
        "pt": ttensor,
        "und": b"undecoded",
        "custom": b"nothing",
    }

    # write the encoded data using the default encoder
    data = [sample]
    ds = ray.data.from_items(data).repartition(1)
    ds.write_webdataset(path=tmp_path, try_create_dir=True)
    ds = ray.data.read_webdataset(
        paths=[str(tmp_path)], override_num_blocks=1, expand_json=True
    )
    record = ds.take(1)
    assert [1, 2] == record[0]["a"]


def test_webdataset_suffixes(ray_start_2_cpus, tmp_path):
    path = os.path.join(tmp_path, "bar_000000.tar")
    with TarWriter(path) as tf:
        for i in range(100):
            tf.write(f"{i}.txt", str(i).encode("utf-8"))
            tf.write(f"{i}.test.txt", str(i**2).encode("utf-8"))
            tf.write(f"{i}.cls", str(i**2).encode("utf-8"))
            tf.write(f"{i}.test.cls2", str(i**2).encode("utf-8"))
    assert os.path.exists(path)
    assert len(glob.glob(f"{tmp_path}/*.tar")) == 1

    # test simple suffixes
    ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["txt", "cls"])
    samples = ds.take(100)
    assert len(samples) == 100
    for i, sample in enumerate(samples):
        assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls"}

    # test fnmatch patterns for suffixes
    ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["*.txt", "*.cls"])
    samples = ds.take(100)
    assert len(samples) == 100
    for i, sample in enumerate(samples):
        assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls", "test.txt"}

    # test selection function
    def select(name):
        return name.endswith("txt")

    ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=select)
    samples = ds.take(100)
    assert len(samples) == 100
    for i, sample in enumerate(samples):
        assert set(sample.keys()) == {"__url__", "__key__", "txt", "test.txt"}

    # test filerename
    def renamer(name):
        result = name.replace("txt", "text")
        print("***", name, result)
        return result

    ds = ray.data.read_webdataset(paths=[str(tmp_path)], filerename=renamer)
    samples = ds.take(100)
    assert len(samples) == 100
    for i, sample in enumerate(samples):
        assert set(sample.keys()) == {
            "__url__",
            "__key__",
            "text",
            "cls",
            "test.text",
            "test.cls2",
        }


def test_webdataset_write(ray_start_2_cpus, tmp_path):
    print(ray.available_resources())
    data = [dict(__key__=str(i), a=str(i), b=str(i**2)) for i in range(100)]
    ds = ray.data.from_items(data).repartition(1)
    ds.write_webdataset(path=tmp_path, try_create_dir=True)
    paths = glob.glob(f"{tmp_path}/*.tar")
    assert len(paths) == 1
    with open(paths[0], "rb") as stream:
        tf = tarfile.open(fileobj=stream)
        for i in range(100):
            assert tf.extractfile(f"{i}.a").read().decode("utf-8") == str(i)
            assert tf.extractfile(f"{i}.b").read().decode("utf-8") == str(i**2)


def custom_decoder(sample):
    for key, value in sample.items():
        if key == "png":
            # check that images have already been decoded
            assert not isinstance(value, bytes)
        elif key.endswith("custom"):
            sample[key] = "custom-value"
    return sample


def test_webdataset_coding(ray_start_2_cpus, tmp_path, allow_unsafe_deserialization):
    import numpy as np
    import PIL.Image
    import torch

    image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
    gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
    dstruct = dict(a=[1], b=dict(c=2), d="hello")
    ttensor = torch.tensor([1, 2, 3]).numpy()

    sample = {
        "__key__": "foo",
        "jpg": image,
        "gray.png": gray,
        "mp": dstruct,
        "json": dstruct,
        "pt": ttensor,
        "und": b"undecoded",
        "custom": b"nothing",
    }

    # write the encoded data using the default encoder
    data = [sample]
    ds = ray.data.from_items(data).repartition(1)
    ds.write_webdataset(path=tmp_path, try_create_dir=True)

    # read the encoded data using the default decoder
    paths = glob.glob(f"{tmp_path}/*.tar")
    assert len(paths) == 1
    path = paths[0]
    assert os.path.exists(path)
    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
    samples = ds.take(1)
    assert len(samples) == 1
    for sample in samples:
        assert isinstance(sample, dict), sample
        assert sample["__key__"] == "foo"
        assert isinstance(sample["jpg"], np.ndarray)
        assert sample["jpg"].shape == (100, 100, 3)
        assert isinstance(sample["gray.png"], np.ndarray)
        assert sample["gray.png"].shape == (100, 100)
        assert isinstance(sample["mp"], dict)
        assert sample["mp"]["a"] == [1]
        assert sample["mp"]["b"]["c"] == 2
        assert isinstance(sample["json"], dict)
        assert sample["json"]["a"] == [1]
        assert isinstance(sample["pt"], np.ndarray)
        assert sample["pt"].tolist() == [1, 2, 3]

    # test the format argument to the default decoder and multiple decoders
    ds = ray.data.read_webdataset(
        paths=[str(tmp_path)], decoder=["PIL", custom_decoder]
    )
    samples = ds.take(1)
    assert len(samples) == 1
    for sample in samples:
        assert isinstance(sample, dict), sample
        assert sample["__key__"] == "foo"
        assert isinstance(sample["jpg"], PIL.Image.Image)
        assert isinstance(sample["gray.png"], PIL.Image.Image)
        assert isinstance(sample["und"], bytes)
        assert sample["und"] == b"undecoded"
        assert sample["custom"] == "custom-value"


def test_webdataset_decoding(ray_start_2_cpus, tmp_path):
    import numpy as np
    import torch

    image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
    gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
    dstruct = dict(a=np.nan, b=dict(c=2), d="hello", e={"img_filename": "for_test.jpg"})
    ttensor = torch.tensor([1, 2, 3]).numpy()

    sample = {
        "__key__": "foo",
        "jpg": image,
        "gray.png": gray,
        "mp": dstruct,
        "json": dstruct,
        "pt": ttensor,
        "und": b"undecoded",
        "custom": b"nothing",
    }

    # write the encoded data using the default encoder
    data = [sample]
    ds = ray.data.from_items(data).repartition(1)
    ds.write_webdataset(path=tmp_path, try_create_dir=True)

    ds = ray.data.read_webdataset(
        paths=[str(tmp_path)],
        override_num_blocks=1,
        decoder=None,
    )
    samples = ds.take(1)
    import json

    meta_json = json.loads(samples[0]["json"].decode("utf-8"))
    assert meta_json["e"]["img_filename"] == "for_test.jpg"


@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
def test_write_min_rows_per_file(tmp_path, ray_start_2_cpus, min_rows_per_file):
    ray.data.from_items(
        [{"id": str(i)} for i in range(100)], override_num_blocks=20
    ).write_webdataset(tmp_path, min_rows_per_file=min_rows_per_file)

    for filename in os.listdir(tmp_path):
        dataset = wds.WebDataset(os.path.join(tmp_path, filename))
        assert len(list(dataset)) == min_rows_per_file


@pytest.mark.parametrize(
    "filename",
    ["000000.pkl", "000000.pickle", "000000.pt", "000000.pth"],
)
def test_default_decoder_rejects_unsafe_extensions(
    ray_start_2_cpus, tmp_path, filename
):
    path = os.path.join(tmp_path, "unsafe.tar")
    with TarWriter(path) as tf:
        tf.write(filename, b"fake-payload")

    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
    with pytest.raises(Exception, match="Refusing to"):
        ds.take_all()


def test_default_decoder_allows_unsafe_with_env_var(
    ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
):
    path = os.path.join(tmp_path, "trusted.tar")
    with TarWriter(path) as tf:
        tf.write("000000.pkl", pickle.dumps({"key": "value"}))

    ds = ray.data.read_webdataset(paths=[str(tmp_path)])
    rows = ds.take_all()

    assert len(rows) == 1
    assert rows[0]["pkl"] == {"key": "value"}


def test_custom_decoder_bypasses_unsafe_guard(ray_start_2_cpus, tmp_path):
    path = os.path.join(tmp_path, "custom.tar")
    with TarWriter(path) as tf:
        tf.write("000000.pkl", pickle.dumps({"key": "value"}))

    def safe_pkl_decoder(sample):
        sample = dict(sample)
        for key, value in sample.items():
            if key == "pkl":
                sample[key] = pickle.loads(value)
        return sample

    ds = ray.data.read_webdataset(paths=[str(tmp_path)], decoder=safe_pkl_decoder)
    rows = ds.take_all()

    assert len(rows) == 1
    assert rows[0]["pkl"] == {"key": "value"}


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))