from __future__ import annotations

import sys
from importlib.metadata import EntryPoint, EntryPoints
from itertools import starmap
from unittest import mock

import pytest

from xarray.backends import common, plugins
from xarray.core.options import OPTIONS
from xarray.tests import (
    has_h5netcdf,
    has_netCDF4,
    has_pydap,
    has_scipy,
    has_zarr,
)

# Do not import list_engines here, this will break the lazy tests

importlib_metadata_mock = "importlib.metadata"


class DummyBackendEntrypointArgs(common.BackendEntrypoint):
    def open_dataset(filename_or_obj, *args):  # type: ignore[override]
        pass


class DummyBackendEntrypointKwargs(common.BackendEntrypoint):
    def open_dataset(filename_or_obj, **kwargs):  # type: ignore[override]
        pass


class DummyBackendEntrypoint1(common.BackendEntrypoint):
    def open_dataset(self, filename_or_obj, *, decoder):  # type: ignore[override]
        pass


class DummyBackendEntrypoint2(common.BackendEntrypoint):
    def open_dataset(self, filename_or_obj, *, decoder):  # type: ignore[override]
        pass


@pytest.fixture
def dummy_duplicated_entrypoints():
    specs = [
        ["engine1", "xarray.tests.test_plugins:backend_1", "xarray.backends"],
        ["engine1", "xarray.tests.test_plugins:backend_2", "xarray.backends"],
        ["engine2", "xarray.tests.test_plugins:backend_1", "xarray.backends"],
        ["engine2", "xarray.tests.test_plugins:backend_2", "xarray.backends"],
    ]
    eps = list(starmap(EntryPoint, specs))
    return eps


@pytest.mark.filterwarnings("ignore:Found")
def test_remove_duplicates(dummy_duplicated_entrypoints) -> None:
    with pytest.warns(RuntimeWarning):
        entrypoints = plugins.remove_duplicates(dummy_duplicated_entrypoints)
    assert len(entrypoints) == 2


def test_broken_plugin() -> None:
    broken_backend = EntryPoint(
        "broken_backend",
        "xarray.tests.test_plugins:backend_1",
        "xarray.backends",
    )
    with pytest.warns(RuntimeWarning) as record:
        _ = plugins.build_engines(EntryPoints([broken_backend]))
    assert len(record) == 1
    message = str(record[0].message)
    assert "Engine 'broken_backend'" in message


def test_remove_duplicates_warnings(dummy_duplicated_entrypoints) -> None:
    with pytest.warns(RuntimeWarning) as record:
        _ = plugins.remove_duplicates(dummy_duplicated_entrypoints)

    assert len(record) == 2
    message0 = str(record[0].message)
    message1 = str(record[1].message)
    assert "entrypoints" in message0
    assert "entrypoints" in message1


@mock.patch(
    f"{importlib_metadata_mock}.EntryPoint.load", mock.MagicMock(return_value=None)
)
def test_backends_dict_from_pkg() -> None:
    specs = [
        ["engine1", "xarray.tests.test_plugins:backend_1", "xarray.backends"],
        ["engine2", "xarray.tests.test_plugins:backend_2", "xarray.backends"],
    ]
    entrypoints = list(starmap(EntryPoint, specs))
    engines = plugins.backends_dict_from_pkg(entrypoints)
    assert len(engines) == 2
    assert engines.keys() == {"engine1", "engine2"}


def test_set_missing_parameters() -> None:
    backend_1 = DummyBackendEntrypoint1
    backend_2 = DummyBackendEntrypoint2
    backend_2.open_dataset_parameters = ("filename_or_obj",)
    engines = {"engine_1": backend_1, "engine_2": backend_2}
    plugins.set_missing_parameters(engines)

    assert len(engines) == 2
    assert backend_1.open_dataset_parameters == ("filename_or_obj", "decoder")
    assert backend_2.open_dataset_parameters == ("filename_or_obj",)

    backend_kwargs = DummyBackendEntrypointKwargs
    backend_kwargs.open_dataset_parameters = ("filename_or_obj", "decoder")
    plugins.set_missing_parameters({"engine": backend_kwargs})
    assert backend_kwargs.open_dataset_parameters == ("filename_or_obj", "decoder")

    backend_args = DummyBackendEntrypointArgs
    backend_args.open_dataset_parameters = ("filename_or_obj", "decoder")
    plugins.set_missing_parameters({"engine": backend_args})
    assert backend_args.open_dataset_parameters == ("filename_or_obj", "decoder")

    # reset
    backend_1.open_dataset_parameters = None
    backend_1.open_dataset_parameters = None
    backend_kwargs.open_dataset_parameters = None
    backend_args.open_dataset_parameters = None


def test_set_missing_parameters_raise_error() -> None:
    backend = DummyBackendEntrypointKwargs
    with pytest.raises(TypeError):
        plugins.set_missing_parameters({"engine": backend})

    backend_args = DummyBackendEntrypointArgs
    with pytest.raises(TypeError):
        plugins.set_missing_parameters({"engine": backend_args})


