import sys
import tempfile
from pathlib import Path

import numpy as np
import pytest
from cluster_tools import get_executor

import webknossos as wk
from webknossos import Annotation, SegmentationLayer
from webknossos.annotation.volume_layer import VolumeLayerEditMode
from webknossos.dataset import DataFormat
from webknossos.geometry import BoundingBox, Vec3Int

from .constants import TESTDATA_DIR, TESTOUTPUT_DIR


def test_annotation_from_wkw_zip_file() -> None:
    annotation = wk.Annotation.load(
        TESTDATA_DIR
        / "annotations"
        / "l4dense_motta_et_al_demo_v2__explorational__4a6356.zip"
    )

    assert annotation.dataset_name == "l4dense_motta_et_al_demo_v2"
    assert annotation.organization_id == "scalable_minds"
    assert annotation.dataset_id == "67b7c81b0100004d00f5ab31"
    assert annotation.owner_name == "Philipp Otto"
    assert annotation.annotation_id == "61c20205010000cc004a6356"
    assert "timestamp" in annotation.metadata
    assert len(list(annotation.get_volume_layer_names())) == 1
    assert len(list(annotation.skeleton.flattened_trees())) == 1

    annotation.save(TESTOUTPUT_DIR / "test_dummy.zip")
    copied_annotation = wk.Annotation.load(TESTOUTPUT_DIR / "test_dummy.zip")

    assert copied_annotation.dataset_name == "l4dense_motta_et_al_demo_v2"
    assert copied_annotation.organization_id == "scalable_minds"
    assert copied_annotation.dataset_id == "67b7c81b0100004d00f5ab31"
    assert copied_annotation.owner_name == "Philipp Otto"
    assert copied_annotation.annotation_id == "61c20205010000cc004a6356"
    assert "timestamp" in copied_annotation.metadata
    assert len(list(copied_annotation.get_volume_layer_names())) == 1
    assert len(list(copied_annotation.skeleton.flattened_trees())) == 1

    copied_annotation.add_volume_layer(
        name="new_volume_layer",
        dtype=np.uint32,
    )
    assert len(list(copied_annotation.get_volume_layer_names())) == 2
    copied_annotation.delete_volume_layer(volume_layer_name="new_volume_layer")
    assert len(list(copied_annotation.get_volume_layer_names())) == 1

    with annotation.temporary_volume_layer_copy() as volume_layer:
        input_annotation_mag = volume_layer.get_finest_mag()
        voxel_id = input_annotation_mag.read(
            absolute_offset=Vec3Int(2830, 4356, 1792), size=Vec3Int.full(1)
        )

        assert voxel_id == 2504698


def test_annotation_from_zarr3_zip_file() -> None:
    annotation = wk.Annotation.load(
        TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
    )

    with annotation.temporary_volume_layer_copy() as volume_layer:
        assert volume_layer.data_format == DataFormat.Zarr3
        assert volume_layer.bounding_box == BoundingBox(
            (3072, 3072, 512), (1024, 1024, 1024)
        )
        input_annotation_mag = volume_layer.get_mag("2-2-1")
        voxel_id = input_annotation_mag.read(
            absolute_offset=Vec3Int(3630, 3502, 1024), size=Vec3Int(2, 2, 1)
        )

        assert np.array_equiv(voxel_id, 1)


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="Windows does not like file handles that are not properly closed. Should probably be fixed in the future.",
)
def test_annotation_from_nml_file() -> None:
    snapshot_path = TESTDATA_DIR / "nmls" / "generated_annotation_snapshot.nml"

    annotation = wk.Annotation.load(snapshot_path)

    assert annotation.dataset_name == "My Dataset"
    assert annotation.organization_id is None
    assert len(list(annotation.skeleton.flattened_trees())) == 3

    annotation.save(TESTOUTPUT_DIR / "test_dummy.zip")
    copied_annotation = wk.Annotation.load(TESTOUTPUT_DIR / "test_dummy.zip")
    assert copied_annotation.dataset_name == "My Dataset"
    assert copied_annotation.organization_id is None
    assert len(list(copied_annotation.skeleton.flattened_trees())) == 3


