"""Tests for imageio's freeimage plugin"""

import os
import shutil
import sys
import platform
import time
import warnings

import imageio.plugins
import imageio.v3 as iio3
import imageio.v2 as iio
import numpy as np
import pytest
from imageio import core
from imageio.core import IS_PYPY, get_platform
from imageio.plugins._freeimage import FNAME_PER_PLATFORM, download
from pytest import raises
from conftest import deprecated_test


@pytest.fixture(scope="module")
def vendored_lib(tmp_path_factory):
    """Download FreeImage binaries via raw.githubusercontent.com (with retries)."""
    plat = get_platform()
    if not plat or plat not in FNAME_PER_PLATFORM:
        pytest.skip("No FreeImage binary for this platform")

    lib_dir = tmp_path_factory.mktemp("freeimage_dir")
    last_exc = None
    for attempt in range(4):
        try:
            download(directory=str(lib_dir))
            break
        except Exception as exc:
            last_exc = exc
            time.sleep(0.5 * (2**attempt))
    else:
        pytest.skip(f"Could not download FreeImage binaries: {last_exc}")

    freeimage_dir = lib_dir / "freeimage"
    if not freeimage_dir.is_dir() or not any(freeimage_dir.iterdir()):
        pytest.skip("FreeImage download did not produce expected files")

    yield freeimage_dir


# TODD: update fixture once we transition into v3
@pytest.fixture(scope="module")
def setup_library(tmp_path_factory, vendored_lib):
    # Checks if freeimage is installed by the system
    from imageio.plugins.freeimage import fi

    use_imageio_binary = not fi.has_lib()

    # During this test, pretend that FI is the default format
    with warnings.catch_warnings(record=True):
        iio.formats.sort("-FI")

    if use_imageio_binary:
        if sys.platform.startswith("win"):
            user_dir_env = "LOCALAPPDATA"
        else:
            user_dir_env = "IMAGEIO_USERDIR"

        # Setup from test_images/freeimage
        ud = tmp_path_factory.mktemp("userdir")
        old_user_dir = os.getenv(user_dir_env, None)
        os.environ[user_dir_env] = str(ud)
        os.makedirs(ud, exist_ok=True)

        add = core.appdata_dir("imageio")
        os.makedirs(add, exist_ok=True)
        shutil.copytree(vendored_lib, os.path.join(add, "freeimage"))

    try:
        fi.load_freeimage()
    except OSError:
        pytest.skip("Could not find a compatible FreeImage version for this OS.")

    yield

    if use_imageio_binary:
        if old_user_dir is not None:
            os.environ[user_dir_env] = old_user_dir
        else:
            del os.environ[user_dir_env]

    # Sort formats back to normal
    with warnings.catch_warnings(record=True):
        iio.formats.sort()


# Create test images LUMINANCE
im0 = np.zeros((42, 32), np.uint8)
im0[:16, :] = 200
im1 = np.zeros((42, 32, 1), np.uint8)
im1[:16, :] = 200
# Create test image RGB
im3 = np.zeros((42, 32, 3), np.uint8)
im3[:16, :, 0] = 250
im3[:, :16, 1] = 200
im3[50:, :16, 2] = 100
# Create test image RGBA
im4 = np.zeros((42, 32, 4), np.uint8)
im4[:16, :, 0] = 250
im4[:, :16, 1] = 200
im4[50:, :16, 2] = 100
im4[:, :, 3] = 255
im4[20:, :, 3] = 120


def get_ref_im(colors, crop, isfloat):
    """Get reference image with
    * colors: 0, 1, 3, 4
    * cropping: 0-> none, 1-> crop, 2-> crop with non-contiguous data
    * float: False, True
    """
    assert colors in (0, 1, 3, 4)
    assert crop in (0, 1, 2)
    assert isfloat in (False, True)
    rim = [im0, im1, None, im3, im4][colors]
    if isfloat:
        rim = rim.astype(np.float32) / 255.0
    if crop == 1:
        rim = rim[:-1, :-1].copy()
    elif crop == 2:
        rim = rim[:-1, :-1]
    return rim


