"""
Helper functions for testing.
"""

import importlib
import inspect
import string
from pathlib import Path

import xarray as xr
from pygmt.exceptions import GMTImageComparisonFailure


def check_figures_equal(*, extensions=("png",), tol=0.0, result_dir="result_images"):
    """
    Decorator for test cases that generate and compare two figures.

    The decorated function must return two arguments, *fig_ref* and *fig_test*,
    these two figures will then be saved and compared against each other.

    This decorator is practically identical to matplotlib's check_figures_equal
    function, but adapted for PyGMT figures. See also the original code at
    https://matplotlib.org/3.3.1/api/testing_api.html#
    matplotlib.testing.decorators.check_figures_equal

    Parameters
    ----------
    extensions : list
        The extensions to test. Default is ["png"].
    tol : float
        The RMS threshold above which the test is considered failed.
    result_dir : str
        The directory where the figures will be stored.

    Examples
    --------

    >>> import pytest
    >>> import shutil
    >>> from pygmt import Figure
    >>> from pathlib import Path

    >>> @check_figures_equal(result_dir="tmp_result_images")
    ... def test_check_figures_equal():
    ...     fig_ref = Figure()
    ...     fig_ref.basemap(projection="X5c", region=[0, 5, 0, 5], frame=True)
    ...     fig_test = Figure()
    ...     fig_test.basemap(
    ...         projection="X5c", region=[0, 5, 0, 5], frame=["WrStZ", "af"]
    ...     )
    ...     return fig_ref, fig_test
    >>> test_check_figures_equal()
    >>> assert len(list(Path("tmp_result_images").iterdir())) == 0
    >>> shutil.rmtree(path="tmp_result_images")  # cleanup folder if tests pass

    >>> @check_figures_equal(result_dir="tmp_result_images")
    ... def test_check_figures_unequal():
    ...     fig_ref = Figure()
    ...     fig_ref.basemap(projection="X5c", region=[0, 6, 0, 6], frame=True)
    ...     fig_test = Figure()
    ...     fig_test.basemap(projection="X5c", region=[0, 3, 0, 3], frame=True)
    ...     return fig_ref, fig_test
    >>> with pytest.raises(GMTImageComparisonFailure):
    ...     test_check_figures_unequal()
    >>> for suffix in ["", "-expected", "-failed-diff"]:
    ...     assert (
    ...         Path("tmp_result_images") / f"test_check_figures_unequal{suffix}.png"
    ...     ).exists()
    >>> shutil.rmtree(path="tmp_result_images")  # cleanup folder if tests pass
    """
    allowed_chars = set(string.digits + string.ascii_letters + "_-[]()")
    keyword_only = inspect.Parameter.KEYWORD_ONLY

    def decorator(func):
        import pytest  #  noqa: PLC0415
        from matplotlib.testing.compare import compare_images  # noqa: PLC0415

        Path(result_dir).mkdir(parents=True, exist_ok=True)
        old_sig = inspect.signature(func)

        @pytest.mark.parametrize("ext", extensions)
        def wrapper(*args, ext="png", request=None, **kwargs):
            if "ext" in old_sig.parameters:
                kwargs["ext"] = ext
            if "request" in old_sig.parameters:
                kwargs["request"] = request
            try:
                file_name = "".join(c for c in request.node.name if c in allowed_chars)
            except AttributeError:  # 'NoneType' object has no attribute 'node'
                file_name = func.__name__

            fig_ref, fig_test = None, None
            try:
                fig_ref, fig_test = func(*args, **kwargs)
                ref_image_path = Path(result_dir) / f"{file_name}-expected.{ext}"
                test_image_path = Path(result_dir) / f"{file_name}.{ext}"
                fig_ref.savefig(ref_image_path)
                fig_test.savefig(test_image_path)

                # Code below is adapted for PyGMT, and is originally based on
                # matplotlib.testing.decorators._raise_on_image_difference
                err = compare_images(
                    expected=ref_image_path,
                    actual=test_image_path,
                    tol=tol,
                    in_decorator=True,
                )
                if err is None:  # Images are the same
                    ref_image_path.unlink()
                    test_image_path.unlink()
                else:  # Images are not the same
                    for key in ["actual", "expected", "diff"]:
                        err[key] = Path(err[key]).relative_to(".")
                    msg = (
                        f"images not close (RMS {err['rms']:.3f}):\n"
                        f"\t{err['actual']}\n"
                        f"\t{err['expected']}"
                    )
                    raise GMTImageComparisonFailure(msg)
            finally:
                del fig_ref
                del fig_test

        parameters = [
            param
            for param in old_sig.parameters.values()
            if param.name not in {"fig_test", "fig_ref"}
        ]
        if "ext" not in old_sig.parameters:
            parameters += [inspect.Parameter("ext", keyword_only)]
        if "request" not in old_sig.parameters:
            parameters += [inspect.Parameter("request", keyword_only)]
        new_sig = old_sig.replace(parameters=parameters)
        wrapper.__signature__ = new_sig

        # reach a bit into pytest internals to hoist the marks from
        # our wrapped function
        new_marks = getattr(func, "pytestmark", []) + wrapper.pytestmark
        wrapper.pytestmark = new_marks

        return wrapper

    return decorator


def load_static_earth_relief() -> xr.DataArray:
    """
    Load the static_earth_relief.nc file for internal testing.

    Returns
    -------
    data
        A grid of Earth relief for internal tests.
    """
    return xr.load_dataarray(
        "@static_earth_relief.nc", engine="gmt", raster_kind="grid"
    )


def skip_if_no(package):
    """
    Generic function to help skip tests when required packages are not present on the
    testing system.

    This function returns a pytest mark with a skip condition that will be
    evaluated during test collection. An attempt will be made to import the
    specified ``package``.

    The mark can be used as either a decorator for a test class or to be
    applied to parameters in pytest.mark.parametrize calls or parametrized
    fixtures. Use pytest.importorskip if an imported moduled is later needed
    or for test functions.

    If the import is unsuccessful, then the test function (or test case when
    used in conjunction with parametrization) will be skipped.

    Adapted from
    https://github.com/pandas-dev/pandas/blob/v2.1.4/pandas/util/_test_decorators.py#L121

    Parameters
    ----------
    package : str
        The name of the required package.

    Returns
    -------
    pytest.MarkDecorator
        A pytest.mark.skipif to use as either a test decorator or a
        parametrization mark.
    """
    import pytest  # noqa: PLC0415

    try:
        _ = importlib.import_module(name=package)
        has_package = True
    except ImportError:
        has_package = False
    return pytest.mark.skipif(not has_package, reason=f"Could not import '{package}'")