def test_annotation_from_file_with_multi_volume() -> None:
    annotation = wk.Annotation.load(
        TESTDATA_DIR / "annotations" / "multi_volume_example_CREMI.zip"
    )

    volume_names = sorted(annotation.get_volume_layer_names())

    assert volume_names == ["Volume", "Volume_2"]

    # Read from first layer
    with annotation.temporary_volume_layer_copy(
        volume_layer_name=volume_names[0]
    ) as layer:
        read_voxel = layer.get_finest_mag().read(
            absolute_offset=(590, 512, 16),
            size=(1, 1, 1),
        )
        assert read_voxel == 7718, (
            f"Expected to see voxel id 7718, but saw {read_voxel} instead."
        )

        read_voxel = layer.get_finest_mag().read(
            absolute_offset=(490, 512, 16),
            size=(1, 1, 1),
        )
        # When viewing the annotation online, this segment id will be 284.
        # However, this is fallback data which is not included in this annotation.
        # Therefore, we expect to read a 0 here.
        assert read_voxel == 0, (
            f"Expected to see voxel id 0, but saw {read_voxel} instead."
        )

    # Read from second layer
    with annotation.temporary_volume_layer_copy(
        volume_layer_name=volume_names[1]
    ) as layer:
        read_voxel = layer.get_finest_mag().read(
            absolute_offset=(590, 512, 16),
            size=(1, 1, 1),
        )
        assert read_voxel == 1, (
            f"Expected to see voxel id 1, but saw {read_voxel} instead."
        )

        read_voxel = layer.get_finest_mag().read(
            absolute_offset=(490, 512, 16),
            size=(1, 1, 1),
        )
        assert read_voxel == 0, (
            f"Expected to see voxel id 0, but saw {read_voxel} instead."
        )

    # Reading from not-existing layer should raise an error
    with pytest.raises(AssertionError):
        with annotation.temporary_volume_layer_copy(
            volume_layer_name="not existing name"
        ) as layer:
            pass


@pytest.mark.use_proxay
def test_dataset_access_via_annotation() -> None:
    # This is a regression test for a bug occurring when the dataset name was changed
    # while it was referenced in an annotation.

    # load a remote dataset
    remote_ds = wk.Dataset.open_remote(
        "http://localhost:9000/datasets/Organization_X/l4_sample"
    )

    # upload an annotation with a reference to a remote dataset
    annotation_from_file = wk.Annotation.load(
        TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
    )
    annotation_from_file.organization_id = "Organization_X"
    annotation_from_file.dataset_name = remote_ds.name
    annotation_from_file.dataset_id = remote_ds._dataset_id
    url = annotation_from_file.upload()

    # change the name of the remote dataset
    remote_ds.name = "some_other_name"

    # check if the annotations dataset can be accessed
    wk.Annotation.open_as_remote_dataset(url)

    # change the name of the remote dataset back to the original name
    remote_ds.name = "l4_sample"


@pytest.mark.use_proxay
def test_remote_annotation_list() -> None:
    path = TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
    annotation_from_file = wk.Annotation.load(path)
    annotation_from_file.organization_id = "Organization_X"
    annotation_from_file.upload()
    annotations = wk.AnnotationInfo.get_remote_annotations()

    assert annotation_from_file.name in [a.name for a in annotations]