def assert_close(im1, im2, tol=0.0):
    if im1.ndim == 3 and im1.shape[-1] == 1:
        im1 = im1.reshape(im1.shape[:-1])
    if im2.ndim == 3 and im2.shape[-1] == 1:
        im2 = im2.reshape(im2.shape[:-1])
    assert im1.shape == im2.shape
    diff = im1.astype("float32") - im2.astype("float32")
    diff[15:17, :] = 0  # Mask edge artifacts
    diff[:, 15:17] = 0
    assert np.abs(diff).max() <= tol
    # import visvis as vv
    # vv.subplot(121); vv.imshow(im1); vv.subplot(122); vv.imshow(im2)


def test_download():
    # this is a regression test
    # see: https://github.com/imageio/imageio/issues/690

    assert hasattr(imageio.plugins.freeimage, "download")


def test_get_ref_im():
    """A test for our function to get test images"""

    crop = 0
    for f in (False, True):
        for colors in (0, 1, 3, 4):
            rim = get_ref_im(0, crop, f)
            assert rim.flags.c_contiguous is True
            assert rim.shape[:2] == (42, 32)

    crop = 1
    for f in (False, True):
        for colors in (0, 1, 3, 4):
            rim = get_ref_im(0, crop, f)
            assert rim.flags.c_contiguous is True
            assert rim.shape[:2] == (41, 31)

    if IS_PYPY:
        return "PYPY cannot have non-contiguous data"

    crop = 2
    for f in (False, True):
        for colors in (0, 1, 3, 4):
            rim = get_ref_im(0, crop, f)
            assert rim.flags.c_contiguous is False
            assert rim.shape[:2] == (41, 31)


@pytest.mark.needs_internet
def test_get_fi_lib(vendored_lib, tmp_userdir):
    from imageio.plugins._freeimage import get_freeimage_lib

    add = core.appdata_dir("imageio")
    os.makedirs(add, exist_ok=True)
    shutil.copytree(vendored_lib, os.path.join(add, "freeimage"))

    lib = get_freeimage_lib()
    assert os.path.isfile(lib)


@deprecated_test
def test_freeimage_format(setup_library, test_images, tmp_path):
    fnamebase = str(tmp_path / "test")

    # Format
    F = iio.formats["PNG-FI"]
    assert F.name == "PNG-FI"

    # Reader
    R = F.get_reader(core.Request(test_images / "chelsea.png", "ri"))
    assert len(R) == 1
    assert isinstance(R.get_meta_data(), dict)
    assert isinstance(R.get_meta_data(0), dict)
    raises(IndexError, R.get_data, 2)
    raises(IndexError, R.get_meta_data, 2)

    # Writer
    W = F.get_writer(core.Request(fnamebase + ".png", "wi"))
    W.append_data(im0)
    W.set_meta_data({"foo": 3})
    raises(RuntimeError, W.append_data, im0)


@pytest.mark.needs_internet
def test_freeimage_lib(setup_library):
    fi = imageio.plugins.freeimage.fi

    # Error messages
    imageio.plugins._freeimage.fi._messages.append("this is a test")
    assert imageio.plugins._freeimage.fi.get_output_log()
    imageio.plugins._freeimage.fi._show_any_warnings()
    imageio.plugins._freeimage.fi._get_error_message()

    # Test getfif
    raises(ValueError, fi.getFIF, "foo.png", "x")  # mode must be r or w
    raises(ValueError, fi.getFIF, "foo.notvalid", "w")  # invalid ext
    raises(ValueError, fi.getFIF, "foo.iff", "w")  # We cannot write iff


