import logging
import os
from pathlib import Path
from typing import Any

import fsspec
import numpy as np
import pandas as pd
import pyarrow.fs
import pytest
import zarr
from pytest_lazy_fixtures import lf as lazy_fixture

import ray
from ray.data._internal.datasource import zarrv2_datasource
from ray.data.block import BlockAccessor
from ray.data.tests.conftest import *  # noqa: F401, F403


def _execute_read_tasks(tasks) -> pd.DataFrame:
    frames = [
        BlockAccessor.for_block(block).to_pandas() for task in tasks for block in task()
    ]
    return pd.concat(frames, ignore_index=True)


def _reconstruct_array(df: pd.DataFrame, array_name: str) -> np.ndarray:
    """Concatenate all chunks of one array from a long-form result frame."""
    sub = df[df["array"] == array_name].sort_values(
        "chunk_index", key=lambda col: col.map(tuple)
    )
    return np.concatenate(list(sub["chunk"]), axis=0)


def _write_real_zarr_store(
    store_path: Path,
    arrays: dict,  # {name: (data, chunks)}
) -> Path:
    """Write a real Zarr v2 store from numpy arrays and consolidate metadata."""
    root = zarr.open_group(str(store_path), mode="w")
    for name, (data, chunks) in arrays.items():
        root.create_dataset(name, data=data, chunks=chunks, dtype=data.dtype)
    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
    return store_path


@pytest.fixture
def zarrv2_group_store(tmp_path) -> Path:
    """Two arrays at the store root, both 2-D and 1-D, axis-0-aligned (shape[0]==5)."""
    return _write_real_zarr_store(
        tmp_path / "group.zarr",
        {
            "images": (np.arange(20, dtype="<i4").reshape(5, 4), (2, 4)),
            "nested": (np.arange(5, dtype="|u1"), (2,)),
        },
    )


@pytest.fixture
def zarrv2_root_store(tmp_path) -> Path:
    """Single-array store with the array sitting directly at the store root."""
    store_path = tmp_path / "root.zarr"
    arr = zarr.open(
        str(store_path),
        mode="w",
        shape=(5, 4),
        chunks=(2, 4),
        dtype="<i4",
    )
    arr[:] = np.arange(20, dtype="<i4").reshape(5, 4)
    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
    return store_path


@pytest.fixture
def local_fsspec_fs():
    """fsspec local filesystem (for parametrized cross-fs read tests)."""
    return fsspec.filesystem("file")


@pytest.fixture
def heterogeneous_zarrv2_store(tmp_path) -> Path:
    """A store mixing different ranks, shape[0]s, dtypes, and native chunk sizes.

    Mirrors the UMI-style real-world layout where ``data/*`` arrays share an
    axis-0 timestep count but differ in everything else, and ``meta/*``
    arrays live in a separate axis-0 universe entirely. The chunk-per-row
    datasource handles all of these in one read; nothing has to align.
    """
    store_path = tmp_path / "heterogeneous.zarr"
    root = zarr.open_group(str(store_path), mode="w")
    # 4-D image tensor with tiny axis-0 chunks (1 image per chunk).
    root.create_dataset(
        "data/camera0_rgb",
        data=np.arange(20 * 2 * 2 * 3, dtype="|u1").reshape(20, 2, 2, 3),
        chunks=(1, 2, 2, 3),
    )
    # 2-D pose array, same shape[0]=20, much larger axis-0 chunks (10).
    root.create_dataset(
        "data/robot0_eef_pos",
        data=np.arange(20 * 3, dtype="<f4").reshape(20, 3),
        chunks=(10, 3),
    )
    # Episode-boundary metadata: separate axis-0 universe.
    root.create_dataset(
        "meta/episode_ends",
        data=np.array([5, 12, 20], dtype="<i8"),
        chunks=(3,),
    )
    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
    return store_path


@pytest.fixture
def unconsolidated_zarrv2_store(tmp_path) -> Path:
    """Two arrays at the store root, no ``.zmetadata``.

    Exercises the no-``.zmetadata`` code paths (per-array ``.zarray``
    discovery and full-store walk) — the common shape of real-world stores
    behind plain HTTPS or other listing-less filesystems.
    """
    store_path = tmp_path / "unconsolidated.zarr"
    root = zarr.open_group(str(store_path), mode="w")
    root.create_dataset(
        "images", data=np.arange(20, dtype="<i4").reshape(5, 4), chunks=(2, 4)
    )
    root.create_dataset("nested", data=np.arange(5, dtype="|u1"), chunks=(2,))
    return store_path