@pytest.mark.use_proxay
@pytest.mark.skipif(sys.platform == "win32", reason="too slow on windows")
def test_annotation_upload_download_roundtrip() -> None:
    path = TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
    annotation_from_file = wk.Annotation.load(path)
    annotation_from_file.organization_id = "Organization_X"
    url = annotation_from_file.upload()
    annotation = wk.Annotation.download(url)
    assert annotation.dataset_name == "l4_sample"
    assert len(list(annotation.skeleton.flattened_trees())) == 1

    mag = wk.Mag("16-16-4")
    node_bbox = wk.BoundingBox.from_points(
        next(annotation.skeleton.flattened_trees()).get_node_positions()
    ).align_with_mag(mag, ceil=True)
    ds = annotation.get_remote_annotation_dataset()

    mag_view = ds.layers["Volume"].get_mag(mag)
    annotated_data = mag_view.read(absolute_bounding_box=node_bbox)
    assert annotated_data.size > 10
    # assert (annotated_data == 1).all()
    assert mag_view.read(absolute_offset=(0, 0, 0), size=(16, 16, 4))[0, 0, 0, 0] == 0
    assert (
        mag_view.read(absolute_offset=(3600, 3488, 1024), size=(16, 16, 4))[0, 0, 0, 0]
        == 1
    )
    segment_info = annotation.get_volume_layer_segments("Volume")[1]
    assert segment_info.anchor_position == (3395, 3761, 1024)
    segment_info.name = "Test Segment"
    segment_info.color = (1, 0, 0, 1)

    annotation.save(TESTOUTPUT_DIR / "test_dummy_downloaded.zip")
    annotation = wk.Annotation.load(TESTOUTPUT_DIR / "test_dummy_downloaded.zip")
    assert annotation.dataset_name == "l4_sample"
    assert len(list(annotation.skeleton.flattened_trees())) == 1
    segment_info = annotation.get_volume_layer_segments("Volume")[1]
    assert segment_info.anchor_position == (3395, 3761, 1024)
    assert segment_info.name == "Test Segment"
    assert segment_info.color == (1, 0, 0, 1)


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="Windows does not like file handles that are not properly closed. Should probably be fixed in the future.",
)
def test_reading_bounding_boxes() -> None:
    def check_properties(annotation: wk.Annotation) -> None:
        assert len(annotation.user_bounding_boxes) == 2
        assert annotation.user_bounding_boxes[0].topleft.x == 2371
        assert annotation.user_bounding_boxes[0].name == "Bounding box 1"
        assert annotation.user_bounding_boxes[0].is_visible

        assert annotation.user_bounding_boxes[1] == BoundingBox(
            (371, 4063, 1676), (891, 579, 232)
        )
        assert annotation.user_bounding_boxes[1].name == "Bounding box 2"
        assert not annotation.user_bounding_boxes[1].is_visible
        assert annotation.user_bounding_boxes[1].color == (
            0.2705882489681244,
            0.6470588445663452,
            0.19607843458652496,
            1.0,
        )

    # Check loading checked-in file
    input_path = TESTDATA_DIR / "annotations" / "bounding-boxes-example.zip"
    annotation = wk.Annotation.load(input_path)
    check_properties(annotation)

    # Check exporting and re-reading checked-in file (roundtrip)
    with tempfile.TemporaryDirectory() as tmp_dir:
        output_path = Path(tmp_dir) / "serialized.zip"
        annotation.save(output_path)

        annotation_deserialized = wk.Annotation.load(output_path)
        check_properties(annotation_deserialized)


@pytest.mark.use_proxay
def test_bounding_box_roundtrip() -> None:
    ds = wk.Dataset.open_remote("l4_sample")

    annotation_before = wk.Annotation(
        name="test_bounding_box_roundtrip",
        dataset_name=ds.name,
        voxel_size=ds.voxel_size,
    )
    group = annotation_before.skeleton.add_group("a group")
    tree = group.add_tree("a tree")
    tree.add_node(position=(0, 0, 0), comment="node 1")

    annotation_before.user_bounding_boxes = [
        wk.BoundingBox((1024, 512, 128), (64, 64, 64))
    ]
    color = (0.5, 0, 0.2, 1)
    annotation_before.task_bounding_box = (
        wk.BoundingBox((10, 10, 10), (5, 5, 5)).with_name("task_bbox").with_color(color)
    )

    annotation_url = annotation_before.upload()
    annotation_after = wk.Annotation.download(annotation_url)

    # task bounding box is appended to user bounding boxes when uploading a normal annotation:
    assert (
        annotation_after.user_bounding_boxes
        == annotation_before.user_bounding_boxes + [annotation_before.task_bounding_box]
    )


