from unittest.mock import MagicMock, patch

import datasets
import pyarrow
import pytest
import requests
from packaging.version import Version

import ray
from ray.data.dataset import Dataset, MaterializedDataset
from ray.tests.conftest import *  # noqa


@pytest.fixture
def mock_hf_dataset():
    """Create a mock HuggingFace dataset for testing."""
    texts = [
        "Climate change is a serious threat to our planet",
        "We need to take action on global warming",
        "Renewable energy is the future",
        "Fossil fuels are destroying the environment",
        "Solar power is becoming more affordable",
        "Wind energy is growing rapidly",
        "Electric vehicles are the way forward",
        "Carbon emissions must be reduced",
        "Green technology is advancing quickly",
        "Sustainability is important for future generations",
        "Climate science is well established",
        "Ocean levels are rising due to warming",
        "Extreme weather events are increasing",
        "Biodiversity loss is accelerating",
        "Deforestation contributes to climate change",
        "Clean energy jobs are growing",
        "Energy efficiency saves money",
        "Public transportation reduces emissions",
        "Plant-based diets help the environment",
        "Recycling is essential for sustainability",
    ]

    # Create labels array with exactly the same length as texts
    labels = [i % 2 for i in range(len(texts))]  # Alternating 0s and 1s

    return datasets.Dataset.from_dict(
        {
            "text": texts,
            "label": labels,
        }
    )


@pytest.fixture
def mock_hf_dataset_dict(mock_hf_dataset):
    """Create a mock HuggingFace DatasetDict for testing."""
    return datasets.DatasetDict({"train": mock_hf_dataset})


@pytest.fixture
def mock_hf_iterable_dataset():
    """Create a mock HuggingFace IterableDataset for testing."""
    texts = [
        "Streaming climate tweet 1: The planet is warming",
        "Streaming climate tweet 2: Renewable energy is key",
        "Streaming climate tweet 3: We must act now",
        "Streaming climate tweet 4: Solar panels everywhere",
        "Streaming climate tweet 5: Wind turbines are beautiful",
        "Streaming climate tweet 6: Electric cars are the future",
        "Streaming climate tweet 7: Carbon neutral by 2050",
        "Streaming climate tweet 8: Green energy revolution",
        "Streaming climate tweet 9: Climate action needed",
        "Streaming climate tweet 10: Sustainable development",
        "Streaming climate tweet 11: Ocean conservation",
        "Streaming climate tweet 12: Forest protection",
        "Streaming climate tweet 13: Clean air matters",
        "Streaming climate tweet 14: Water conservation",
        "Streaming climate tweet 15: Biodiversity protection",
    ]

    labels = [1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1]

    dataset = datasets.Dataset.from_dict(
        {
            "text": texts,
            "label": labels,
        }
    )
    iterable_dataset = dataset.to_iterable_dataset()
    iterable_dataset.expected_count = len(texts)
    return iterable_dataset


@pytest.fixture
def mock_parquet_urls():
    """Fixture providing mock parquet URLs for testing."""
    return [
        "https://huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
        "https://huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
    ]


@pytest.fixture
def mock_resolved_urls():
    """Fixture providing mock resolved URLs (after HTTP redirects) for testing."""
    return [
        "https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
        "https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
    ]


@pytest.fixture
def mock_ray_dataset(mock_hf_dataset):
    """Fixture providing a mock Ray dataset that matches the mock HuggingFace dataset."""
    return ray.data.from_items(
        [
            {"text": text, "label": label}
            for text, label in zip(mock_hf_dataset["text"], mock_hf_dataset["label"])
        ]
    )


@pytest.fixture
def mock_successful_http_responses(mock_parquet_urls):
    """Fixture providing mock successful HTTP responses for URL resolution."""
    mock_responses = []
    for url in mock_parquet_urls:
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.url = url
        mock_responses.append(mock_response)
    return mock_responses


@pytest.fixture
def mock_redirected_http_responses(mock_parquet_urls, mock_resolved_urls):
    """Fixture providing mock HTTP responses that simulate redirects."""
    mock_responses = []
    for original_url, resolved_url in zip(mock_parquet_urls, mock_resolved_urls):
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.url = resolved_url
        mock_responses.append(mock_response)
    return mock_responses


@pytest.fixture
def mock_huggingface_datasource():
    """Fixture providing the HuggingFaceDatasource class for mocking."""
    from ray.data._internal.datasource.huggingface_datasource import (
        HuggingFaceDatasource,
    )

    return HuggingFaceDatasource


def verify_http_requests(mock_requests_head, expected_urls):
    """Verify that HTTP requests were made correctly."""
    assert mock_requests_head.call_count == len(expected_urls)

    for i, url in enumerate(expected_urls):
        call_args = mock_requests_head.call_args_list[i]
        assert call_args[0][0] == url
        assert call_args[1]["allow_redirects"] is True
        assert call_args[1]["timeout"] == 5


