strategies.py

import math
import sys
from collections.abc import Callable, Mapping
from typing import Any, Literal

import hypothesis.extra.numpy as npst
import hypothesis.strategies as st
import numpy as np
import numpy.typing as npt
from hypothesis import event
from hypothesis.strategies import SearchStrategy

import zarr
from zarr.abc.store import RangeByteRequest, Store
from zarr.codecs.bytes import BytesCodec
from zarr.core.array import Array
from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding
from zarr.core.common import JSON, AccessModeLiteral, ZarrFormat
from zarr.core.dtype import get_data_type_from_native_dtype
from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata
from zarr.core.sync import sync
from zarr.storage import MemoryStore, StoreLike
from zarr.storage._common import _dereference_path
from zarr.storage._utils import normalize_path
from zarr.testing.models import ArrayNode, GroupNode, Node
from zarr.types import AnyArray

TrueOrFalse = Literal[True, False]


def default_fs_case_insensitive() -> bool:
    """Return whether the current platform defaults to a case-insensitive filesystem.

    macOS APFS and Windows NTFS are case-insensitive by default; Linux
    filesystems are typically case-sensitive. Used as the default for the
    ``trees()`` strategy so sibling names won't collide when the tree is
    materialized into a filesystem-backed store.
    """
    return sys.platform in ("darwin", "win32")

# Copied from Xarray
_attr_keys = st.text(st.characters(), min_size=1)
_attr_values = st.recursive(
    st.none() | st.booleans() | st.text(st.characters(), max_size=5),
    lambda children: st.lists(children) | st.dictionaries(_attr_keys, children),
    max_leaves=3,
)