@pytest.fixture
def aligned_zarrv2_store(tmp_path) -> Path:
    """Three arrays sharing ``shape[0]=8``, different ranks and native chunks.

    Models the UMI-style case where data arrays co-stride on the timestep
    axis but differ in everything else.
    """
    store_path = tmp_path / "aligned.zarr"
    root = zarr.open_group(str(store_path), mode="w")
    root.create_dataset(
        "img",
        data=np.arange(8 * 4 * 4 * 3, dtype="|u1").reshape(8, 4, 4, 3),
        chunks=(2, 4, 4, 3),
    )
    root.create_dataset(
        "state",
        data=np.arange(8 * 3, dtype="<f4").reshape(8, 3),
        chunks=(4, 3),  # different native axis-0 chunks than img
    )
    root.create_dataset(
        "label",
        data=np.arange(8, dtype="<i8"),
        chunks=(8,),
    )
    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
    return store_path


@pytest.fixture
def zarr_zip_store(tmp_path) -> Path:
    """A small Zarr store packed into a ``.zip`` for URL-detection tests."""
    src = tmp_path / "src.zarr"
    _write_real_zarr_store(
        src,
        {
            "data": (np.arange(12, dtype="<i4").reshape(6, 2), (3, 2)),
        },
    )
    zip_path = tmp_path / "store.zarr.zip"
    import shutil

    shutil.make_archive(
        base_name=str(tmp_path / "store.zarr"),
        format="zip",
        root_dir=str(src),
    )
    assert zip_path.exists()
    return zip_path


# ---------------------------------------------------------------------------
# Metadata discovery
# ---------------------------------------------------------------------------


def test_normalizes_requested_root_array_path(zarrv2_root_store):
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(zarrv2_root_store),
        array_paths=[""],
    )
    assert list(datasource._metadata_by_path) == [""]


def test_normalizes_requested_array_paths(zarrv2_group_store):
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(zarrv2_group_store),
        array_paths=["images/", "nested"],
    )
    assert list(datasource._metadata_by_path) == ["images", "nested"]


def test_rejects_missing_array_paths(zarrv2_group_store):
    with pytest.raises(
        ValueError,
        match=r"Array\(s\) not found: 'missing'\. Available: 'images', 'nested'",
    ):
        zarrv2_datasource.ZarrV2Datasource(
            str(zarrv2_group_store),
            array_paths=["missing"],
        )


def test_loads_per_array_zarray_without_zmetadata(unconsolidated_zarrv2_store):
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(unconsolidated_zarrv2_store),
        array_paths=["images", "nested"],
    )
    assert set(datasource._metadata_by_path) == {"images", "nested"}


def test_full_scan_discovers_arrays_without_zmetadata(unconsolidated_zarrv2_store):
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(unconsolidated_zarrv2_store),
        allow_full_metadata_scan=True,
    )
    assert set(datasource._metadata_by_path) == {"images", "nested"}


def test_requires_array_paths_or_full_scan_when_unconsolidated(
    unconsolidated_zarrv2_store,
):
    with pytest.raises(
        ValueError,
        match=(
            r"No array_paths were provided and this Zarr store does not "
            r"contain \.zmetadata"
        ),
    ):
        zarrv2_datasource.ZarrV2Datasource(str(unconsolidated_zarrv2_store))


def test_array_paths_missing_zarray_file_raises_value_error(
    unconsolidated_zarrv2_store,
):
    with pytest.raises(
        ValueError,
        match=r"Array path 'missing' not found",
    ):
        zarrv2_datasource.ZarrV2Datasource(
            str(unconsolidated_zarrv2_store),
            array_paths=["missing"],
        )


def test_local_scheme_pins_reads_to_driver_node(zarrv2_group_store):
    """``local://`` stores can't be read distributed; plain/cloud paths can."""
    local = zarrv2_datasource.ZarrV2Datasource("local://" + str(zarrv2_group_store))
    assert local.supports_distributed_reads is False

    plain = zarrv2_datasource.ZarrV2Datasource(str(zarrv2_group_store))
    assert plain.supports_distributed_reads is True


