from __future__ import annotations

import numpy
import pytest
import h5py
from nxtomomill.utils.h5pyutils import resolve_link
from silx.io.url import DataUrl

# start fixture definition


def _create_raw_dataset(
    file_path, data_shape=(10, 10), data_type=numpy.uint32, linspace_start=0
):
    """Utility function to create a raw dataset without any link"""
    with h5py.File(file_path, "w") as f:
        f.create_dataset(
            "data",
            data=numpy.linspace(
                linspace_start,
                linspace_start + 100,
                num=numpy.prod(data_shape),
                dtype=data_type,
            ).reshape(data_shape),
        )


@pytest.fixture
def raw_dataset(tmp_path):
    """Fixture that creates a dataset without any link"""
    target_file = tmp_path / "target.h5"
    _create_raw_dataset(target_file)
    return target_file


@pytest.fixture
def raw_datasets(tmp_path):
    """Fixture that creates 3 several datasets without any link"""
    files = []
    for i in range(3):
        file = tmp_path / f"target_{i}.h5"
        _create_raw_dataset(file, linspace_start=(i * 100))
        files.append(file)
    return files


@pytest.fixture
def dataset_with_external_link(tmp_path, raw_dataset):
    """Fixture that creates a dataset with an external link"""
    source_file = tmp_path / "source.h5"
    with h5py.File(source_file, "w") as f:
        f["external_link"] = h5py.ExternalLink(str(raw_dataset), "data")

    return source_file, raw_dataset


@pytest.fixture
def dataset_with_single_virtual_source(tmp_path, raw_dataset):
    """Fixture that creates a dataset with a single virtual source. The Virtual source is a is_regular_hyperslab"""
    vds_dataset_file = tmp_path / "vds_dataset.h5"
    with h5py.File(vds_dataset_file, mode="w") as o_h5s:
        layout = h5py.VirtualLayout(shape=(10, 10), dtype=numpy.uint32)
        with h5py.File(raw_dataset, mode="r") as i_h5s:
            layout[:] = h5py.VirtualSource(path_or_dataset=i_h5s["data"])
        o_h5s.create_virtual_dataset("vds_data", layout, fillvalue=-5)
    return vds_dataset_file, raw_dataset


@pytest.fixture
def dataset_with_single_virtual_source_2(tmp_path, dataset_with_single_virtual_source):
    """As 'dataset_with_single_virtual_source' but adds one more level of indirection between the VDS and the original dataset."""
    vds_dataset_file_l1, raw_dataset = dataset_with_single_virtual_source
    vds_dataset_file_l2 = tmp_path / "vds_dataset_2.h5"
    with h5py.File(vds_dataset_file_l2, mode="w") as o_h5s:
        layout = h5py.VirtualLayout(shape=(10, 10), dtype=numpy.uint32)
        with h5py.File(vds_dataset_file_l1, mode="r") as i_h5s:
            layout[:] = h5py.VirtualSource(path_or_dataset=i_h5s["vds_data"])
        o_h5s.create_virtual_dataset("vds_data", layout, fillvalue=-5)
    return vds_dataset_file_l2, raw_dataset


# Define a "dataset" fixture that uses request.param to return the correct dataset
@pytest.fixture
def dataset_with_vds(request):
    # request.param will be set to the values in the parametrize decorator
    return request.getfixturevalue(request.param)


@pytest.fixture
def dataset_virtual_source_with_irregular_hyperslab(
    tmp_path,
    raw_dataset,
):
    """Fixture that creates a dataset with a single virtual source. The Virtual source is a is_regular_hyperslab"""
    # create VDS
    vds_dataset_file = tmp_path / "vds_dataset.h5"
    with h5py.File(vds_dataset_file, mode="w") as o_h5s:
        layout = h5py.VirtualLayout(shape=(5, 10), dtype=numpy.uint32)
        with h5py.File(raw_dataset, mode="r") as i_h5s:
            layout[:] = h5py.VirtualSource(path_or_dataset=i_h5s["data"][::2])
        o_h5s.create_virtual_dataset("vds_data", layout, fillvalue=-5)
    return vds_dataset_file, raw_dataset


@pytest.fixture
def dataset_with_several_virtual_source(raw_datasets, tmp_path):
    """Fixture that creates a dataset with a single virtual source. All virtual sources are regular hyperslab containing the entire original dataset"""
    file_with_vds = tmp_path / "dataset_with_several_vds.h5"
    with h5py.File(file_with_vds, mode="w") as o_h5s:
        layout = h5py.VirtualLayout(shape=(30, 10), dtype=numpy.uint32)
        for i, raw_dataset in enumerate(raw_datasets):
            with h5py.File(raw_dataset, mode="r") as i_h5s:
                index_start, index_stop = i * 10, (i + 1) * 10
                layout[index_start:index_stop] = h5py.VirtualSource(
                    path_or_dataset=i_h5s["data"]
                )
        o_h5s.create_virtual_dataset("vds_data", layout, fillvalue=-5)
    return file_with_vds, raw_datasets