@st.composite
def keys(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str:
    return draw(st.lists(node_names, min_size=1, max_size=max_num_nodes).map("/".join))


@st.composite
def paths(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str:
    return draw(st.just("/") | keys(max_num_nodes=max_num_nodes))


def dtypes() -> st.SearchStrategy[np.dtype[Any]]:
    return (
        npst.boolean_dtypes()
        | npst.integer_dtypes(endianness="=")
        | npst.unsigned_integer_dtypes(endianness="=")
        | npst.floating_dtypes(endianness="=")
        | npst.complex_number_dtypes(endianness="=")
        | npst.byte_string_dtypes(endianness="=")
        | npst.unicode_string_dtypes(endianness="=")
        | npst.datetime64_dtypes(endianness="=")
        | npst.timedelta64_dtypes(endianness="=")
    )


def v3_dtypes() -> st.SearchStrategy[np.dtype[Any]]:
    return dtypes()


def v2_dtypes() -> st.SearchStrategy[np.dtype[Any]]:
    return dtypes()


def safe_unicode_for_dtype(dtype: np.dtype[np.str_]) -> st.SearchStrategy[str]:
    """Generate UTF-8-safe text constrained to max_len of dtype."""
    # account for utf-32 encoding (i.e. 4 bytes/character)
    max_len = max(1, dtype.itemsize // 4)

    return st.text(
        alphabet=st.characters(
            exclude_categories=["Cs"],  # Avoid *technically allowed* surrogates
            min_codepoint=32,
        ),
        min_size=1,
        max_size=max_len,
    )


def clear_store(x: Store) -> Store:
    sync(x.clear())
    return x


# From https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#node-names
# 1. must not be the empty string ("")
# 2. must not include the character "/"
# 3. must not be a string composed only of period characters, e.g. "." or ".."
# 4. must not start with the reserved prefix "__"
zarr_key_chars = st.sampled_from(
    ".-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz"
)
node_names = (
    st.text(zarr_key_chars, min_size=1)
    .filter(lambda t: t not in (".", "..") and not t.startswith("__"))
    .filter(lambda name: name.lower() != "zarr.json")
)
short_node_names = (
    st.text(zarr_key_chars, max_size=3, min_size=1)
    .filter(lambda t: t not in (".", "..") and not t.startswith("__"))
    .filter(lambda name: name.lower() != "zarr.json")
)
array_names = node_names
attrs: st.SearchStrategy[Mapping[str, JSON] | None] = st.none() | st.dictionaries(
    _attr_keys, _attr_values
)
# st.builds will only call a new store constructor for different keyword arguments
# i.e. stores.examples() will always return the same object per Store class.
# So we map a clear to reset the store.
stores = st.builds(MemoryStore, st.just({})).map(clear_store)
compressors = st.sampled_from([None, "default"])
zarr_formats: st.SearchStrategy[ZarrFormat] = st.sampled_from([3, 2])
# We de-prioritize arrays having dim sizes 0, 1, 2
array_shapes = npst.array_shapes(max_dims=4, min_side=3, max_side=5) | npst.array_shapes(
    max_dims=4, min_side=0
)


@st.composite
def dimension_names(draw: st.DrawFn, *, ndim: int | None = None) -> list[None | str] | None:
    simple_text = st.text(zarr_key_chars, min_size=0)
    return draw(st.none() | st.lists(st.none() | simple_text, min_size=ndim, max_size=ndim))  # type: ignore[arg-type]


@st.composite
def array_metadata(
    draw: st.DrawFn,
    *,
    array_shapes: Callable[..., st.SearchStrategy[tuple[int, ...]]] = npst.array_shapes,
    zarr_formats: st.SearchStrategy[ZarrFormat] = zarr_formats,
    attributes: SearchStrategy[Mapping[str, JSON] | None] = attrs,
) -> ArrayV2Metadata | ArrayV3Metadata:
    zarr_format = draw(zarr_formats)
    # separator = draw(st.sampled_from(['/', '\\']))
    shape = draw(array_shapes())
    ndim = len(shape)
    np_dtype = draw(dtypes())
    dtype = get_data_type_from_native_dtype(np_dtype)
    fill_value = draw(npst.from_dtype(np_dtype))
    if zarr_format == 2:
        chunk_shape = draw(array_shapes(min_dims=ndim, max_dims=ndim, min_side=1))
        return ArrayV2Metadata(
            shape=shape,
            chunks=chunk_shape,
            dtype=dtype,
            fill_value=fill_value,
            order=draw(st.sampled_from(["C", "F"])),
            attributes=draw(attributes),  # type: ignore[arg-type]
            dimension_separator=draw(st.sampled_from([".", "/"])),
            filters=None,
            compressor=None,
        )
    else:
        chunk_grid = draw(chunk_grids(shape=shape))
        return ArrayV3Metadata(
            shape=shape,
            data_type=dtype,
            chunk_grid=chunk_grid,
            fill_value=fill_value,
            attributes=draw(attributes),  # type: ignore[arg-type]
            dimension_names=draw(dimension_names(ndim=ndim)),
            chunk_key_encoding=DefaultChunkKeyEncoding(separator="/"),  # FIXME
            codecs=[BytesCodec()],
            storage_transformers=(),
        )


@st.composite
def numpy_arrays(
    draw: st.DrawFn,
    *,
    shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes,
    dtype: np.dtype[Any] | None = None,
) -> npt.NDArray[Any]:
    """
    Generate numpy arrays that can be saved in the provided Zarr format.
    """
    if dtype is None:
        dtype = draw(dtypes())
    if np.issubdtype(dtype, np.str_):
        safe_unicode_strings = safe_unicode_for_dtype(dtype)
        return draw(npst.arrays(dtype=dtype, shape=shapes, elements=safe_unicode_strings))

    return draw(npst.arrays(dtype=dtype, shape=shapes))


@st.composite
def chunk_shapes(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[int, ...]:
    # We want this strategy to shrink towards arrays with smaller number of chunks
    # 1. st.integers() shrinks towards smaller values. So we use that to generate number of chunks
    numchunks = draw(
        st.tuples(
            *[
                st.integers(min_value=0 if size == 0 else 1, max_value=max(size, 1))
                for size in shape
            ]
        )
    )
    # 2. and now generate the chunks tuple
    # Chunk sizes must be >= 1 per spec; for zero-extent dimensions use 1.
    chunks = tuple(
        max(1, size // nchunks) if nchunks > 0 else 1
        for size, nchunks in zip(shape, numchunks, strict=True)
    )

    for c in chunks:
        event("chunk size", c)

    if any((c != 0 and s % c != 0) for s, c in zip(shape, chunks, strict=True)):
        event("smaller last chunk")

    return chunks


@st.composite
def shard_shapes(
    draw: st.DrawFn, *, shape: tuple[int, ...], chunk_shape: tuple[int, ...]
) -> tuple[int, ...]:
    # We want this strategy to shrink towards arrays with smaller number of shards
    # shards must be an integral number of chunks
    assert all(c != 0 for c in chunk_shape)
    numchunks = tuple(s // c for s, c in zip(shape, chunk_shape, strict=True))
    multiples = tuple(draw(st.integers(min_value=1, max_value=nc)) for nc in numchunks)
    return tuple(m * c for m, c in zip(multiples, chunk_shape, strict=True))


@st.composite
def np_array_and_chunks(
    draw: st.DrawFn,
    *,
    arrays: st.SearchStrategy[npt.NDArray[Any]] = numpy_arrays(),  # noqa: B008
) -> tuple[np.ndarray[Any, Any], tuple[int, ...]]:
    """A hypothesis strategy to generate small sized random arrays.

    Returns: a tuple of the array and a suitable random chunking for it.
    """
    array = draw(arrays)
    return (array, draw(chunk_shapes(shape=array.shape)))


@st.composite
def arrays(
    draw: st.DrawFn,
    *,
    shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes,
    compressors: st.SearchStrategy = compressors,
    stores: st.SearchStrategy[StoreLike] = stores,
    paths: st.SearchStrategy[str] = paths(),  # noqa: B008
    array_names: st.SearchStrategy = array_names,
    arrays: st.SearchStrategy | None = None,
    attrs: st.SearchStrategy = attrs,
    zarr_formats: st.SearchStrategy = zarr_formats,
    open_mode: AccessModeLiteral = "w",
) -> AnyArray:
    store = draw(stores, label="store")
    path = draw(paths, label="array parent")
    name = draw(array_names, label="array name")
    attributes = draw(attrs, label="attributes")
    zarr_format = draw(zarr_formats, label="zarr format")
    if arrays is None:
        arrays = numpy_arrays(shapes=shapes)
    nparray = draw(arrays, label="array data")
    dim_names: None | list[str | None] = None

    # For v3 arrays, optionally use RectilinearChunkGridMetadata
    chunk_grid_meta: RegularChunkGridMetadata | RectilinearChunkGridMetadata | None = None
    shard_shape = None
    if zarr_format == 3:
        chunk_grid_meta = draw(chunk_grids(shape=nparray.shape), label="chunk grid")

        # Sharding is only supported with regular chunk grids, and has complex
        # divisibility constraints that don't play well with hypothesis shrinking.
        # Disabled for now — sharding should be tested separately.

        dim_names = draw(dimension_names(ndim=nparray.ndim), label="dimension names")
    else:
        dim_names = None

    # test that None works too.
    fill_value = draw(st.one_of([st.none(), npst.from_dtype(nparray.dtype)]))
    # compressor = draw(compressors)

    expected_attrs = {} if attributes is None else attributes

    array_path = _dereference_path(path, name)
    root = zarr.open_group(store, mode=open_mode, zarr_format=zarr_format)

    # Convert chunk grid metadata to a form create_array accepts:
    # - RegularChunkGridMetadata -> flat tuple of ints
    # - RectilinearChunkGridMetadata -> nested list of ints (triggers rectilinear path)
    # - v2 -> flat tuple of ints
    chunks_param: tuple[int, ...] | list[list[int]]
    if zarr_format == 3 and chunk_grid_meta is not None:
        if isinstance(chunk_grid_meta, RectilinearChunkGridMetadata):
            chunks_param = [
                list(dim) if isinstance(dim, tuple) else [dim]
                for dim in chunk_grid_meta.chunk_shapes
            ]
        else:
            chunks_param = chunk_grid_meta.chunk_shape
    else:
        chunks_param = draw(chunk_shapes(shape=nparray.shape), label="chunk shape")

    a = root.create_array(
        array_path,
        shape=nparray.shape,
        chunks=chunks_param,
        shards=shard_shape,
        dtype=nparray.dtype,
        attributes=attributes,
        # compressor=compressor,  # FIXME
        fill_value=fill_value,
        dimension_names=dim_names,
    )

    assert isinstance(a, Array)
    if a.metadata.zarr_format == 3:
        assert a.fill_value is not None
    assert a.name is not None
    assert a.path == normalize_path(array_path)
    assert a.name == "/" + a.path
    assert isinstance(root[array_path], Array)
    assert nparray.shape == a.shape

    # Verify chunks — for rectilinear grids, .chunks raises
    if zarr_format == 3:
        if isinstance(a.metadata.chunk_grid, RectilinearChunkGridMetadata):
            assert shard_shape is None
        else:
            assert isinstance(a.metadata.chunk_grid, RegularChunkGridMetadata)
            assert a.metadata.chunk_grid.chunk_shape == a.chunks
            assert shard_shape == a.shards

    assert a.basename == name, (a.basename, name)
    assert dict(a.attrs) == expected_attrs

    a[:] = nparray

    return a


@st.composite
def simple_arrays(
    draw: st.DrawFn,
    *,
    shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes,
) -> Any:
    return draw(
        arrays(
            shapes=shapes,
            paths=paths(max_num_nodes=2),
            array_names=short_node_names,
            attrs=st.none(),
            compressors=st.sampled_from([None, "default"]),
        )
    )


@st.composite
def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[int]]:
    """Generate valid rectilinear chunk shapes for a given array shape.

    Uses two modes per dimension:
    - "expanded": random divider points create arbitrary chunk sizes
    - "rle": uniform chunks with optional remainder, optionally shuffled

    Keeps max chunks per dimension <= 20 to avoid performance issues
    in property tests. With higher dimensions, the total chunk count
    grows multiplicatively.
    """
    chunk_shapes: list[list[int]] = []
    for size in shape:
        assert size > 0
        if size > 1:
            mode = draw(st.sampled_from(["expanded", "rle"]))
            if mode == "expanded":
                event("rectilinear expanded")
                max_chunks = min(size - 1, 20)
                nchunks = draw(st.integers(min_value=1, max_value=max_chunks))
                dividers = sorted(
                    draw(
                        st.lists(
                            st.integers(min_value=1, max_value=size - 1),
                            min_size=nchunks - 1,
                            max_size=nchunks - 1,
                            unique=True,
                        )
                    )
                )
                chunk_shapes.append(
                    [a - b for a, b in zip(dividers + [size], [0] + dividers, strict=False)]
                )
            else:
                # RLE mode: uniform chunks with optional remainder
                max_chunk_size = min(size, 20)
                chunk_size = draw(st.integers(min_value=1, max_value=max_chunk_size))
                n_full = size // chunk_size
                remainder = size % chunk_size
                chunks_list = [chunk_size] * n_full
                if remainder > 0:
                    chunks_list.append(remainder)
                # Optionally shuffle to create non-contiguous duplicate patterns
                if draw(st.booleans()):
                    event("rectilinear rle shuffled")
                    chunks_list = draw(st.permutations(chunks_list))
                else:
                    event("rectilinear rle")
                chunk_shapes.append(list(chunks_list))
        else:
            chunk_shapes.append([1])
    return chunk_shapes


@st.composite
def chunk_grids(
    draw: st.DrawFn, *, shape: tuple[int, ...]
) -> RegularChunkGridMetadata | RectilinearChunkGridMetadata:
    """Generate either a RegularChunkGridMetadata or RectilinearChunkGridMetadata.

    This strategy depends on the global state of the config having rectilinear chunk grids enabled or not.
    This means that it may be a possible source of a hypothesis FlakyStrategy error due dependence
    on global state. However, in practice this seems unlikely to happen.

    This allows property tests to exercise both chunk grid types.
    """
    # RectilinearChunkGridMetadata doesn't support zero-sized dimensions,
    # so use RegularChunkGridMetadata if any dimension is 0
    if any(s == 0 for s in shape):
        event("using RegularChunkGridMetadata (zero-sized dimensions)")
        return RegularChunkGridMetadata(chunk_shape=draw(chunk_shapes(shape=shape)))

    if zarr.config.get("array.rectilinear_chunks") and draw(st.booleans()):
        chunks = draw(rectilinear_chunks(shape=shape))
        event("using RectilinearChunkGridMetadata")
        return RectilinearChunkGridMetadata(chunk_shapes=tuple(tuple(dim) for dim in chunks))
    else:
        event("using RegularChunkGridMetadata")
        return RegularChunkGridMetadata(chunk_shape=draw(chunk_shapes(shape=shape)))


# Rectilinear arrays need min_side >= 1 so every dimension has at least one element
_rectilinear_shapes = npst.array_shapes(max_dims=3, min_side=1, max_side=20)


@st.composite
def rectilinear_arrays(
    draw: st.DrawFn,
    *,
    shapes: st.SearchStrategy[tuple[int, ...]] = _rectilinear_shapes,
) -> Any:
    """Generate a zarr v3 array with rectilinear (variable) chunk grid."""
    shape = draw(shapes)
    chunk_shapes = draw(rectilinear_chunks(shape=shape))

    np_dtype = draw(dtypes())
    nparray = draw(numpy_arrays(shapes=st.just(shape), dtype=np_dtype))
    fill_value = draw(st.one_of([st.none(), npst.from_dtype(np_dtype)]))
    dim_names = draw(dimension_names(ndim=len(shape)))

    store = MemoryStore()
    with zarr.config.set({"array.rectilinear_chunks": True}):
        a = zarr.create_array(
            store=store,
            shape=shape,
            chunks=chunk_shapes,
            dtype=np_dtype,
            fill_value=fill_value,
            dimension_names=dim_names,
        )
        a[:] = nparray

    return a


def is_negative_slice(idx: Any) -> bool:
    return isinstance(idx, slice) and idx.step is not None and idx.step < 0


@st.composite
def end_slices(draw: st.DrawFn, *, shape: tuple[int, ...]) -> Any:
    """
    A strategy that slices ranges that include the last chunk.
    This is intended to stress-test handling of a possibly smaller last chunk.
    """
    slicers = []
    for size in shape:
        start = draw(st.integers(min_value=size // 2, max_value=size - 1))
        length = draw(st.integers(min_value=0, max_value=size - start))
        slicers.append(slice(start, start + length))
    event("drawing end slice")
    return tuple(slicers)


@st.composite
def basic_indices(
    draw: st.DrawFn,
    *,
    shape: tuple[int, ...],
    min_dims: int = 0,
    max_dims: int | None = None,
    allow_newaxis: TrueOrFalse = False,
    allow_ellipsis: TrueOrFalse = True,
) -> Any:
    """Basic indices without unsupported negative slices."""
    strategy = npst.basic_indices(
        shape=shape,
        min_dims=min_dims,
        max_dims=max_dims,
        allow_newaxis=allow_newaxis,
        allow_ellipsis=allow_ellipsis,
    ).filter(
        lambda idxr: (
            not (
                is_negative_slice(idxr)
                or (isinstance(idxr, tuple) and any(is_negative_slice(idx) for idx in idxr))
            )
        )
    )
    if math.prod(shape) >= 3:
        strategy = end_slices(shape=shape) | strategy
    return draw(strategy)


@st.composite
def orthogonal_indices(
    draw: st.DrawFn, *, shape: tuple[int, ...]
) -> tuple[tuple[np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]:
    """
    Strategy that returns
    (1) a tuple of integer arrays used for orthogonal indexing of Zarr arrays.
    (2) a tuple of integer arrays that can be used for equivalent indexing of numpy arrays
    """
    zindexer = []
    npindexer = []
    ndim = len(shape)
    for axis, size in enumerate(shape):
        if size != 0:
            strategy = npst.integer_array_indices(
                shape=(size,), result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1)
            ) | basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False)
        else:
            strategy = basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False)

        val = draw(
            strategy
            # bare ints, slices
            .map(lambda x: (x,) if not isinstance(x, tuple) else x)
            # skip empty tuple
            .filter(bool)
        )
        (idxr,) = val
        if isinstance(idxr, int):
            idxr = np.array([idxr])
        zindexer.append(idxr)
        if isinstance(idxr, slice):
            idxr = np.arange(*idxr.indices(size))
        elif isinstance(idxr, (tuple, int)):
            idxr = np.array(idxr)
        newshape = [1] * ndim
        newshape[axis] = idxr.size
        npindexer.append(idxr.reshape(newshape))

    # casting the output of broadcast_arrays is needed for numpy < 2
    return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer))


def key_ranges(
    keys: SearchStrategy[str] = node_names, max_size: int = sys.maxsize
) -> SearchStrategy[list[tuple[str, RangeByteRequest]]]:
    """
    Function to generate key_ranges strategy for get_partial_values()
    returns list strategy w/ form::

        [(key, (range_start, range_end)),
         (key, (range_start, range_end)),...]
    """

    def make_request(start: int, length: int) -> RangeByteRequest:
        return RangeByteRequest(start, end=min(start + length, max_size))

    byte_ranges = st.builds(
        make_request,
        start=st.integers(min_value=0, max_value=max_size),
        length=st.integers(min_value=0, max_value=max_size),
    )
    key_tuple = st.tuples(keys, byte_ranges)
    return st.lists(key_tuple, min_size=1, max_size=10)


@st.composite
def complex_rectilinear_arrays(
    draw: st.DrawFn,
    *,
    stores: st.SearchStrategy[StoreLike] = stores,
    paths: st.SearchStrategy[str] = paths(),  # noqa: B008
    array_names: st.SearchStrategy = array_names,
    attrs: st.SearchStrategy = attrs,
) -> tuple[npt.NDArray[Any], AnyArray]:
    """Generate a rectilinear array with many small chunks.

    The shape is derived from the chunk edges (5-10 chunks per dim,
    sizes 1-5), exercising higher chunk counts than ``rectilinear_arrays``.
    """
    ndim = draw(st.integers(min_value=1, max_value=3))
    nchunks = draw(st.integers(min_value=5, max_value=10))
    dim_chunks = st.lists(st.integers(min_value=1, max_value=5), min_size=nchunks, max_size=nchunks)
    chunk_shapes = draw(st.lists(dim_chunks, min_size=ndim, max_size=ndim))

    shape = tuple(sum(dim) for dim in chunk_shapes)
    nparray = draw(numpy_arrays(shapes=st.just(shape)))
    dim_names = draw(dimension_names(ndim=ndim))
    fill_value = draw(st.one_of([st.none(), npst.from_dtype(nparray.dtype)]))
    attributes = draw(attrs)

    store = draw(stores, label="store")
    path = draw(paths, label="array parent")
    name = draw(array_names, label="array name")
    array_path = _dereference_path(path, name)

    root = zarr.open_group(store, mode="w", zarr_format=3)
    with zarr.config.set({"array.rectilinear_chunks": True}):
        a = root.create_array(
            array_path,
            shape=shape,
            chunks=chunk_shapes,
            dtype=nparray.dtype,
            fill_value=fill_value,
            dimension_names=dim_names,
            attributes=attributes,
        )
    a[:] = nparray
    return nparray, a


@st.composite
def chunk_paths(draw: st.DrawFn, ndim: int, numblocks: tuple[int, ...], subset: bool = True) -> str:
    blockidx = draw(
        st.tuples(*tuple(st.integers(min_value=0, max_value=max(0, b - 1)) for b in numblocks))
    )
    subset_slicer = slice(draw(st.integers(min_value=0, max_value=ndim))) if subset else slice(None)
    return "/".join(map(str, blockidx[subset_slicer]))


# ---------------------------------------------------------------------------
# Name strategies — pool-based derivation for prefix collisions
# ---------------------------------------------------------------------------

# Short affix drawn from the zarr key alphabet for prefix/suffix collisions.
affix = st.text(zarr_key_chars, min_size=1, max_size=3)
separators = st.sampled_from(["_", "-", "."])


def similar_name(
    non_sibling_names: set[str],
    sibling_names: set[str],
    *,
    case_insensitive: bool | None = None,
) -> st.SearchStrategy[str]:
    """Strategy that picks a name similar to existing ones, for prefix collisions.

    Either an affixed variant of a sibling name (e.g. ``"foo"`` → ``"foo-bar"``)
    or an exact copy of a non-sibling (e.g. cousin) name.

    Parameters
    ----------
    non_sibling_names : set[str]
        Names from elsewhere in the tree (cousins, ancestors, etc.).
        These may be reused exactly as the generated name.
    sibling_names : set[str]
        Names of nodes at the same level as the one being generated.
        These are used as bases for affixed variants (prefix/suffix collisions).
    case_insensitive : bool | None
        If ``True``, produced names will not differ from ``sibling_names``
        only in letter case. Required when the target store backs onto a
        case-insensitive filesystem (macOS APFS, Windows NTFS default),
        where ``foo`` and ``FOO`` resolve to the same path. Defaults to the
        current platform's filesystem behavior when ``None``.

    Examples
    --------
    Given a tree like::

        /
        ├── alpha/
        │   ├── x
        │   └── y
        └── beta/
            ├── z
            └── ?   ← generating a new name here

    ``sibling_names = {"z"}``, ``non_sibling_names = {"alpha", "x", "y", "beta"}``.
    The strategy might produce ``"z_0"`` (affixed sibling) or ``"x"`` (reused cousin).
    or ``beta`` or a new random name entirely.
    """
    if case_insensitive is None:
        case_insensitive = default_fs_case_insensitive()
    siblings = sorted(sibling_names)
    non_siblings = sorted(non_sibling_names - sibling_names)

    strategies = []
    if bool(siblings):
        # if there are any named siblings we can affix a sibling name
        # choosing to not affix all names in the tree (e.g. cousin names) because
        # that doesn't seem likely to bring a bug, and would expand the search space.
        strategies.append(
            st.sampled_from(siblings).flatmap(
                lambda base: st.one_of(
                    separators.flatmap(lambda sep: affix.map(lambda afx: base + sep + afx)),
                    separators.flatmap(lambda sep: affix.map(lambda afx: afx + sep + base)),
                )
            )
        )
    if bool(non_siblings):
        strategies.append(st.sampled_from(non_siblings))
    key = str.casefold if case_insensitive else (lambda n: n)
    forbidden = {key(n) for n in sibling_names}
    return st.one_of(*strategies).filter(lambda name: key(name) not in forbidden)


@st.composite
def unique_sibling_names(
    draw: st.DrawFn,
    existing_names: set[str],
    num_names: int,
    existing_siblings: set[str] | None = None,
    *,
    case_insensitive: bool | None = None,
) -> list[str]:
    """Draw *num_names* unique names, biased toward collisions with existing ones.

    Parameters
    ----------
    existing_names : set[str]
        All names already present in the tree. Used to generate
        similar-looking candidates (affixed siblings, reused cousins).
    num_names : int
        Number of unique names to generate.
    existing_siblings : set[str] | None
        Names already present at the destination that must not be reused.
        Used by valid_moves to avoid collisions with existing children.

    Returns
    -------
    list[str]
        The generated names, unique among themselves and not in existing_siblings.
    """
    if case_insensitive is None:
        case_insensitive = default_fs_case_insensitive()
    generated_names: set[str] = set()
    already_taken = existing_siblings or set()
    key = str.casefold if case_insensitive else (lambda n: n)

    for _ in range(num_names):
        excluded = generated_names | already_taken
        forbidden = {key(n) for n in excluded}
        # Filter the whole strategy — similar_name can produce collisions.
        generated_names.add(
            draw(
                (
                    st.one_of(
                        node_names,
                        similar_name(
                            existing_names, excluded, case_insensitive=case_insensitive
                        ),
                    )
                    if bool(existing_names) | bool(generated_names)
                    else node_names
                ).filter(
                    lambda name_, f=forbidden, k=key: k(name_) not in f  # type: ignore[misc]
                )
            )
        )
    return list(generated_names)


# ---------------------------------------------------------------------------
# Tree skeleton + naming
# ---------------------------------------------------------------------------


def skeletons(*, max_leaves: int = 50, max_children: int = 4) -> st.SearchStrategy[GroupNode]:
    """Unnamed tree skeletons via st.recursive.

    Always returns a GroupNode (the root group). Child names are placeholder
    indices ("0", "1", ...); real names are assigned later by ``trees``.
    """
    leaves = st.just(ArrayNode(shape=(1,), dtype=np.dtype("i4")))

    def extend(children: st.SearchStrategy[Node]) -> st.SearchStrategy[GroupNode]:
        return st.lists(children, min_size=1, max_size=max_children).map(
            lambda child_list: GroupNode(
                children={str(i): child for i, child in enumerate(child_list)}
            )
        )

    # Wrap in extend so the top level is always a GroupNode (the root group).
    return extend(st.recursive(leaves, extend, max_leaves=max_leaves))


@st.composite
def trees(
    draw: st.DrawFn,
    *,
    max_leaves: st.SearchStrategy[int] = st.integers(min_value=5, max_value=50),  # noqa: B008
    max_children: st.SearchStrategy[int] = st.integers(min_value=1, max_value=4),  # noqa: B008
    case_insensitive: bool | None = None,
) -> GroupNode:
    """Strategy producing a GroupNode tree descriptor.

    Uses st.recursive for the tree structure (good structural shrinking)
    and @composite for name assignment (pool-based prefix collisions).

    Parameters
    ----------
    case_insensitive : bool | None
        If ``True``, sibling names will not differ only in letter case, so
        the tree can be materialized into a store backed by a case-insensitive
        filesystem (macOS APFS, Windows NTFS default), where ``foo`` and
        ``FOO`` resolve to the same path. Defaults to the current platform's
        filesystem behavior when ``None``.

    Examples
    --------
    >>> @given(tree=trees())
    ... def test_something(tree):
    ...     reference = tree.materialize(zarr.storage.MemoryStore())
    ...     under_test = tree.materialize(my_icechunk_store())
    ...     # compare...
    """
    if case_insensitive is None:
        case_insensitive = default_fs_case_insensitive()

    def rebuild_with_names(
        group: GroupNode, existing_names: set[str]
    ) -> tuple[GroupNode, set[str]]:
        new_names = draw(
            unique_sibling_names(
                existing_names,
                num_names=len(group.children),
                case_insensitive=case_insensitive,
            )
        )
        existing_names = existing_names | set(new_names)
        children: dict[str, Node] = {}
        for name, child in zip(new_names, group.children.values(), strict=True):
            if isinstance(child, GroupNode):
                child, existing_names = rebuild_with_names(child, existing_names)
            children[name] = child
        return GroupNode(children=children), existing_names

    # Two-step generation: first draw the tree structure (skeletons uses
    # st.recursive which gives good structural shrinking), then assign real
    # names via @composite (which allows pool-based derivation for realistic
    # prefix collisions). Doing both in one step would sacrifice either
    # structural shrinking or name similarity.
    skeleton = draw(skeletons(max_leaves=draw(max_leaves), max_children=draw(max_children)))
    result, _ = rebuild_with_names(skeleton, set())
    return result