def test_png(setup_library, test_images, tmp_path):
    fnamebase = str(tmp_path / "test")

    for isfloat in (False, True):
        for crop in (0, 1, 2):
            for colors in (0, 1, 3, 4):
                fname = fnamebase + "%i.%i.%i.png" % (isfloat, crop, colors)
                rim = get_ref_im(colors, crop, isfloat)
                iio.imsave(fname, rim)
                im = iio.imread(fname)
                mul = 255 if isfloat else 1
                assert_close(rim * mul, im, 0.1)  # lossless

    # Run exact same test, but now in pypy backup mode
    try:
        imageio.plugins._freeimage.TEST_NUMPY_NO_STRIDES = True
        for isfloat in (False, True):
            for crop in (0, 1, 2):
                for colors in (0, 1, 3, 4):
                    fname = fnamebase + "%i.%i.%i.png" % (isfloat, crop, colors)
                    rim = get_ref_im(colors, crop, isfloat)
                    iio.imsave(fname, rim)
                    im = iio.imread(fname)
                    mul = 255 if isfloat else 1
                    assert_close(rim * mul, im, 0.1)  # lossless
    finally:
        imageio.plugins._freeimage.TEST_NUMPY_NO_STRIDES = False

    # Parameters
    im = iio.imread(test_images / "chelsea.png", ignoregamma=True)
    iio.imsave(fnamebase + ".png", im, interlaced=True)

    # Parameter fail
    raises(
        TypeError,
        iio.imread,
        test_images / "chelsea.png",
        notavalidk=True,
    )
    raises(TypeError, iio.imsave, fnamebase + ".png", im, notavalidk=True)

    # Compression
    iio.imsave(fnamebase + "1.png", im, compression=0)
    iio.imsave(fnamebase + "2.png", im, compression=9)
    s1 = os.stat(fnamebase + "1.png").st_size
    s2 = os.stat(fnamebase + "2.png").st_size
    assert s2 < s1
    # Fail
    raises(ValueError, iio.imsave, fnamebase + ".png", im, compression=12)

    # Quantize
    if sys.platform.startswith("darwin"):
        return  # quantization segfaults on my osx VM
    iio.imsave(fnamebase + "1.png", im, quantize=256)
    iio.imsave(fnamebase + "2.png", im, quantize=4)

    im = iio.imread(fnamebase + "2.png")  # touch palette read code
    s1 = os.stat(fnamebase + "1.png").st_size
    s2 = os.stat(fnamebase + "2.png").st_size
    assert s1 > s2
    # Fail
    fname = fnamebase + "1.png"
    raises(ValueError, iio.imsave, fname, im[:, :, :3], quantize=300)
    raises(ValueError, iio.imsave, fname, im[:, :, 0], quantize=100)


