"""Tests for ZipStore support in zarrnii."""

import os
import tempfile

import dask.array as da
import numpy as np
import pytest

from zarrnii import ZarrNii


class TestZipStoreSupport:
    """Test ZipStore functionality for OME-Zarr files."""

    @pytest.fixture
    def sample_zarrnii(self):
        """Create a sample ZarrNii object for testing."""
        # Create test data
        data = da.ones((1, 16, 32, 32), dtype=np.uint16, chunks=(1, 8, 16, 16))
        return ZarrNii.from_darr(data, axes_order="ZYX", orientation="RAS")

    def test_save_to_zip_file(self, sample_zarrnii):
        """Test saving OME-Zarr to a .zip file."""
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = os.path.join(tmpdir, "test.ome.zarr.zip")

            # Save to ZIP file
            result = sample_zarrnii.to_ome_zarr(zip_path, max_layer=2)

            # Should return self for method chaining
            assert result is sample_zarrnii

            # Check that ZIP file was created
            assert os.path.exists(zip_path)
            assert os.path.getsize(zip_path) > 0

            # ZIP file should be readable as zip
            import zipfile

            with zipfile.ZipFile(zip_path, "r") as zf:
                files = zf.namelist()
                assert ".zgroup" in files
                assert ".zattrs" in files
                # Should have multiscale levels
                assert any("0" in name for name in files)

    def test_load_from_zip_file(self, sample_zarrnii):
        """Test loading OME-Zarr from a .zip file."""
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = os.path.join(tmpdir, "test.ome.zarr.zip")

            # Save to ZIP file first
            sample_zarrnii.to_ome_zarr(zip_path, max_layer=2)

            # Load from ZIP file
            loaded_zarrnii = ZarrNii.from_ome_zarr(zip_path, level=0)

            # Verify loaded data
            assert loaded_zarrnii.darr.shape == sample_zarrnii.darr.shape
            assert loaded_zarrnii.axes_order == sample_zarrnii.axes_order
            assert loaded_zarrnii.orientation == sample_zarrnii.orientation

            # Data should be the same
            np.testing.assert_array_equal(
                loaded_zarrnii.darr.compute(), sample_zarrnii.darr.compute()
            )

    def test_round_trip_zip_storage(self, sample_zarrnii):
        """Test round-trip save and load with ZIP storage."""
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = os.path.join(tmpdir, "roundtrip.ome.zarr.zip")

            # Round trip
            sample_zarrnii.to_ome_zarr(zip_path, max_layer=3)
            loaded = ZarrNii.from_ome_zarr(zip_path, level=0)

            # Verify all properties are preserved
            assert loaded.darr.shape == sample_zarrnii.darr.shape
            assert loaded.axes_order == sample_zarrnii.axes_order
            assert loaded.orientation == sample_zarrnii.orientation

            # Verify data integrity
            original_data = sample_zarrnii.darr.compute()
            loaded_data = loaded.darr.compute()
            np.testing.assert_array_equal(loaded_data, original_data)

    def test_zip_vs_regular_directory_equivalence(self, sample_zarrnii):
        """Test that ZIP and regular directory storage produce equivalent results."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Save to regular directory
            regular_path = os.path.join(tmpdir, "regular.ome.zarr")
            sample_zarrnii.to_ome_zarr(regular_path, max_layer=2)

            # Save to ZIP file
            zip_path = os.path.join(tmpdir, "zipped.ome.zarr.zip")
            sample_zarrnii.to_ome_zarr(zip_path, max_layer=2)

            # Load both and compare
            loaded_regular = ZarrNii.from_ome_zarr(regular_path, level=0)
            loaded_zip = ZarrNii.from_ome_zarr(zip_path, level=0)

            # Should have identical data
            np.testing.assert_array_equal(
                loaded_regular.darr.compute(), loaded_zip.darr.compute()
            )

            # Should have identical metadata
            assert loaded_regular.axes_order == loaded_zip.axes_order
            assert loaded_regular.orientation == loaded_zip.orientation
            assert loaded_regular.darr.shape == loaded_zip.darr.shape

    def test_multi_level_zip_access(self, sample_zarrnii):
        """Test accessing different pyramid levels in ZIP files."""
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = os.path.join(tmpdir, "multilevel.ome.zarr.zip")

            # Save with multiple levels
            sample_zarrnii.to_ome_zarr(zip_path, max_layer=3)

            # Load different levels
            level0 = ZarrNii.from_ome_zarr(zip_path, level=0)
            level1 = ZarrNii.from_ome_zarr(zip_path, level=1)
            level2 = ZarrNii.from_ome_zarr(zip_path, level=2)

            # Verify pyramid structure
            assert level0.darr.shape == sample_zarrnii.darr.shape
            # Each level should be roughly half the size in index 2 or 3 (y or x)
            assert level1.darr.shape[2] <= level0.darr.shape[2] // 2 + 1
            assert level2.darr.shape[2] <= level1.darr.shape[2] // 2 + 1

    def test_zip_file_extension_detection(self, sample_zarrnii):
        """Test that .zip extension is properly detected for ZipStore handling."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Test various zip extensions
            extensions = [".zip", ".ome.zarr.zip", ".zarr.zip"]

            for ext in extensions:
                zip_path = os.path.join(tmpdir, f"test{ext}")

                # Should successfully save and load
                sample_zarrnii.to_ome_zarr(zip_path)
                loaded = ZarrNii.from_ome_zarr(zip_path)

                assert loaded.darr.shape == sample_zarrnii.darr.shape
                np.testing.assert_array_equal(
                    loaded.darr.compute(), sample_zarrnii.darr.compute()
                )

    def test_zip_error_handling(self):
        """Test error handling for invalid ZIP files."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create an empty file with .zip extension
            invalid_zip = os.path.join(tmpdir, "invalid.ome.zarr.zip")
            with open(invalid_zip, "w") as f:
                f.write("not a zip file")

            # Should raise appropriate error when trying to read
            with pytest.raises(Exception):
                ZarrNii.from_ome_zarr(invalid_zip)