def test_consolidation_detected_via_open_consolidated(
    zarrv2_group_store, unconsolidated_zarrv2_store
):
    """``_consolidated`` reflects whether ``.zmetadata`` actually opened."""
    consolidated = zarrv2_datasource.ZarrV2Datasource(
        str(zarrv2_group_store), array_paths=["images"]
    )
    assert consolidated._consolidated is True

    unconsolidated = zarrv2_datasource.ZarrV2Datasource(
        str(unconsolidated_zarrv2_store), array_paths=["images"]
    )
    assert unconsolidated._consolidated is False


def test_array_paths_rejects_group_path(tmp_path):
    """Requesting a group path (not an array) on an unconsolidated store errors."""
    store_path = tmp_path / "withgroup.zarr"
    root = zarr.open_group(str(store_path), mode="w")
    grp = root.create_group("grp")
    grp.create_dataset("inner", data=np.arange(4, dtype="<i4"), chunks=(2,))
    # Not consolidated -> the per-array ``.zarray`` lookup path.
    with pytest.raises(ValueError, match="is a group, not an array"):
        zarrv2_datasource.ZarrV2Datasource(str(store_path), array_paths=["grp"])


def test_root_array_rejects_non_root_array_paths(zarrv2_root_store):
    """A single root-level array rejects array_paths that aren't the root ''."""
    with pytest.raises(ValueError, match="single root-level array"):
        zarrv2_datasource.ZarrV2Datasource(
            str(zarrv2_root_store), array_paths=["missing"]
        )


# ---------------------------------------------------------------------------
# chunk_shapes validation
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "chunk_shapes, match",
    [
        ("invalid", "positive integers"),
        ({"images": 1}, "positive integers"),
        ({"does_not_exist": [2]}, "Unknown array path"),
    ],
)
def test_rejects_invalid_chunk_shapes(zarrv2_group_store, chunk_shapes, match):
    with pytest.raises(ValueError, match=match):
        zarrv2_datasource.ZarrV2Datasource(
            str(zarrv2_group_store), chunk_shapes=chunk_shapes
        )


@pytest.mark.parametrize(
    "chunk_shapes,array_paths,expected",
    [
        # No chunk_shapes: every array reads at its native chunk size.
        # 4-D image with tiny chunks coexists with 2-D pose with big chunks —
        # nothing is forced into a shared min/max.
        (
            None,
            None,
            {
                "data/camera0_rgb": (1, 2, 2, 3),
                "data/robot0_eef_pos": (10, 3),
                "meta/episode_ends": (3,),
            },
        ),
        # ``[5]`` prefix overrides axis 0 across arrays of all ranks at once.
        (
            [5],
            None,
            {
                "data/camera0_rgb": (5, 2, 2, 3),
                "data/robot0_eef_pos": (5, 3),
                "meta/episode_ends": (5,),
            },
        ),
        # Length-2 prefix overrides axes 0+1; needs every selected array to
        # have rank >= 2, so we filter out ``meta/episode_ends`` (rank 1).
        (
            [5, 1],
            ["data/camera0_rgb", "data/robot0_eef_pos"],
            {
                "data/camera0_rgb": (5, 1, 2, 3),
                "data/robot0_eef_pos": (5, 1),
            },
        ),
        # Per-array overrides may retile only some arrays while others keep
        # their native chunks.
        (
            {
                "data/camera0_rgb": [5],
                "data/robot0_eef_pos": [7],
            },
            None,
            {
                "data/camera0_rgb": (5, 2, 2, 3),
                "data/robot0_eef_pos": (7, 3),
                "meta/episode_ends": (3,),
            },
        ),
    ],
)
def test_chunk_shapes_resolution_across_mixed_rank(
    heterogeneous_zarrv2_store, chunk_shapes, array_paths, expected
):
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(heterogeneous_zarrv2_store),
        chunk_shapes=chunk_shapes,
        array_paths=array_paths,
    )
    assert datasource._array_chunks == expected


# ---------------------------------------------------------------------------
# align_axis_0 (wide-form mode)
# ---------------------------------------------------------------------------


