"""Utility functions for zarr array copying and rechunking."""

from __future__ import annotations

import itertools
from collections.abc import Generator

import zarr
from zarr.codecs import BloscCodec


def tuple_max(*tuples: tuple[int, ...]) -> tuple[int, ...]:
    """Return element-wise maximum across multiple same-length tuples.

    Example:
        tuple_max((1, 8, 4), (2, 4, 8)) -> (2, 8, 8)
    """
    return tuple(max(values) for values in zip(*tuples))


def iter_slices(
    bounds: tuple[slice, ...],
    chunk_shape: tuple[int, ...],
) -> Generator[tuple[slice, ...], None, None]:
    """Yield tuples of slices that tile *bounds* in *chunk_shape* increments.

    Parameters
    ----------
    bounds:
        A tuple of slices describing the region of the array to iterate over.
        Each slice must have integer start/stop values (step is ignored).
    chunk_shape:
        Chunk sizes along each dimension.  Must have the same length as *bounds*.

    Yields
    ------
    tuple[slice, ...]
        Successive chunk-aligned slices that together cover *bounds*.
    """
    ranges = []
    for sl, chunk in zip(bounds, chunk_shape):
        starts = range(sl.start, sl.stop, chunk)
        ranges.append([slice(s, min(s + chunk, sl.stop)) for s in starts])
    yield from itertools.product(*ranges)


def gen_array(
    group: zarr.Group,
    template: zarr.Array,
    chunks: tuple[int, ...],
    name: str = "tmp",
    compressors: object = None,
) -> zarr.Array:
    """Create a new zarr array in *group* shaped like *template* but with *chunks*.

    Parameters
    ----------
    group:
        Zarr group in which to create the array.
    template:
        Source array; its shape, dtype, and dimension names are copied.
    chunks:
        Chunk shape for the new array.
    name:
        Name of the array within *group*.
    compressors:
        Zarr v3 compressor codec(s) to use.  Defaults to Blosc/lz4 level-1
        (fast, moderate compression) when *None*.

    Returns
    -------
    zarr.Array
    """
    if compressors is None:
        compressors = BloscCodec(cname="lz4", clevel=1, shuffle="noshuffle")

    return group.create_array(
        name,
        shape=template.shape,
        chunks=chunks,
        dtype=template.dtype,
        compressors=compressors,
        overwrite=True,
    )