def verify_read_parquet_call(mock_read_parquet, expected_urls):
    """Verify that read_parquet was called with correct parameters."""
    mock_read_parquet.assert_called_once()
    call_args = mock_read_parquet.call_args

    # Check that the parquet URLs were passed
    assert call_args[0][0] == expected_urls

    # Check that the filesystem is HTTPFileSystem
    assert "filesystem" in call_args[1]
    assert "HTTPFileSystem" in str(type(call_args[1]["filesystem"]))

    # Check that retry_exceptions includes FileNotFoundError and ClientResponseError
    assert "ray_remote_args" in call_args[1]
    assert FileNotFoundError in call_args[1]["ray_remote_args"]["retry_exceptions"]


def verify_dataset_creation(ds, mock_hf_dataset):
    """Verify that the dataset was created successfully."""
    assert isinstance(ds, MaterializedDataset)
    assert ds.count() == mock_hf_dataset.num_rows


def setup_parquet_mocks(
    mock_huggingface_datasource,
    mock_parquet_urls,
    mock_http_responses,
    mock_ray_dataset,
):
    """Setup common mocking pattern for parquet-based tests."""
    patches = []

    # Mock the list_parquet_urls_from_dataset method
    datasource_patch = patch.object(
        mock_huggingface_datasource,
        "list_parquet_urls_from_dataset",
        return_value=mock_parquet_urls,
    )
    patches.append(datasource_patch)

    # Mock the requests.head calls
    requests_patch = patch("requests.head")
    patches.append(requests_patch)

    # Mock the read_parquet function
    read_parquet_patch = patch("ray.data.read_api.read_parquet")
    patches.append(read_parquet_patch)

    # Start all patches
    datasource_mock = datasource_patch.start()
    requests_mock = requests_patch.start()
    read_parquet_mock = read_parquet_patch.start()

    # Configure mocks
    requests_mock.side_effect = mock_http_responses
    read_parquet_mock.return_value = mock_ray_dataset

    return datasource_mock, requests_mock, read_parquet_mock, patches


def hfds_assert_equals(hfds: datasets.Dataset, ds: Dataset):
    hfds_table = hfds.data.table
    ds_table = pyarrow.concat_tables([ray.get(tbl) for tbl in ds.to_arrow_refs()])

    sorting = [(name, "descending") for name in hfds_table.column_names]
    hfds_table = hfds_table.sort_by(sorting)
    ds_table = ds_table.sort_by(sorting)

    assert hfds_table.equals(ds_table)


@pytest.mark.parametrize("num_par", [1, 4])
def test_from_huggingface(mock_hf_dataset_dict, ray_start_regular_shared, num_par):
    # Check that DatasetDict is not directly supported.
    assert isinstance(mock_hf_dataset_dict, datasets.DatasetDict)
    with pytest.raises(
        DeprecationWarning,
        match="You provided a Hugging Face DatasetDict",
    ):
        ray.data.from_huggingface(mock_hf_dataset_dict)

    ray_datasets = {
        "train": ray.data.from_huggingface(
            mock_hf_dataset_dict["train"], override_num_blocks=num_par
        ),
    }

    assert isinstance(ray_datasets["train"], ray.data.Dataset)
    hfds_assert_equals(mock_hf_dataset_dict["train"], ray_datasets["train"])

    # Test reading in a split Hugging Face dataset yields correct individual datasets
    base_hf_dataset = mock_hf_dataset_dict["train"]
    hf_dataset_split = base_hf_dataset.train_test_split(test_size=0.2)
    ray_dataset_split_train = ray.data.from_huggingface(hf_dataset_split["train"])
    assert ray_dataset_split_train.count() == hf_dataset_split["train"].num_rows


@pytest.mark.skipif(
    datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
    reason="IterableDataset.iter() added in 2.8.0",
)
@pytest.mark.skipif(
    Version(pyarrow.__version__) < Version("8.0.0"),
    reason="pyarrow.Table.to_reader() added in 8.0.0",
)
# Note, pandas is excluded here because IterableDatasets do not support pandas format.
@pytest.mark.parametrize(
    "batch_format",
    [None, "numpy", "arrow", "torch", "tensorflow", "jax"],
)
def test_from_huggingface_streaming(
    mock_hf_iterable_dataset, batch_format, ray_start_regular_shared
):
    hfds = mock_hf_iterable_dataset.with_format(batch_format)
    assert isinstance(hfds, datasets.IterableDataset)

    ds = ray.data.from_huggingface(hfds)
    expected_count = mock_hf_iterable_dataset.expected_count
    assert ds.count() == expected_count


@pytest.mark.skipif(
    datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
    reason="IterableDataset.iter() added in 2.8.0",
)
def test_from_huggingface_dynamic_generated(ray_start_regular_shared):
    # https://github.com/ray-project/ray/issues/49529
    # Mock the dynamic dataset loading
    mock_dataset = datasets.Dataset.from_dict(
        {
            "text": [
                "dynamic tweet 1",
                "dynamic tweet 2",
                "dynamic tweet 3",
                "dynamic tweet 4",
                "dynamic tweet 5",
            ],
            "label": [0, 1, 0, 1, 0],
        }
    )
    mock_iterable = mock_dataset.to_iterable_dataset()

    with patch("datasets.load_dataset", return_value=mock_iterable):
        hfds = datasets.load_dataset(
            "dataset-org/dream",
            split="test",
            streaming=True,
            trust_remote_code=True,
        )
        ds = ray.data.from_huggingface(hfds)
        ds.take(1)