def test_align_axis_0_emits_wide_rows(ray_start_regular_shared, aligned_zarrv2_store):
    """Wide-row schema: ``t_start``, ``t_stop``, one column per selected array."""
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(aligned_zarrv2_store),
        align_axis_0=True,
        chunk_shapes=[4],
    )
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
    assert set(df.columns) == {"t_start", "t_stop", "img", "state", "label"}
    # shape[0]=8, chunk_shapes=[4] -> 2 rows.
    assert len(df) == 2
    # Reconstruct each array by concatenating slices in order.
    img_recon = np.concatenate(list(df["img"]), axis=0)
    assert img_recon.shape == (8, 4, 4, 3)
    state_recon = np.concatenate(list(df["state"]), axis=0)
    assert state_recon.shape == (8, 3)
    label_recon = np.concatenate(list(df["label"]), axis=0)
    assert label_recon.shape == (8,)
    # t_start/t_stop are correct.
    starts = sorted(df["t_start"].tolist())
    stops = sorted(df["t_stop"].tolist())
    assert starts == [0, 4]
    assert stops == [4, 8]


def test_align_axis_0_column_set(ray_start_regular_shared, aligned_zarrv2_store):
    """array_paths selects which arrays are read; aligned mode emits one column
    per selected array (plus t_start/t_stop)."""
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(aligned_zarrv2_store),
        array_paths=["img", "state"],
        align_axis_0=True,
        chunk_shapes=[4],
    )
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
    assert set(df.columns) == {"t_start", "t_stop", "img", "state"}


def test_align_axis_0_rejects_misaligned_shape0(heterogeneous_zarrv2_store):
    """Misalignment raises with the per-array shape[0] breakdown."""
    with pytest.raises(
        ValueError,
        match=r"All selected arrays must share shape\[0\]",
    ):
        zarrv2_datasource.ZarrV2Datasource(
            str(heterogeneous_zarrv2_store),
            align_axis_0=True,
            chunk_shapes=[5],
        )


def test_align_axis_0_rejects_divergent_axis_0_chunks(aligned_zarrv2_store):
    """If aligned arrays end up with different axis-0 chunks, error clearly.

    The native chunks differ (img=2, state=4, label=8) — without a
    ``chunk_shapes`` re-tile they all stay at native, and the validator
    catches the mismatch.
    """
    with pytest.raises(
        ValueError, match="Aligned arrays must share the same axis-0 chunk size"
    ):
        zarrv2_datasource.ZarrV2Datasource(
            str(aligned_zarrv2_store),
            align_axis_0=True,
        )


# ---------------------------------------------------------------------------
# overlap (aligned-mode lookahead)
# ---------------------------------------------------------------------------


def test_overlap_extends_chunk_data(ray_start_regular_shared, aligned_zarrv2_store):
    """``overlap=N`` makes each row's per-array slice cover ``N`` extra timesteps.

    Aligned store has shape[0]=8, ``chunk_shapes=[4]`` -> rows own [0,4) and [4,8).
    With ``overlap=2``, row 0's data covers [0,6) and row 1's data covers [4,8)
    (clipped at the store end since 4+4+2 > 8).
    """
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(aligned_zarrv2_store),
        align_axis_0=True,
        chunk_shapes=[4],
        overlap=2,
    )
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
    # Ownership unchanged: 2 rows of width 4 each.
    assert sorted(zip(df["t_start"], df["t_stop"])) == [(0, 4), (4, 8)]
    # Data extents: row 0 has 6 timesteps, row 1 has 4 (clipped at shape[0]=8).
    rows = sorted(df.to_dict("records"), key=lambda r: r["t_start"])
    assert rows[0]["img"].shape[0] == 6  # 4 owned + 2 overlap
    assert rows[0]["state"].shape[0] == 6
    assert rows[1]["img"].shape[0] == 4  # 4 owned + 0 overlap (clipped)
    assert rows[1]["state"].shape[0] == 4


def test_overlap_requires_align_axis_0(aligned_zarrv2_store):
    """``overlap`` in long-form (no ``align_axis_0``) is a clear error."""
    with pytest.raises(ValueError, match="overlap requires align_axis_0=True"):
        zarrv2_datasource.ZarrV2Datasource(
            str(aligned_zarrv2_store),
            overlap=2,
        )