def test_empty_volume_annotation() -> None:
    a = wk.Annotation.load(TESTDATA_DIR / "annotations" / "empty_volume_annotation.zip")
    with a.temporary_volume_layer_copy() as layer:
        assert layer.bounding_box == wk.BoundingBox.empty()
        assert layer.largest_segment_id == 0
        assert set(layer.mags) == set(
            [
                wk.Mag(1),
                wk.Mag((2, 2, 1)),
                wk.Mag((4, 4, 1)),
                wk.Mag((8, 8, 1)),
                wk.Mag((16, 16, 1)),
            ]
        )


@pytest.mark.parametrize(
    "nml_path",
    [
        TESTDATA_DIR / "annotations" / "nml_with_volumes.nml",
        TESTDATA_DIR / "annotations" / "nml_with_volumes.zip",
    ],
)
def test_nml_with_volumes(nml_path: Path) -> None:
    if nml_path.suffix == ".zip":
        with pytest.warns(UserWarning, match="location is not referenced in the NML"):
            a = wk.Annotation.load(nml_path)
    else:
        a = wk.Annotation.load(nml_path)
    segment_info = a.get_volume_layer_segments("segmentation")
    assert set(segment_info) == set([2504698])
    assert segment_info[2504698] == wk.SegmentInformation(
        name="test_segment",
        anchor_position=Vec3Int(3581, 3585, 1024),
        color=None,
        metadata={},
    )


def test_segment_metadata(tmp_path: Path) -> None:
    annotation = wk.Annotation.load(
        TESTDATA_DIR / "annotations" / "nml_with_volumes.zip"
    )
    annotation.get_volume_layer_segments("segmentation")[2504698].metadata[
        "test_segment"
    ] = "test"
    annotation.save(tmp_path / "test.zip")
    tmp_annotation = wk.Annotation.load(tmp_path / "test.zip")
    assert (
        tmp_annotation.get_volume_layer_segments("segmentation")[2504698].metadata[
            "test_segment"
        ]
        == "test"
    )


def test_tree_metadata(tmp_path: Path) -> None:
    annotation = wk.Annotation.load(
        TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
    )
    list(annotation.skeleton.flattened_trees())[0].metadata["test_tree"] = "test"
    annotation.save(tmp_path / "test.zip")
    tmp_annotation = wk.Annotation.load(tmp_path / "test.zip")
    assert (
        list(tmp_annotation.skeleton.flattened_trees())[0].metadata["test_tree"]
        == "test"
    )


@pytest.mark.parametrize(
    "edit_mode", [VolumeLayerEditMode.MEMORY, VolumeLayerEditMode.TEMPORARY_DIRECTORY]
)
@pytest.mark.parametrize("executor", ["sequential", "multiprocessing"])
def test_edit_volume_annotation(edit_mode: VolumeLayerEditMode, executor: str) -> None:
    dtype = np.uint32
    data = np.ones((1, 10, 10, 10), dtype=dtype)
    ann = wk.Annotation(
        name="my_annotation",
        dataset_name="sample_dataset",
        voxel_size=(11.2, 11.2, 25.0),
    )

    volume_layer = ann.add_volume_layer(
        name="segmentation",
        dtype=dtype,
    )
    if edit_mode == VolumeLayerEditMode.MEMORY and executor == "multiprocessing":
        with pytest.raises(ValueError, match="SequentialExecutor"):
            with volume_layer.edit(
                edit_mode=edit_mode, executor=get_executor(executor)
            ) as seg_layer:
                pass
    else:
        with volume_layer.edit(
            edit_mode=edit_mode, executor=get_executor(executor)
        ) as seg_layer:
            assert isinstance(seg_layer, SegmentationLayer)
            mag = seg_layer.add_mag(1)
            mag.write(data, absolute_offset=(0, 0, 0), allow_resize=True)
        with volume_layer.edit(edit_mode=edit_mode) as seg_layer:
            assert len(seg_layer.mags) == 1
            mag = seg_layer.get_mag(1)
            read_data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
            assert np.array_equal(data, read_data)