@pytest.mark.parametrize("override_num_blocks", [1, 2, 4, 8])
def test_from_huggingface_override_num_blocks(
    mock_hf_dataset, ray_start_regular_shared, override_num_blocks
):
    """Test that override_num_blocks works correctly with HuggingFace datasets."""
    hf_train = mock_hf_dataset

    ds_subset = ray.data.from_huggingface(
        hf_train, override_num_blocks=override_num_blocks
    )

    assert isinstance(ds_subset, MaterializedDataset)

    # Verify number of blocks
    assert ds_subset.num_blocks() == override_num_blocks

    # Verify data integrity
    assert ds_subset.count() == hf_train.num_rows
    hfds_assert_equals(hf_train, ds_subset)

    # Test with a smaller subset to test edge cases
    small_size = max(override_num_blocks * 3, 10)
    hf_small = hf_train.select(range(min(small_size, hf_train.num_rows)))
    ds_small = ray.data.from_huggingface(
        hf_small, override_num_blocks=override_num_blocks
    )

    # Verify number of blocks
    assert ds_small.num_blocks() == override_num_blocks

    # Verify data integrity
    assert ds_small.count() == hf_small.num_rows
    hfds_assert_equals(hf_small, ds_small)


def test_from_huggingface_with_parquet_files(
    mock_hf_dataset,
    ray_start_regular_shared,
    mock_parquet_urls,
    mock_ray_dataset,
    mock_successful_http_responses,
    mock_huggingface_datasource,
):
    """Test the distributed read path when parquet file URLs are available."""
    datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
        mock_huggingface_datasource,
        mock_parquet_urls,
        mock_successful_http_responses,
        mock_ray_dataset,
    )

    try:
        ds = ray.data.from_huggingface(mock_hf_dataset)

        # Verify HTTP requests
        verify_http_requests(requests_mock, mock_parquet_urls)

        # Verify read_parquet call
        verify_read_parquet_call(read_parquet_mock, mock_parquet_urls)

        # Verify dataset creation
        verify_dataset_creation(ds, mock_hf_dataset)

    finally:
        # Stop all patches
        for patch_obj in patches:
            patch_obj.stop()


def test_from_huggingface_with_resolved_urls(
    mock_hf_dataset,
    ray_start_regular_shared,
    mock_parquet_urls,
    mock_resolved_urls,
    mock_ray_dataset,
    mock_redirected_http_responses,
    mock_huggingface_datasource,
):
    """Test the URL resolution logic when HTTP redirects are encountered."""
    datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
        mock_huggingface_datasource,
        mock_parquet_urls,
        mock_redirected_http_responses,
        mock_ray_dataset,
    )

    try:
        ds = ray.data.from_huggingface(mock_hf_dataset)

        # Verify HTTP requests
        verify_http_requests(requests_mock, mock_parquet_urls)

        # Verify read_parquet call with resolved URLs
        verify_read_parquet_call(read_parquet_mock, mock_resolved_urls)

        # Verify dataset creation
        verify_dataset_creation(ds, mock_hf_dataset)

    finally:
        # Stop all patches
        for patch_obj in patches:
            patch_obj.stop()


def test_from_huggingface_url_resolution_failures(
    mock_hf_dataset,
    ray_start_regular_shared,
    mock_parquet_urls,
    mock_ray_dataset,
    mock_huggingface_datasource,
):
    """Test URL resolution failures fall back to single node read."""
    # Convert the mock dataset to an IterableDataset so it uses the read_datasource fallback
    mock_iterable_dataset = mock_hf_dataset.to_iterable_dataset()

    with patch.object(
        mock_huggingface_datasource,
        "list_parquet_urls_from_dataset",
        return_value=mock_parquet_urls,
    ):
        # Mock the requests.head calls to simulate failures
        with patch("requests.head") as mock_requests_head:
            # Configure mock to raise an exception for all URLs
            mock_requests_head.side_effect = requests.RequestException(
                "Connection failed"
            )

            # Mock the fallback path
            with patch("ray.data.read_api.read_datasource") as mock_read_datasource:
                mock_read_datasource.return_value = mock_ray_dataset

                ds = ray.data.from_huggingface(mock_iterable_dataset)

                # Verify that requests.head was called for each URL
                assert mock_requests_head.call_count == len(mock_parquet_urls)

                # Verify that the fallback read_datasource was called
                mock_read_datasource.assert_called_once()

                # Verify the dataset was created successfully via fallback
                verify_dataset_creation(ds, mock_hf_dataset)


if __name__ == "__main__":
    import sys

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