def test_overlap_rejects_negative_and_non_int(aligned_zarrv2_store):
    bad_values: list[Any] = [-1, 1.5, "two"]

    for bad in bad_values:
        with pytest.raises(ValueError, match="overlap must be a non-negative integer"):
            zarrv2_datasource.ZarrV2Datasource(
                str(aligned_zarrv2_store),
                align_axis_0=True,
                chunk_shapes=[4],
                overlap=bad,
            )


def test_chunk_shapes_rejected_when_longer_than_smallest_array(
    heterogeneous_zarrv2_store,
):
    """A shared ``chunk_shapes`` override longer than a target rank is an error."""
    with pytest.raises(
        ValueError,
        match=r"chunk_shapes override for array .* has 2 axes but array of shape .* has rank 1",
    ):
        zarrv2_datasource.ZarrV2Datasource(
            str(heterogeneous_zarrv2_store),
            chunk_shapes=[2, 2],  # OK for 2-D and 4-D, fails for 1-D episode_ends
        )


# ---------------------------------------------------------------------------
# Filesystem handling
# ---------------------------------------------------------------------------


def test_accepts_pyarrow_fs_filesystem(zarrv2_group_store):
    """A pyarrow.fs.FileSystem passed in is wrapped into fsspec internally."""
    datasource = zarrv2_datasource.ZarrV2Datasource(
        str(zarrv2_group_store),
        filesystem=pyarrow.fs.LocalFileSystem(),
    )
    from fsspec.spec import AbstractFileSystem

    assert isinstance(datasource._fs, AbstractFileSystem)
    assert set(datasource._metadata_by_path) == {"images", "nested"}


def test_rejects_unsupported_filesystem_type():
    """Filesystem that's neither pyarrow.fs nor fsspec raises ``TypeError``."""
    with pytest.raises(
        TypeError,
        match=r"filesystem must be pyarrow\.fs\.FileSystem or",
    ):
        zarrv2_datasource.ZarrV2Datasource(
            "/tmp/some.zarr",
            filesystem="not-a-filesystem",
        )


# ---------------------------------------------------------------------------
# .zarr.zip URL support
# ---------------------------------------------------------------------------


def test_reads_zarr_zip_local_path(ray_start_regular_shared, zarr_zip_store):
    """A local ``.zarr.zip`` path auto-wires fsspec's ZipFileSystem."""
    datasource = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store))
    # The store has one array "data" of shape (6, 2) chunks (3, 2) -> 2 chunks.
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
    assert len(df) == 2
    assert set(df["array"]) == {"data"}
    recon = _reconstruct_array(df, "data")
    np.testing.assert_array_equal(recon, np.arange(12, dtype="<i4").reshape(6, 2))


# ---------------------------------------------------------------------------
# Read task generation and execution (end-to-end)
# ---------------------------------------------------------------------------


def test_get_read_tasks_batches_chunks_by_parallelism(tmp_path):
    """5 chunks split across parallelism=3 produces batches [2, 2, 1]."""
    store_path = tmp_path / "store.zarr"
    _write_real_zarr_store(
        store_path,
        {"images": (np.arange(5 * 4, dtype="<i4").reshape(5, 4), (1, 4))},
    )
    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))

    read_tasks = datasource.get_read_tasks(parallelism=3)

    assert len(read_tasks) == 3
    assert [task.metadata.num_rows for task in read_tasks] == [2, 2, 1]
    assert all(task.metadata.input_files == (str(store_path),) for task in read_tasks)


def test_long_form_chunk_index_order_matches_grid(ray_start_regular_shared, tmp_path):
    """Lazy grid-range tasks emit chunk_index in the same row-major order as a
    full grid enumeration (regression guard for the lazy-unravel refactor)."""
    from itertools import product

    store_path = tmp_path / "order.zarr"
    # shape (6, 4), chunks (2, 2) -> grid (3, 2) = 6 chunks.
    _write_real_zarr_store(
        store_path, {"a": (np.arange(6 * 4, dtype="<i4").reshape(6, 4), (2, 2))}
    )
    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
    # parallelism=2 -> two flat-index ranges; concatenated they must be in order.
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
    got = [tuple(int(x) for x in ci) for ci in df["chunk_index"]]
    assert got == list(product(range(3), range(2)))