@pytest.mark.needs_internet
def test_png_dtypes(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    # See issue #44

    # Two images, one 0-255, one 0-200
    im1 = np.zeros((100, 100, 3), dtype="uint8")
    im2 = np.zeros((100, 100, 3), dtype="uint8")
    im1[20:80, 20:80, :] = 255
    im2[20:80, 20:80, :] = 200

    fname = fnamebase + ".dtype.png"

    # uint8
    iio.imsave(fname, im1)
    assert_close(im1, iio.imread(fname))
    iio.imsave(fname, im2)
    assert_close(im2, iio.imread(fname))

    # float scaled
    iio.imsave(fname, im1 / 255.0)
    assert_close(im1, iio.imread(fname))
    iio.imsave(fname, im2 / 255.0)
    assert_close(im2, iio.imread(fname))

    # float not scaled
    iio.imsave(fname, im1 * 1.0)
    assert_close(im1, iio.imread(fname))
    iio.imsave(fname, im2 * 1.0)
    assert_close(im1, iio.imread(fname))  # scaled

    # int16
    iio.imsave(fname, im1.astype("int16"))
    assert_close(im1, iio.imread(fname))
    iio.imsave(fname, im2.astype("int16"))
    assert_close(im1, iio.imread(fname))  # scaled


@pytest.mark.needs_internet
def test_jpg(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    for isfloat in (False, True):
        for crop in (0, 1, 2):
            for colors in (0, 1, 3):
                fname = fnamebase + "%i.%i.%i.jpg" % (isfloat, crop, colors)
                rim = get_ref_im(colors, crop, isfloat)
                iio.imsave(fname, rim)
                im = iio.imread(fname)
                mul = 255 if isfloat else 1
                assert_close(rim * mul, im, 1.1)  # lossy

    # No alpha in JPEG
    raises(Exception, iio.imsave, fname, im4)

    # Parameters
    iio.imsave(fnamebase + ".jpg", im3, progressive=True, optimize=True, baseline=True)

    # Parameter fail
    raises(
        TypeError,
        iio.imread,
        fnamebase + ".jpg",
        notavalidkwarg=True,
        format="JPEG-FI",
    )
    raises(
        TypeError,
        iio.imsave,
        fnamebase + ".jpg",
        im,
        notavalidk=True,
        format="JPEG-FI",
    )

    # Compression
    iio.imsave(fnamebase + "1.jpg", im3, quality=10)
    iio.imsave(fnamebase + "2.jpg", im3, quality=90)
    s1 = os.stat(fnamebase + "1.jpg").st_size
    s2 = os.stat(fnamebase + "2.jpg").st_size
    assert s2 > s1
    raises(ValueError, iio.imsave, fnamebase + ".jpg", im, quality=120)


def test_jpg_more(setup_library, test_images, tmp_path):
    fnamebase = str(tmp_path / "test")

    # Test broken JPEG
    fname = fnamebase + "_broken.jpg"
    open(fname, "wb").write(b"this is not an image")
    raises(Exception, iio.imread, fname)
    #
    img = get_ref_im(3, 0, 0)
    bb = iio.imsave(iio.RETURN_BYTES, img, "JPEG-FI")
    with open(fname, "wb") as f:
        f.write(bb[:400])
        f.write(b" ")
        f.write(bb[400:])
    raises(Exception, iio.imread, fname)

    # Test EXIF stuff
    fname = test_images / "rommel.jpg"
    im = iio.imread(fname)
    assert im.shape[0] > im.shape[1]
    im = iio.imread(fname, exifrotate=False)
    assert im.shape[0] < im.shape[1]
    im = iio.imread(fname, exifrotate=2)  # Rotation in Python
    assert im.shape[0] > im.shape[1]
    # Write the jpg and check that exif data is maintained
    if sys.platform.startswith("darwin"):
        return  # segfaults on my osx VM, why?
    iio.imsave(fnamebase + "rommel.jpg", im)
    im = iio.imread(fname)
    assert im.meta.EXIF_MAIN


@pytest.mark.needs_internet
def test_bmp(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    for isfloat in (False, True):
        for crop in (0, 1, 2):
            for colors in (0, 1, 3, 4):
                fname = fnamebase + "%i.%i.%i.bmp" % (isfloat, crop, colors)
                rim = get_ref_im(colors, crop, isfloat)
                iio.imsave(fname, rim, format="BMP-FI")
                im = iio.imread(fname, format="BMP-FI")
                mul = 255 if isfloat else 1
                assert_close(rim * mul, im, 0.1)  # lossless

    # Compression
    iio.imsave(fnamebase + "1.bmp", im3, compression=False, format="BMP-FI")
    iio.imsave(fnamebase + "2.bmp", im3, compression=True, format="BMP-FI")
    s1 = os.stat(fnamebase + "1.bmp").st_size
    s2 = os.stat(fnamebase + "2.bmp").st_size
    assert s1 + s2  # todo: bug in FreeImage? assert s1 < s2

    # Parameter fail
    raises(
        TypeError,
        iio.imread,
        fnamebase + "1.bmp",
        notavalidkwarg=True,
        format="BMP-FI",
    )
    raises(
        TypeError,
        iio.imsave,
        fnamebase + "1.bmp",
        im,
        notavalidk=True,
        format="BMP-FI",
    )


@pytest.mark.needs_internet
def test_gif(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    # The not-animated gif

    for isfloat in (False, True):
        for crop in (0, 1, 2):
            for colors in (0, 3, 4):
                if colors > 1 and sys.platform.startswith("darwin"):
                    continue  # quantize fails, see also png
                fname = fnamebase + "%i.%i.%i.gif" % (isfloat, crop, colors)
                rim = get_ref_im(colors, crop, isfloat)
                iio.imsave(fname, rim, format="GIF-FI")
                im = iio.imread(fname, format="GIF-FI")
                mul = 255 if isfloat else 1
                if colors in (0, 1):
                    im = im[:, :, 0]
                else:
                    im = im[:, :, :3]
                    rim = rim[:, :, :3]
                assert_close(rim * mul, im, 1.1)  # lossless

    # Parameter fail
    raises(TypeError, iio.imread, fname, notavalidkwarg=True, format="GIF-FI")
    raises(
        TypeError,
        iio.imsave,
        fnamebase + "1.gif",
        im,
        notavalidk=True,
        format="GIF-FI",
    )


@pytest.mark.skipif(
    sys.platform.startswith("darwin"),
    reason="On OSX quantization of freeimage is unstable",
)
@pytest.mark.needs_internet
def test_animated_gif(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    # Get images
    im = get_ref_im(4, 0, 0)
    ims = []
    for i in range(10):
        im = im.copy()
        im[:, -5:, 0] = i * 20
        ims.append(im)

    # Store - animated GIF always poops out RGB
    for isfloat in (False, True):
        for colors in (3, 4):
            ims1 = ims[:]
            if isfloat:
                ims1 = [x.astype(np.float32) / 256 for x in ims1]
            ims1 = [x[:, :, :colors] for x in ims1]
            fname = fnamebase + ".animated.%i.gif" % colors
            iio.mimsave(fname, ims1, duration=0.2, format="GIF-FI")
            # Retrieve
            ims2 = iio.mimread(fname, format="GIF-FI")
            ims1 = [x[:, :, :3] for x in ims]  # fresh ref
            ims2 = [x[:, :, :3] for x in ims2]  # discart alpha
            for im1, im2 in zip(ims1, ims2):
                assert_close(im1, im2, 1.1)

    # We can also store grayscale
    fname = fnamebase + ".animated.%i.gif" % 1
    iio.mimsave(fname, [x[:, :, 0] for x in ims], duration=0.2, format="GIF-FI")
    iio.mimsave(fname, [x[:, :, :1] for x in ims], duration=0.2, format="GIF-FI")

    # Irragular duration. You probably want to check this manually (I did)
    duration = [0.1 for i in ims]
    for i in [2, 5, 7]:
        duration[i] = 0.5
    iio.mimsave(
        fnamebase + ".animated_irr.gif", ims, duration=duration, format="GIF-FI"
    )

    # Other parameters
    iio.mimsave(fnamebase + ".animated.loop2.gif", ims, loop=2, fps=20, format="GIF-FI")
    R = iio.read(fnamebase + ".animated.loop2.gif", format="GIF-FI")
    W = iio.save(
        fnamebase + ".animated.palettes100.gif",
        palettesize=100,
        format="GIF-FI",
    )
    assert W._palettesize == 128
    # Fail
    raises(IndexError, R.get_meta_data, -1)
    raises(
        ValueError,
        iio.mimsave,
        fname,
        ims,
        quantizer="foo",
        format="GIF-FI",
    )
    raises(ValueError, iio.mimsave, fname, ims, duration="foo", format="GIF-FI")

    # Add one duplicate image to ims to touch subractangle with not change
    ims.append(ims[-1])

    # Test subrectangles
    iio.mimsave(fnamebase + ".subno.gif", ims, subrectangles=False, format="GIF-FI")
    iio.mimsave(fnamebase + ".subyes.gif", ims, subrectangles=True, format="GIF-FI")
    s1 = os.stat(fnamebase + ".subno.gif").st_size
    s2 = os.stat(fnamebase + ".subyes.gif").st_size
    assert s2 < s1

    # Meta (dummy, because always {}
    assert isinstance(iio.read(fname).get_meta_data(), dict)


@pytest.mark.needs_internet
def test_ico(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    for isfloat in (False, True):
        for crop in (0,):
            for colors in (1, 3, 4):
                fname = fnamebase + "%i.%i.%i.ico" % (isfloat, crop, colors)
                rim = get_ref_im(colors, crop, isfloat)
                rim = rim[:32, :32]  # ico needs nice size
                iio.imsave(fname, rim, format="ICO-FI")
                im = iio.imread(fname, format="ICO-FI")
                mul = 255 if isfloat else 1
                assert_close(rim * mul, im, 0.1)  # lossless

    # Meta data
    R = iio.read(fnamebase + "0.0.1.ico", format="ICO-FI")
    assert isinstance(R.get_meta_data(0), dict)
    assert isinstance(R.get_meta_data(None), dict)  # But this print warning
    R.close()
    writer = iio.save(fnamebase + "I.ico", format="ICO-FI")
    writer.set_meta_data({})
    writer.close()

    # Parameters. Note that with makealpha, RGBA images are read in incorrectly
    im = iio.imread(fnamebase + "0.0.1.ico", makealpha=True, format="ICO-FI")
    assert im.ndim == 3 and im.shape[-1] == 4

    # Parameter fail
    raises(TypeError, iio.imread, fname, notavalidkwarg=True, format="ICO-FI")
    raises(
        TypeError,
        iio.imsave,
        fnamebase + "1.ico",
        im,
        notavalidk=True,
        format="ICO-FI",
    )


# Skip on Windows xref: https://github.com/imageio/imageio/issues/21
@pytest.mark.skipif(
    platform.system() == "Windows",
    reason="Windows has a known issue with multi-icon files",
)
@pytest.mark.needs_internet
def test_multi_icon_ico(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    im = get_ref_im(4, 0, 0)[:32, :32]
    ims = [np.repeat(np.repeat(im, i, 1), i, 0) for i in (1, 2)]  # SegF on win
    ims = im, np.column_stack((im, im)), np.vstack((im, im))  # error on win
    iio.mimsave(fnamebase + "I2.ico", ims, format="ICO-FI")
    ims2 = iio.mimread(fnamebase + "I2.ico", format="ICO-FI")
    for im1, im2 in zip(ims, ims2):
        assert_close(im1, im2, 0.1)


@pytest.mark.skip("MNG seems broken in FreeImage")
def test_mng(setup_library, test_images):
    iio.imread(test_images / "mngexample.mng")


@pytest.mark.needs_internet
def test_pnm(setup_library, tmp_path):
    fnamebase = str(tmp_path / "test")

    for useAscii in (True, False):
        for crop in (0, 1, 2):
            for colors in (0, 1, 3):
                fname = fnamebase
                fname += "%i.%i.%i.ppm" % (useAscii, crop, colors)
                rim = get_ref_im(colors, crop, isfloat=False)
                iio.imsave(fname, rim, use_ascii=useAscii, format="PPM-FI")
                im = iio.imread(fname, format="PPM-FI")
                assert_close(rim, im, 0.1)  # lossless

                # Parameter fail
                raises(
                    TypeError,
                    iio.imread,
                    fname,
                    notavalidkwarg=True,
                    format="PPM-FI",
                )
                raises(
                    TypeError,
                    iio.imsave,
                    fname,
                    im,
                    notavalidk=True,
                    format="PPM-FI",
                )


def test_gamma_correction(setup_library, test_images):
    fname = test_images / "kodim03.png"

    # Load image three times
    im1 = iio.imread(fname, format="PNG-FI")
    im2 = iio.imread(fname, ignoregamma=True, format="PNG-FI")
    im3 = iio.imread(fname, ignoregamma=False, format="PNG-FI")

    # Default is to ignore gamma
    assert np.all(im1 == im2)

    # Test result depending of application of gamma
    assert im1.mean() == im2.mean()

    # TODO: We have assert im2.mean() == im3.mean()
    # But this is wrong, we want: assert im2.mean() < im3.mean()

    # test_regression_302
    for im in (im1, im2, im3):
        assert im.shape == (512, 768, 3) and im.dtype == "uint8"


@pytest.mark.needs_internet
def test_improps(setup_library, test_images):
    props = iio3.improps(test_images / "kodim03.png", plugin="PNG-FI")

    assert props.shape == (512, 768, 3)
    assert props.dtype == np.uint8
    assert props.is_batch is False


@pytest.mark.needs_internet
def test_exr_write(setup_library):
    expected = np.full((128, 128, 3), 0.42, dtype=np.float32)
    buffer = iio3.imwrite("<bytes>", expected, extension=".exr")

    actual = iio3.imread(buffer)

    np.allclose(actual, expected)