@mock.patch(
    f"{importlib_metadata_mock}.EntryPoint.load",
    mock.MagicMock(return_value=DummyBackendEntrypoint1),
)
def test_build_engines() -> None:
    dummy_pkg_entrypoint = EntryPoint(
        "dummy", "xarray.tests.test_plugins:backend_1", "xarray_backends"
    )
    backend_entrypoints = plugins.build_engines(EntryPoints([dummy_pkg_entrypoint]))

    assert isinstance(backend_entrypoints["dummy"], DummyBackendEntrypoint1)
    assert backend_entrypoints["dummy"].open_dataset_parameters == (
        "filename_or_obj",
        "decoder",
    )


@mock.patch(
    f"{importlib_metadata_mock}.EntryPoint.load",
    mock.MagicMock(return_value=DummyBackendEntrypoint1),
)
def test_build_engines_sorted() -> None:
    dummy_pkg_entrypoints = EntryPoints(
        [
            EntryPoint(
                "dummy2", "xarray.tests.test_plugins:backend_1", "xarray.backends"
            ),
            EntryPoint(
                "dummy1", "xarray.tests.test_plugins:backend_1", "xarray.backends"
            ),
        ]
    )
    backend_entrypoints = list(plugins.build_engines(dummy_pkg_entrypoints))

    indices = []
    for be in OPTIONS["netcdf_engine_order"]:
        try:
            index = backend_entrypoints.index(be)
            backend_entrypoints.pop(index)
            indices.append(index)
        except ValueError:
            pass

    assert set(indices) < {0, -1}
    assert list(backend_entrypoints) == sorted(backend_entrypoints)


@mock.patch(
    "xarray.backends.plugins.list_engines",
    mock.MagicMock(return_value={"dummy": DummyBackendEntrypointArgs()}),
)
def test_no_matching_engine_found() -> None:
    with pytest.raises(ValueError, match=r"did not find a match in any"):
        plugins.guess_engine("not-valid")

    with pytest.raises(ValueError, match=r"found the following matches with the input"):
        plugins.guess_engine("foo.nc")


@mock.patch(
    "xarray.backends.plugins.list_engines",
    mock.MagicMock(return_value={}),
)
def test_engines_not_installed() -> None:
    with pytest.raises(ValueError, match=r"xarray is unable to open"):
        plugins.guess_engine("not-valid")

    with pytest.raises(ValueError, match=r"found the following matches with the input"):
        plugins.guess_engine("foo.nc")


def test_lazy_import() -> None:
    """Test that some modules are imported in a lazy manner.

    When importing xarray these should not be imported as well.
    Only when running code for the first time that requires them.
    """
    deny_list = [
        "cubed",
        "cupy",
        # "dask",  # TODO: backends.locks is not lazy yet :(
        "dask.array",
        "dask.distributed",
        "flox",
        "h5netcdf",
        "matplotlib",
        "nc_time_axis",
        "netCDF4",
        "numbagg",
        "pint",
        "pydap",
        # "scipy",  # TODO: xarray.backends.scipy_ is currently not lazy
        "sparse",
        "zarr",
    ]
    # ensure that none of the above modules has been imported before
    modules_backup = {}
    for pkg in list(sys.modules.keys()):
        for mod in deny_list + ["xarray"]:
            if pkg.startswith(mod):
                modules_backup[pkg] = sys.modules[pkg]
                del sys.modules[pkg]
                break

    try:
        import xarray  # noqa: F401
        from xarray.backends import list_engines

        list_engines()

        # ensure that none of the modules that are supposed to be
        # lazy loaded are loaded when importing xarray
        is_imported = set()
        for pkg in sys.modules:
            for mod in deny_list:
                if pkg.startswith(mod):
                    is_imported.add(mod)
                    break
        assert len(is_imported) == 0, (
            f"{is_imported} have been imported but should be lazy"
        )

    finally:
        # restore original
        sys.modules.update(modules_backup)


def test_list_engines() -> None:
    from xarray.backends import list_engines

    engines = list_engines()
    assert list_engines.cache_info().currsize == 1

    assert ("scipy" in engines) == has_scipy
    assert ("h5netcdf" in engines) == has_h5netcdf
    assert ("netcdf4" in engines) == has_netCDF4
    assert ("pydap" in engines) == has_pydap
    assert ("zarr" in engines) == has_zarr
    assert "store" in engines


def test_refresh_engines() -> None:
    from xarray.backends import list_engines, refresh_engines

    EntryPointMock1 = mock.MagicMock()
    EntryPointMock1.name = "test1"
    EntryPointMock1.load.return_value = DummyBackendEntrypoint1

    return_value = EntryPoints([EntryPointMock1])

    with mock.patch("xarray.backends.plugins.entry_points", return_value=return_value):
        list_engines.cache_clear()
        engines = list_engines()
    assert "test1" in engines
    assert isinstance(engines["test1"], DummyBackendEntrypoint1)

    EntryPointMock2 = mock.MagicMock()
    EntryPointMock2.name = "test2"
    EntryPointMock2.load.return_value = DummyBackendEntrypoint2

    return_value2 = EntryPoints([EntryPointMock2])

    with mock.patch("xarray.backends.plugins.entry_points", return_value=return_value2):
        refresh_engines()
        engines = list_engines()
    assert "test1" not in engines
    assert "test2" in engines
    assert isinstance(engines["test2"], DummyBackendEntrypoint2)

    # reset to original
    refresh_engines()