def test_per_task_row_limit_caps_chunks_read(
    ray_start_regular_shared, tmp_path, monkeypatch
):
    """per_task_row_limit bounds how many chunks a task actually reads, so a
    downstream ``limit`` doesn't pull the whole batch's I/O."""
    store_path = tmp_path / "limit.zarr"
    _write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (1,))})
    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))

    reads = []
    real_read_chunk = zarrv2_datasource._read_chunk

    def _spy(*args, **kwargs):
        reads.append(1)
        return real_read_chunk(*args, **kwargs)

    monkeypatch.setattr(zarrv2_datasource, "_read_chunk", _spy)

    # parallelism=1 -> one task batching all 10 chunks; cap it at 3.
    tasks = datasource.get_read_tasks(parallelism=1, per_task_row_limit=3)
    blocks = [block for task in tasks for block in task()]

    total_rows = sum(BlockAccessor.for_block(b).num_rows() for b in blocks)
    assert total_rows == 3
    # The fix: only 3 chunks were actually read (not all 10, then truncated).
    assert len(reads) == 3


def test_read_chunk_retries_transient_io(monkeypatch):
    """_read_chunk retries reads whose error matches retry_match (Ray Data's
    DataContext.retried_io_errors), then succeeds."""
    monkeypatch.setattr("time.sleep", lambda *_: None)  # no backoff in the test

    class _FlakyArray:
        attempts = 0

        def __getitem__(self, _idx):
            type(self).attempts += 1
            if self.attempts < 3:
                raise OSError("Connection reset by peer")
            return np.arange(4, dtype="<i4")

    class _Root:
        def __getitem__(self, _name):
            return _FlakyArray()

    out = zarrv2_datasource._read_chunk(
        _Root(),  # pyrefly: ignore[bad-argument-type]
        "x",
        ((0, 4),),
        retry_match=["Connection reset"],
    )
    np.testing.assert_array_equal(out, np.arange(4, dtype="<i4"))
    assert _FlakyArray.attempts == 3  # failed twice, then succeeded


def test_long_form_schema_and_materialization(ray_start_regular_shared, tmp_path):
    """End-to-end: long-form rows are emitted with the expected columns and data."""
    store_path = tmp_path / "aligned.zarr"
    images_src = np.arange(20, dtype="<i4").reshape(5, 4)
    labels_src = np.arange(5, dtype="|u1")
    _write_real_zarr_store(
        store_path,
        {
            "images": (images_src, (2, 4)),
            "labels": (labels_src, (2,)),
        },
    )

    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))

    # Schema is the long-form quad.
    assert list(df.columns) == ["array", "chunk_index", "chunk_slices", "chunk"]
    # 3 chunks for images (5/2), 3 chunks for labels (5/2) = 6 rows total.
    assert len(df) == 6
    assert set(df["array"]) == {"images", "labels"}

    np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
    np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)

    # ``chunk_slices`` matches the actual chunk shape and indexes back to
    # the source array: arr[start:stop, ...] equals the chunk.
    for _, row in df.iterrows():
        slices = row["chunk_slices"]
        chunk = row["chunk"]
        assert len(slices) == chunk.ndim
        for axis, (start, stop) in enumerate(slices):
            assert stop - start == chunk.shape[axis]
        if row["array"] == "images":
            np.testing.assert_array_equal(
                chunk,
                images_src[slices[0][0] : slices[0][1], slices[1][0] : slices[1][1]],
            )


def test_chunk_shapes_override_changes_grid(ray_start_regular_shared, tmp_path):
    """User-supplied chunk_shapes controls the chunk grid and row count."""
    store_path = tmp_path / "tile.zarr"
    src = np.arange(10, dtype="<i4")
    _write_real_zarr_store(store_path, {"data": (src, (2,))})  # native: 5 chunks

    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path), chunk_shapes=[5])
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
    assert sorted(chunk.shape[0] for chunk in df["chunk"]) == [5, 5]


