from pathlib import Path
from os import path, mkdir, rename
import numpy as np
import pytest
from tomoscan.esrf.volume import (
    EDFVolume,
    HDF5Volume,
    JP2KVolume,
    MultiTIFFVolume,
    TIFFVolume,
)
from tomoscan.esrf.volume.jp2kvolume import has_glymur
from nabu.io.writer import merge_hdf5_files
from nabu.io.cast_volume import remove_volume


def test_remove_single_frame_volume(tmpdir):
    """
    Test volume removal for tiff, jp2 and EDF
    """
    # Have to use a not-too-small size because of jp2k
    data = np.arange(10 * 40 * 50, dtype="f").reshape((10, 40, 50))

    volume_classes = [EDFVolume, TIFFVolume, JP2KVolume]
    if not (has_glymur):
        volume_classes.pop()
    for volume_cls in volume_classes:
        ext = volume_cls.DEFAULT_DATA_EXTENSION
        folder = path.join(tmpdir, f"{ext}_vol")
        volume_basename = f"{ext}_basename"

        vol_writer = volume_cls(folder=folder, volume_basename=volume_basename, overwrite=True, start_index=0)
        vol_writer.data = data
        vol_writer.save()

        vol_reader = volume_cls(folder=folder, volume_basename=volume_basename)
        assert path.isdir(folder), f"Expected to find a folder f{folder}"
        remove_volume(vol_reader)
        assert not (path.isdir(folder)), f"Expected to have removed the folder f{folder}"

        vol_writer.save()
        vol_reader = volume_cls(folder=folder, volume_basename=volume_basename)
        Path(path.join(folder, f"unexpected.{ext}")).touch()
        with pytest.raises(RuntimeError) as exc:
            remove_volume(vol_reader, check=True)
        assert "Unexpected files present" in str(exc.value), "Expected check to find extraneous files"


def test_remove_multiframe_volume(tmpdir):
    """
    Test volume removal for "multiframe" formats (HDF5, tiff3D)
    The HDF5 files considered in this test do not have virtual sources
    """
    data = np.arange(3 * 4 * 5, dtype="f").reshape((3, 4, 5))

    for ext, volume_cls in {"h5": HDF5Volume, "tiff": MultiTIFFVolume}.items():
        file_path = path.join(tmpdir, f"{ext}_vol.{ext}")

        init_kwargs = {"file_path": file_path}
        if ext == "h5":
            init_kwargs["data_path"] = "entry"
        vol_writer = volume_cls(**init_kwargs)
        vol_writer.data = data
        vol_writer.save()

        vol_reader = volume_cls(**init_kwargs)
        assert path.isfile(file_path), f"Expected to find a {ext} volume at {file_path}"
        remove_volume(vol_reader)
        assert not (path.isfile(file_path)), f"Expected to have removed f{file_path}"


def test_remove_hdf5_multiple_entries(tmpdir):
    data = np.arange(3 * 4 * 5, dtype="f").reshape((3, 4, 5))
    file_path = path.join(tmpdir, "h5_vol.h5")
    vol_writer_1 = HDF5Volume(file_path=file_path, data_path="entry0000")
    vol_writer_1.data = data
    vol_writer_1.save()
    vol_writer_2 = HDF5Volume(file_path=file_path, data_path="entry0001")
    vol_writer_2.data = data + 10
    vol_writer_2.save()
    vol_reader = HDF5Volume(file_path=file_path, data_path="entry0000")
    with pytest.raises(NotImplementedError) as exc:
        remove_volume(vol_reader, check=True)
    assert "Removing a HDF5 volume with more than one entry is not supported" in str(
        exc.value
    ), "Expected an error message"


def test_remove_nabu_hdf5_reconstruction(tmpdir):
    """
    Test removal of HDF5 reconstruction generated by nabu (i.e with virtual sources)
    """

    entry = "entry"
    process_name = "reconstruction"

    master_file_path = path.join(tmpdir, "sample_naburec.hdf5")
    associated_dir = path.join(tmpdir, "sample_naburec")
    if not (path.isdir(associated_dir)):
        mkdir(associated_dir)

    n_chunks = 5
    local_files = []
    for i in range(n_chunks):
        fname = "sample_naburec_%06d.h5" % i
        partial_rec_abspath = path.join(associated_dir, fname)
        local_files.append(f"sample_naburec/{fname}")
        # local_files.append(fname)
        vol = HDF5Volume(partial_rec_abspath, data_path=f"{entry}/{process_name}")
        vol.data = np.arange(3 * 4 * 5, dtype="f").reshape((3, 4, 5))
        vol.save()

    h5_path = f"{entry}/{process_name}/results/data"

    merge_hdf5_files(
        local_files,
        h5_path,
        master_file_path,
        process_name,
        output_entry=entry,
        output_filemode="a",
        processing_index=0,
        config=None,
        base_dir=path.dirname(associated_dir),
        axis=0,
        overwrite=True,
    )

    assert path.isfile(master_file_path), f"Expected to find the master file at {master_file_path}"
    assert path.isdir(associated_dir)
    for local_file in local_files:
        partial_rec_file = path.join(tmpdir, local_file)
        assert path.isfile(partial_rec_file), f"Expected to find partial file number {i} at {partial_rec_file}"

    # Check that the virtual links are handled properly
    #  sample_rec.hdf5   should reference sample_rec/sample_rec_{i}.h5
    renamed_master_file_path = (
        path.join(path.dirname(master_file_path), path.basename(master_file_path).split(".")[0]) + "_renamed" + ".h5"
    )
    rename(master_file_path, renamed_master_file_path)
    h5_vol = HDF5Volume(file_path=renamed_master_file_path, data_path=f"{entry}/{process_name}")
    with pytest.raises(ValueError) as exc:
        remove_volume(h5_vol)
    expected_error_message = f"The virtual sources in {renamed_master_file_path}:{process_name}/results/data reference the directory sample_naburec, but expected was sample_naburec_renamed"
    assert str(exc.value) == expected_error_message

    # Check removal in normal circumstances
    rename(renamed_master_file_path, master_file_path)
    h5_vol = HDF5Volume(file_path=master_file_path, data_path=f"{entry}/{process_name}")
    remove_volume(h5_vol)
    assert not (path.isfile(master_file_path)), f"Expected to find the master file at {master_file_path}"
    assert not (path.isdir(associated_dir))