def test_edited_volume_annotation_format() -> None:
    import zipfile

    import tensorstore

    path = TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
    ann = Annotation.load(path)
    data = np.ones(shape=(10, 10, 10))

    volume_layer = ann.add_volume_layer(
        name="segmentation",
        dtype=np.uint32,
    )
    with volume_layer.edit() as seg_layer:
        mag_view = seg_layer.add_mag(1)
        mag_view.write(data, allow_resize=True)

    save_path = TESTOUTPUT_DIR / "saved_annotation.zip"
    ann.save(save_path)
    unpack_dir = TESTOUTPUT_DIR / "unpacked_annotation"
    with zipfile.ZipFile(save_path, "r") as zip_ref:
        zip_ref.extractall(unpack_dir)

    # test for the format assumptions as mentioned in https://github.com/scalableminds/webknossos/issues/8604
    ts = tensorstore.open(
        {
            "driver": "zarr3",
            "kvstore": {
                "driver": "zip",
                "path": "volumeAnnotationData/1/",
                "base": {
                    "driver": "file",
                    "path": str(unpack_dir / "data_1_segmentation.zip"),
                },
            },
        },
        create=False,
        open=True,
    ).result()
    metadata = ts.spec().to_json()["metadata"]

    assert metadata["chunk_key_encoding"] == {
        "configuration": {"separator": "."},
        "name": "default",
    }
    assert ["transpose", "bytes", "blosc"] == [
        codec["name"] for codec in metadata["codecs"]
    ]
    data_read = ts.read().result()[0, :10, :10, :10]
    assert np.array_equal(data, data_read)


@pytest.mark.parametrize(
    "edit_mode", [VolumeLayerEditMode.MEMORY, VolumeLayerEditMode.TEMPORARY_DIRECTORY]
)
def test_edited_volume_annotation_save_load(edit_mode: VolumeLayerEditMode) -> None:
    data = np.ones((1, 10, 10, 10))

    ann = wk.Annotation(
        name="my_annotation",
        dataset_name="sample_dataset",
        voxel_size=(11.2, 11.2, 25.0),
    )

    volume_layer = ann.add_volume_layer(name="segmentation", dtype=np.uint32)
    with volume_layer.edit(edit_mode=edit_mode) as seg_layer:
        mag_view = seg_layer.add_mag(1)
        mag_view.write(data, allow_resize=True)

    save_path = TESTOUTPUT_DIR / "annotation_saved.zip"
    ann.save(save_path)
    ann_loaded = Annotation.load(save_path)

    volume_layer_downloaded = ann_loaded.get_volume_layer("segmentation")

    with volume_layer_downloaded.edit(edit_mode=edit_mode) as seg_layer:
        assert len(seg_layer.mags) == 1
        mag = seg_layer.get_mag(1)
        read_data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
        assert np.array_equal(data, read_data)


@pytest.mark.use_proxay
def test_edited_volume_annotation_upload_download() -> None:
    data = np.ones((1, 10, 10, 10))

    ann = Annotation.load(
        TESTDATA_DIR / "annotations" / "l4_sample__explorational__suser__94b271.zip"
    )
    ann.organization_id = "Organization_X"

    volume_layer = ann.add_volume_layer(
        name="segmentation",
        dtype=np.uint32,
    )
    with volume_layer.edit() as seg_layer:
        mag_view = seg_layer.add_mag(1)
        mag_view.write(data, allow_resize=True)

    url = ann.upload()
    ann_downloaded = Annotation.download(
        url,
    )

    assert {layer.name for layer in ann_downloaded._volume_layers} == {
        "Volume",
        "segmentation",
    }

    volume_layer_downloaded = ann_downloaded.get_volume_layer("segmentation")

    with volume_layer_downloaded.edit() as seg_layer:
        assert len(seg_layer.mags) == 1
        mag = seg_layer.get_mag(1)
        read_data = mag.read(absolute_offset=(0, 0, 0), size=(10, 10, 10))
        assert np.array_equal(data, read_data)