def test_heterogeneous_store_emits_one_row_per_chunk(
    ray_start_regular_shared, heterogeneous_zarrv2_store
):
    """Mixed-rank/shape/dtype arrays each contribute their chunk count to the output."""
    datasource = zarrv2_datasource.ZarrV2Datasource(str(heterogeneous_zarrv2_store))
    df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))

    # Expected chunk counts:
    #   data/camera0_rgb       shape=(20,2,2,3) chunks=(1,2,2,3) → 20 chunks
    #   data/robot0_eef_pos    shape=(20,3)     chunks=(10,3)    → 2 chunks
    #   meta/episode_ends      shape=(3,)       chunks=(3,)      → 1 chunk
    counts = df.groupby("array").size().to_dict()
    assert counts == {
        "data/camera0_rgb": 20,
        "data/robot0_eef_pos": 2,
        "meta/episode_ends": 1,
    }


# ---------------------------------------------------------------------------
# Estimator
# ---------------------------------------------------------------------------


def test_estimate_inmemory_data_size(tmp_path):
    """Estimate = sum over arrays of numel * dtype.itemsize."""
    store_path = tmp_path / "est.zarr"
    _write_real_zarr_store(
        store_path,
        {
            "a": (np.zeros((5, 4), dtype="<i4"), (2, 4)),
            "b": (np.zeros(5, dtype="|u1"), (2,)),
        },
    )
    datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
    # 5*4*4 (a) + 5*1 (b) = 80 + 5 = 85
    assert datasource.estimate_inmemory_data_size() == 85


# ---------------------------------------------------------------------------
# Cross-filesystem end-to-end (Ray Data convention)
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "fs",
    [
        None,
        lazy_fixture("local_fs"),  # pyarrow.fs (gets wrapped to fsspec internally)
        lazy_fixture("local_fsspec_fs"),  # native fsspec
    ],
)
def test_read_zarr_basic_across_filesystems(ray_start_regular_shared, fs, local_path):
    """Round-trip a real Zarr store through read_zarr for each filesystem flavor.

    Mirrors the parametrized read-path coverage other Ray Data datasources use
    (lance, parquet, json, hudi, …) — exercises None / pyarrow.fs / fsspec
    input shapes against the same store written to a local path.
    """
    store_path = os.path.join(local_path, "data.zarr")
    images_src = np.arange(20, dtype="<i4").reshape(5, 4)
    labels_src = np.arange(5, dtype="|u1")
    _write_real_zarr_store(
        Path(store_path),
        {
            "images": (images_src, (2, 4)),
            "labels": (labels_src, (2,)),
        },
    )

    ds = ray.data.read_zarr(store_path, filesystem=fs)

    # 3 chunks each for images and labels (5/2 → ceil = 3) → 6 rows total.
    assert ds.count() == 6
    df = pd.DataFrame(ds.take_all())
    np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
    np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)


def test_rejects_zarr_v3(tmp_path, monkeypatch):
    """read_zarr targets zarr-python 2.x; an incompatible v3 install must raise a
    clear, actionable error at construction, not a cryptic ImportError mid-read."""
    monkeypatch.setattr(zarr, "__version__", "3.0.1")
    with pytest.raises(ImportError, match=r"zarr-python 2\.x"):
        zarrv2_datasource.ZarrV2Datasource(str(tmp_path))


def test_explicit_filesystem_strips_uri_scheme(ray_start_regular_shared, tmp_path):
    """An explicit ``filesystem=`` plus a scheme-prefixed path must strip the
    scheme so the store path is backend-relative. Regression: pyarrow
    filesystems can't resolve a ``file://`` / ``gs://`` prefix in the path."""
    store_path = tmp_path / "scheme.zarr"
    _write_real_zarr_store(store_path, {"data": (np.arange(6, dtype="<i4"), (2,))})

    ds = zarrv2_datasource.ZarrV2Datasource(
        f"file://{store_path}", filesystem=pyarrow.fs.LocalFileSystem()
    )
    assert ds._store_path == str(store_path)
    df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
    assert len(df) == 3


def test_get_read_tasks_parallelism_zero(tmp_path):
    """parallelism=0 must not divide by zero; fall back to a single task."""
    store_path = tmp_path / "p0.zarr"
    _write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (2,))})
    ds = zarrv2_datasource.ZarrV2Datasource(str(store_path))
    tasks = ds.get_read_tasks(parallelism=0)
    assert len(tasks) >= 1