@pytest.fixture
def dataset_with_several_subset_virtual_source(raw_datasets, tmp_path):
    """Fixture that creates a dataset with a single virtual source. The first dataset only contains a subset of the VS target dataset"""
    file_with_vds = tmp_path / "dataset_with_several_vds.h5"
    with h5py.File(file_with_vds, mode="w") as o_h5s:
        layout = h5py.VirtualLayout(shape=(15, 10), dtype=numpy.uint32)
        for i, raw_dataset in enumerate(raw_datasets):
            with h5py.File(raw_dataset, mode="r") as i_h5s:
                index_start, index_stop = i * 5, (i * 5) + 5
                vs = h5py.VirtualSource(path_or_dataset=i_h5s["data"])
                layout[index_start:index_stop, :] = vs[:5, :]
        o_h5s.create_virtual_dataset("vds_data", layout, fillvalue=-5)
    return file_with_vds, raw_datasets


# end fixture definition


@pytest.mark.parametrize("open_mode_for_reading", ("safe", "robust"))
def test_external_link_resolution(dataset_with_external_link, open_mode_for_reading):
    """Test external link resolution"""
    source_file, target_file = dataset_with_external_link
    source_url = DataUrl(
        file_path=str(source_file), data_path="/external_link", scheme="silx"
    )
    solved_urls = resolve_link(source_url, open_mode_for_reading=open_mode_for_reading)
    assert len(solved_urls) == 1
    assert solved_urls[0].file_path() == str(target_file.absolute())
    assert solved_urls[0].data_path() == "data"


@pytest.mark.parametrize("open_mode_for_reading", ("safe", "robust"))
def test_raw_data_resolution(raw_dataset, open_mode_for_reading):
    """Test that a file without any link won't be affected by the 'resolve_link' function."""
    dataset_url = DataUrl(file_path=str(raw_dataset), data_path="data", scheme="silx")
    solved_urls = resolve_link(dataset_url, open_mode_for_reading=open_mode_for_reading)
    assert len(solved_urls) == 1
    assert solved_urls[0].path() == dataset_url.path()


@pytest.mark.parametrize("open_mode_for_reading", ("safe", "robust"))
@pytest.mark.parametrize(
    "dataset_with_vds",
    ("dataset_with_single_virtual_source", "dataset_with_single_virtual_source_2"),
    indirect=True,
)
def test_virtual_source_resolution(dataset_with_vds, open_mode_for_reading):
    """Test that a file with a single virtual source won't be affected by the 'resolve_link' function."""
    source_file, target_file = dataset_with_vds
    dataset_url = DataUrl(
        file_path=str(source_file), data_path="vds_data", scheme="silx"
    )
    solved_urls = resolve_link(dataset_url, open_mode_for_reading=open_mode_for_reading)
    assert len(solved_urls) == 1
    assert solved_urls[0].file_path() == str(target_file)
    assert solved_urls[0].data_path() == "/data"


@pytest.mark.parametrize("open_mode_for_reading", ("safe", "robust"))
def test_dataset_with_several_virtual_source(
    dataset_with_several_virtual_source, open_mode_for_reading
):
    """Test a virtual dataset containing several virtual sources. All virtual sources are regular hyperslab containing the entire original dataset"""
    file_with_vds, raw_datasets = dataset_with_several_virtual_source
    solved_urls = resolve_link(
        DataUrl(file_path=str(file_with_vds), data_path="/vds_data", scheme="silx"),
        open_mode_for_reading=open_mode_for_reading,
    )
    assert len(solved_urls) == 3
    for i_url, url in enumerate(solved_urls):
        assert url.file_path() == str(raw_datasets[i_url])
        assert url.data_path() == "/data"


@pytest.mark.parametrize("open_mode_for_reading", ("safe", "robust"))
def test_dataset_with_several_subset_virtual_source(
    dataset_with_several_subset_virtual_source,
    open_mode_for_reading,
):
    """Test a virtual dataset containing several virtual sources. But this time sources are not 'complete' but only take a subset of the original dataset."""
    file_with_vds, _ = dataset_with_several_subset_virtual_source
    url_with_vds = DataUrl(
        file_path=str(file_with_vds), data_path="/vds_data", scheme="silx"
    )
    solved_urls = resolve_link(
        url_with_vds, open_mode_for_reading=open_mode_for_reading
    )
    assert len(solved_urls) == 1
    assert solved_urls[0].path() == url_with_vds.path()