def test_align_axis_0_rejects_scalar_array(tmp_path):
    """align_axis_0=True with a 0-D (scalar) array must raise a clear error
    rather than an IndexError when reading the (empty) axis-0 chunk size."""
    store_path = tmp_path / "scalar.zarr"
    root = zarr.open_group(str(store_path), mode="w")
    root.create_dataset("vec", data=np.arange(8, dtype="<i4"), chunks=(4,))
    root.create_dataset("scalar", data=np.array(42, dtype="<i4"))  # 0-D
    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))

    with pytest.raises(ValueError, match=r"0-D \(scalar\)"):
        zarrv2_datasource.ZarrV2Datasource(str(store_path), align_axis_0=True)


def test_reads_zarr_zip_with_explicit_zip_filesystem(
    ray_start_regular_shared, zarr_zip_store
):
    """A .zip path read through an explicitly-passed fsspec ZipFileSystem must
    resolve the store at the archive root (store path ``""``), not treat the
    ``.zip`` name as an entry inside the archive."""
    zip_fs = fsspec.filesystem("zip", fo=str(zarr_zip_store))
    ds = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store), filesystem=zip_fs)
    assert ds._store_path == ""
    df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
    assert len(df) == 2


def test_align_axis_0_columns_unify_across_blocks(
    ray_start_regular_shared, aligned_zarrv2_store
):
    """Wide-form gives each array its own column, so blocks combine cleanly
    across the dataset even with trailing edge chunks of differing shape -- the
    batch-safe schema for row-aligned arrays."""
    from ray.data._internal.arrow_ops.transform_pyarrow import unify_schemas
    from ray.data.block import BlockAccessor

    ds = zarrv2_datasource.ZarrV2Datasource(
        str(aligned_zarrv2_store), align_axis_0=True, chunk_shapes=[3]
    )
    blocks = [block for task in ds.get_read_tasks(parallelism=64) for block in task()]
    assert len(blocks) > 1  # actually exercise cross-block unification
    schemas = [BlockAccessor.for_block(b).to_arrow().schema for b in blocks]
    unified = unify_schemas(schemas)  # must not raise
    assert {"t_start", "t_stop", "img", "state", "label"}.issubset(set(unified.names))


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))


# ---------------------------------------------------------------------------
# Custom codec registration in Ray workers
# ---------------------------------------------------------------------------


@pytest.fixture
def fresh_ray():
    """A clean Ray for a test that needs its own ``ray.init`` (e.g. a custom
    ``runtime_env``). Unlike ``shutdown_only`` (teardown only), it also shuts
    down any pre-existing cluster, so isolation doesn't depend on test order.
    """
    if ray.is_initialized():
        ray.shutdown()
    yield
    if ray.is_initialized():
        ray.shutdown()


def test_custom_codec_succeeds_with_worker_setup_hook(fresh_ray, tmp_path):
    """Test that we successfully register a custom codec.

    numcodecs' registry is process-local.
    """
    import numcodecs

    def _register_codec():
        import numcodecs
        import numpy as np

        class _RayZarrTestCodec(numcodecs.abc.Codec):
            codec_id = "ray_zarr_test_codec"

            def encode(self, buf):
                return bytes(buf)

            def decode(self, buf, out=None):
                arr = np.frombuffer(buf, dtype=np.uint8)
                if out is not None:
                    out[:] = arr.view(out.dtype)
                    return out
                return arr.copy()

        numcodecs.register_codec(_RayZarrTestCodec)

    # Register driver-side so we can write the store.
    _register_codec()

    store_path = tmp_path / "codec_test.zarr"
    arr = zarr.open(
        str(store_path),
        mode="w",
        shape=(8,),
        chunks=(4,),
        dtype="u1",
        compressor=numcodecs.get_codec({"id": "ray_zarr_test_codec"}),
    )
    arr[:] = np.arange(8, dtype="u1")
    zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))

    ray.init(
        num_cpus=1,
        logging_level=logging.ERROR,
        log_to_driver=False,
        runtime_env={"worker_process_setup_hook": _register_codec},
    )
    ds = ray.data.read_zarr(str(store_path))
    rows = sorted(ds.take_all(), key=lambda r: tuple(r["chunk_index"]))
    recon = np.concatenate([r["chunk"] for r in rows])
    np.testing.assert_array_equal(recon, np.arange(8, dtype="u1"))