# SPDX-FileCopyrightText: Copyright (c) Fideus Labs LLC
# SPDX-License-Identifier: MIT
import sys
import tempfile
from dataclasses import asdict
from pathlib import Path, PurePosixPath
from typing import Optional, Union, Tuple, Dict, List
import warnings

from .methods._metadata import get_method_metadata

if sys.version_info < (3, 10):
    import importlib_metadata
else:
    import importlib.metadata as importlib_metadata

import dask.array
from dask import __version__ as dask_version
import numpy as np
from itkwasm import array_like_to_numpy_array

import zarr
import zarr.storage
from ._zarr_open_array import open_array
from .v04.zarr_metadata import Metadata as Metadata_v04
from .v05.zarr_metadata import Metadata as Metadata_v05
from .rfc4 import is_rfc4_enabled
from .rfc9_zip import is_ozx_path, write_store_to_zip
from ._zarr_types import StoreLike
from ._zarr_kwargs import zarr_kwargs


from .config import config
from .memory_usage import memory_usage
from .methods._support import _dim_scale_factors
from .multiscales import Multiscales
from .rich_dask_progress import NgffProgress, NgffProgressCallback
from .to_multiscales import to_multiscales
from packaging.version import Version

zarr_version = Version(zarr.__version__)
IS_ZARR_V3_PLUS = zarr_version.major >= 3
DASK_SUPPORTS_SHARDING = Version(dask_version) >= Version("2025.12.0")


def _pop_metadata_optionals(metadata_dict, enabled_rfcs: Optional[List[int]] = None):
    for ax in metadata_dict["axes"]:
        if ax["unit"] is None:
            ax.pop("unit")

        # Handle RFC 4: Remove orientation if RFC 4 is not enabled
        if not is_rfc4_enabled(enabled_rfcs) and "orientation" in ax:
            ax.pop("orientation")

    if metadata_dict["coordinateTransformations"] is None:
        metadata_dict.pop("coordinateTransformations")

    if metadata_dict["omero"] is None:
        metadata_dict.pop("omero")

    return metadata_dict


def _prep_for_to_zarr(store: StoreLike, arr: dask.array.Array) -> dask.array.Array:
    try:
        importlib_metadata.distribution("kvikio")
        _KVIKIO_AVAILABLE = True
    except importlib_metadata.PackageNotFoundError:
        _KVIKIO_AVAILABLE = False

    if _KVIKIO_AVAILABLE:
        from kvikio.zarr import GDSStore

        if not isinstance(store, GDSStore):
            arr = dask.array.map_blocks(
                array_like_to_numpy_array,
                arr,
                dtype=arr.dtype,
                meta=np.empty(()),
            )
        return arr
    return dask.array.map_blocks(
        array_like_to_numpy_array, arr, dtype=arr.dtype, meta=np.empty(())
    )


def _numpy_to_zarr_dtype(dtype):
    dtype_map = {
        "bool": "bool",
        "int8": "int8",
        "int16": "int16",
        "int32": "int32",
        "int64": "int64",
        "uint8": "uint8",
        "uint16": "uint16",
        "uint32": "uint32",
        "uint64": "uint64",
        "float16": "float16",
        "float32": "float32",
        "float64": "float64",
        "complex64": "complex64",
        "complex128": "complex128",
    }

    dtype_str = str(dtype)

    # Handle endianness - strip byte order chars
    if dtype_str.startswith(("<", ">", "|")):
        dtype_str = dtype_str[1:]

    # Look up corresponding zarr dtype
    try:
        return dtype_map[dtype_str]
    except KeyError:
        raise ValueError(f"dtype {dtype} cannot be mapped to Zarr v3 core dtype")


def _write_with_tensorstore(
    store_path: str,
    array,
    region,
    chunks,
    zarr_format,
    dimension_names=None,
    internal_chunk_shape=None,
    full_array_shape=None,
    create_dataset=True,
    compressor=None,
    **kwargs,
) -> None:
    """Write array using tensorstore backend"""
    import tensorstore as ts

    # Use full array shape if provided, otherwise use the region array shape
    dataset_shape = full_array_shape if full_array_shape is not None else array.shape

    # Build the base spec
    spec = {
        "kvstore": {
            "driver": "file",
            "path": store_path,
        },
        "metadata": {
            "shape": dataset_shape,
        },
    }

    if zarr_format == 2:
        spec["driver"] = "zarr" if not IS_ZARR_V3_PLUS else "zarr2"
        spec["metadata"]["dimension_separator"] = "/"
        spec["metadata"]["dtype"] = array.dtype.str
        # Only add chunk info when creating the dataset
        if create_dataset:
            spec["metadata"]["chunks"] = chunks
            # Add compression for zarr v2 with TensorStore
            if compressor is not None:
                # TensorStore zarr2 driver uses compressor in metadata
                if hasattr(compressor, "codec_id"):
                    # numcodecs compressor object
                    spec["metadata"]["compressor"] = compressor.get_config()
                else:
                    # Simple compressor name or config dict
                    spec["metadata"]["compressor"] = compressor
    elif zarr_format == 3:
        spec["driver"] = "zarr3"
        spec["metadata"]["data_type"] = _numpy_to_zarr_dtype(array.dtype)
        spec["metadata"]["chunk_key_encoding"] = {
            "name": "default",
            "configuration": {"separator": "/"},
        }
        if dimension_names:
            spec["metadata"]["dimension_names"] = dimension_names
        # Only add chunk info when creating the dataset
        if create_dataset:
            spec["metadata"]["chunk_grid"] = {
                "name": "regular",
                "configuration": {"chunk_shape": chunks},
            }

            # Build codecs list for zarr v3
            codecs = []

            # Helper function to create compression codec
            def create_compression_codec(compressor):
                if compressor is None:
                    return None

                if hasattr(compressor, "codec_id"):
                    # numcodecs compressor object
                    codec_id = compressor.codec_id
                    if codec_id == "gzip":
                        return {
                            "name": "gzip",
                            "configuration": {"level": getattr(compressor, "level", 6)},
                        }
                    elif codec_id == "blosc":
                        return {
                            "name": "blosc",
                            "configuration": {
                                "cname": getattr(compressor, "cname", "lz4"),
                                "clevel": getattr(compressor, "clevel", 5),
                                "shuffle": "shuffle"
                                if getattr(compressor, "shuffle", 1) == 1
                                else "noshuffle",
                            },
                        }
                    elif codec_id == "zstd":
                        return {
                            "name": "zstd",
                            "configuration": {"level": getattr(compressor, "level", 3)},
                        }
                    elif codec_id == "lz4":
                        return {"name": "lz4"}
                    else:
                        # Fallback: try to use the codec_id as name
                        return {"name": codec_id}
                elif isinstance(compressor, str):
                    # Simple codec name
                    return {"name": compressor}
                elif isinstance(compressor, dict):
                    # Already in codec format
                    return compressor
                return None

            # Add sharding codec with inner codecs if needed
            if internal_chunk_shape:
                sharding_config = {"chunk_shape": internal_chunk_shape}

                # If compression is specified, add it as inner codec for sharding
                if compressor is not None:
                    compression_codec = create_compression_codec(compressor)
                    if compression_codec:
                        # For sharding, compression goes in the inner codecs
                        sharding_config["codecs"] = [compression_codec]

                codecs.append(
                    {
                        "name": "sharding_indexed",
                        "configuration": sharding_config,
                    }
                )
            else:
                # No sharding, add compression codec directly if specified
                if compressor is not None:
                    compression_codec = create_compression_codec(compressor)
                    if compression_codec:
                        codecs.append(compression_codec)

            # Set codecs if any were added
            if codecs:
                spec["metadata"]["codecs"] = codecs
    else:
        raise ValueError(f"Unsupported zarr format: {zarr_format}")

    # Try to open existing dataset first, create only if needed
    try:
        if create_dataset:
            dataset = ts.open(spec, create=True, dtype=array.dtype).result()
        else:
            # For existing datasets, use a minimal spec that just specifies the path
            existing_spec = {
                "kvstore": {
                    "driver": "file",
                    "path": store_path,
                },
                "driver": spec["driver"],
            }
            dataset = ts.open(existing_spec, create=False, dtype=array.dtype).result()
    except Exception as e:
        if "ALREADY_EXISTS" in str(e) and create_dataset:
            # Dataset already exists, open it without creating
            existing_spec = {
                "kvstore": {
                    "driver": "file",
                    "path": store_path,
                },
                "driver": spec["driver"],
            }
            dataset = ts.open(existing_spec, create=False, dtype=array.dtype).result()
        else:
            raise

    # Try to write the dask array directly first
    try:
        dataset[region] = array
    except Exception as e:
        # If we encounter dimension mismatch or shape-related errors,
        # compute the array and try again with corrective action
        error_msg = str(e).lower()
        if any(
            keyword in error_msg
            for keyword in [
                "dimension",
                "shape",
                "mismatch",
                "size",
                "extent",
                "rank",
                "invalid",
            ]
        ):
            # Compute the array to get the actual shape
            computed_array = array.compute()

            # Adjust region to match the actual computed array shape if needed
            if len(region) == len(computed_array.shape):
                adjusted_region = tuple(
                    slice(
                        region[i].start or 0,
                        (region[i].start or 0) + computed_array.shape[i],
                    )
                    if isinstance(region[i], slice)
                    else region[i]
                    for i in range(len(region))
                )
            else:
                adjusted_region = region

            # Try writing the computed array with adjusted region
            dataset[adjusted_region] = computed_array
        else:
            # Re-raise the exception if it's not related to dimension/shape issues
            raise


def _validate_ngff_parameters(
    version: str,
    chunks_per_shard: Optional[Union[int, Tuple[int, ...], Dict[str, int]]],
    use_tensorstore: bool,
    store: StoreLike,
) -> None:
    """Validate the parameters for the NGFF Zarr generation."""
    if version != "0.4" and version != "0.5":
        raise ValueError(f"Unsupported version: {version}")

    if chunks_per_shard is not None:
        if version == "0.4":
            raise ValueError(
                "Sharding is only supported for OME-Zarr version 0.5 and later"
            )
        if not use_tensorstore and not IS_ZARR_V3_PLUS:
            raise ValueError(
                "Sharding requires zarr-python version >= 3.0.0b1 for OME-Zarr version >= 0.5"
            )

    if use_tensorstore and not isinstance(store, (str, Path)):
        raise ValueError("use_tensorstore currently requires a path-like store")


def _prepare_metadata(
    multiscales: Multiscales, version: str, enabled_rfcs: Optional[List[int]] = None
) -> Tuple[Union[Metadata_v04, Metadata_v05], Tuple[str, ...], Dict]:
    """Prepare and convert metadata to the proper version format."""
    metadata = multiscales.metadata

    # Convert method enum to lowercase string for the type field
    method_type = None
    method_metadata = None
    if multiscales.method is not None:
        method_type = multiscales.method.value
        method_metadata = get_method_metadata(multiscales.method)

    metadata = metadata.to_version(version)
    metadata.type = method_type
    metadata.metadata = method_metadata

    dimension_names = metadata.dimension_names
    dimension_names_kwargs = (
        {"dimension_names": dimension_names} if version != "0.4" else {}
    )

    return metadata, dimension_names, dimension_names_kwargs


def _create_zarr_root(
    store: StoreLike,
    chunk_store: Optional[StoreLike],
    version: str,
    overwrite: bool,
    metadata_dict: Dict,
) -> zarr.Group:
    """Create and configure the root Zarr group with proper attributes."""
    zarr_format = 2 if version == "0.4" else 3
    format_kwargs = {"zarr_format": zarr_format} if IS_ZARR_V3_PLUS else {}

    if version == "0.4":
        root = zarr.open_group(
            store,
            mode="w" if overwrite else "a",
            chunk_store=chunk_store,
            **format_kwargs,
        )
    else:
        if not IS_ZARR_V3_PLUS:
            raise ValueError(
                "zarr-python version >= 3.0.0b2 required for OME-Zarr version >= 0.5"
            )
        # For version >= 0.5, open root with Zarr v3
        root = zarr.open_group(
            store,
            mode="w" if overwrite else "a",
            chunk_store=chunk_store,
            **format_kwargs,
        )

    if "omero" in metadata_dict:
        root.attrs["omero"] = metadata_dict.pop("omero")

    if version != "0.4":
        # RFC 2, Zarr 3
        root.attrs["ome"] = {"version": version, "multiscales": [metadata_dict]}
    else:
        root.attrs["multiscales"] = [metadata_dict]

    return root


def _configure_sharding(
    arr: dask.array.Array,
    chunks_per_shard: Optional[Union[int, Tuple[int, ...], Dict[str, int]]],
    dims: Tuple[str, ...],
    kwargs: Dict,
) -> Tuple[Dict, Optional[Tuple[int, ...]], dask.array.Array]:
    """Configure sharding parameters if sharding is enabled."""
    if chunks_per_shard is None:
        return {}, None, arr

    c0 = tuple([c[0] for c in arr.chunks])

    if isinstance(chunks_per_shard, int):
        shards = tuple([c * chunks_per_shard for c in c0])
    elif isinstance(chunks_per_shard, (tuple, list)):
        if len(chunks_per_shard) != arr.ndim:
            raise ValueError(f"chunks_per_shard must be a tuple of length {arr.ndim}")
        shards = tuple([c * c0[i] for i, c in enumerate(chunks_per_shard)])
    elif isinstance(chunks_per_shard, dict):
        shards = {d: c * chunks_per_shard.get(d, 1) for d, c in zip(dims, c0)}
        shards = tuple([shards[d] for d in dims])
    else:
        raise ValueError("chunks_per_shard must be an int, tuple, or dict")

    internal_chunk_shape = c0
    arr = arr.rechunk(shards)

    # Configure sharding parameters differently for v2 vs v3
    sharding_kwargs = {}
    if IS_ZARR_V3_PLUS:
        # For Zarr v3, configure sharding as a codec
        # Use chunk_shape for internal chunks and configure sharding via codecs
        sharding_kwargs["chunk_shape"] = internal_chunk_shape
        # Note: sharding codec will be configured separately in the codecs parameter
        # We'll pass the shard shape through a separate key to be handled later
        sharding_kwargs["_shard_shape"] = shards
    else:
        # For zarr v2, use the older API
        sharding_kwargs = {
            "shards": shards,
            "chunks": internal_chunk_shape,
        }

    return sharding_kwargs, internal_chunk_shape, arr


def _write_array_with_tensorstore(
    store_path: str,
    path: str,
    arr: dask.array.Array,
    chunks: Union[Tuple[int, ...], List[int]],
    shards: Optional[Tuple[int, ...]],
    internal_chunk_shape: Optional[Tuple[int, ...]],
    zarr_format: int,
    dimension_names: Optional[Tuple[str, ...]],
    region: Tuple[slice, ...],
    full_array_shape: Optional[Tuple[int, ...]] = None,
    create_dataset: bool = True,
    **kwargs,
) -> None:
    """Write an array using the TensorStore backend."""
    # Extract compressor and other conflicting parameters from kwargs to avoid conflicts
    compressor = kwargs.pop("compressor", None)
    kwargs.pop("chunks", None)  # Remove chunks from kwargs since it's a positional arg

    scale_path = f"{store_path}/{path}"
    if shards is None:
        _write_with_tensorstore(
            scale_path,
            arr,
            region,
            chunks,
            zarr_format=zarr_format,
            dimension_names=dimension_names,
            full_array_shape=full_array_shape,
            create_dataset=create_dataset,
            compressor=compressor,
            **kwargs,
        )
    else:  # Sharding
        _write_with_tensorstore(
            scale_path,
            arr,
            region,
            chunks,
            zarr_format=zarr_format,
            dimension_names=dimension_names,
            internal_chunk_shape=internal_chunk_shape,
            full_array_shape=full_array_shape,
            create_dataset=create_dataset,
            compressor=compressor,
            **kwargs,
        )


def _prepare_zarr_kwargs(to_zarr_kwargs: Dict):
    """Prepare zarr kwargs for dask.array.to_zarr.

    This helper function ensures that correct kwargs are passed on based on which version of zarr
    and dask is being used. The different versions support different sets of arguments. The zarr_kwargs
    are adjusted in place and thus the original is overwritten. This is not a problem given that the
    arguments being adjusted are the same for the zarr store in use.
    """
    is_zarr_f2 = to_zarr_kwargs.get("zarr_format") == 2

    # The zarr v2 case does not have to be checked here as this is done in `_zarr_kwargs.py`.
    # The reason for not doing it here is that it only has one option whereas zarr v3 depends on zarr format being used.
    if IS_ZARR_V3_PLUS and is_zarr_f2:
        if DASK_SUPPORTS_SHARDING:
            # New dask uses chunk_key_encoding
            to_zarr_kwargs["chunk_key_encoding"] = {"name": "v2", "separator": "/"}
            to_zarr_kwargs.pop("dimension_separator", None)
        else:
            # Old dask uses dimension_separator
            to_zarr_kwargs["dimension_separator"] = "/"
            to_zarr_kwargs.pop("chunk_key_encoding", None)

    # New dask doesn't accept zarr_format in zarr_array_kwargs
    if DASK_SUPPORTS_SHARDING:
        to_zarr_kwargs.pop("zarr_format", None)


def _write_array_direct(
    arr: dask.array.Array,
    store: StoreLike,
    path: str,
    sharding_kwargs: Dict,
    zarr_kwargs: Dict,
    format_kwargs: Dict,
    dimension_names_kwargs: Dict,
    region: Optional[Tuple[slice, ...]] = None,
    zarr_array=None,
    **kwargs,
) -> None:
    """Write an array directly using dask.array.to_zarr."""
    arr = _prep_for_to_zarr(store, arr)

    zarr_fmt = format_kwargs.get("zarr_format")

    # Handle sharding kwargs for direct writing
    cleaned_sharding_kwargs = {}

    if sharding_kwargs and "_shard_shape" in sharding_kwargs:
        # For Zarr v3 direct writes, use shards and chunks parameters
        shard_shape = sharding_kwargs["_shard_shape"]
        internal_chunk_shape = sharding_kwargs.get("chunk_shape")

        # Ensure internal_chunk_shape is available
        if internal_chunk_shape is None:
            # Use chunks from arr if available, or default
            internal_chunk_shape = tuple(arr.chunks[i][0] for i in range(arr.ndim))

        # For direct Zarr v3 writes, use shards and chunks
        cleaned_sharding_kwargs["shards"] = shard_shape
        cleaned_sharding_kwargs["chunks"] = internal_chunk_shape

        # Remove internal kwargs
        cleaned_sharding_kwargs.update(
            {
                k: v
                for k, v in sharding_kwargs.items()
                if k not in ["_shard_shape", "chunk_shape"]
            }
        )
    else:
        cleaned_sharding_kwargs = sharding_kwargs

    to_zarr_kwargs = {
        **cleaned_sharding_kwargs,
        **zarr_kwargs,
        **format_kwargs,
        **dimension_names_kwargs,
        **kwargs,
    }

    if zarr_fmt == 3 and zarr_array is None:
        # Zarr v3, use zarr.create_array and assign (whole array or region)
        array = zarr.create_array(
            store=store,
            name=path,
            shape=arr.shape,
            dtype=arr.dtype,
            **to_zarr_kwargs,
        )
        if region is not None:
            array[region] = arr.compute()
        else:
            array[:] = arr.compute()
    else:
        _prepare_zarr_kwargs(to_zarr_kwargs)

        target = (
            zarr_array if (region is not None and zarr_array is not None) else store
        )
        # TODO update this when dask 2026.2.0 comes out which would allow old **kwargs
        if DASK_SUPPORTS_SHARDING:
            dask.array.to_zarr(
                arr,
                target,
                region=region
                if (region is not None and zarr_array is not None)
                else None,
                component=path,
                overwrite=False,
                compute=True,
                return_stored=False,
                zarr_array_kwargs=to_zarr_kwargs,
            )
        else:
            dask.array.to_zarr(
                arr,
                target,
                region=region
                if (region is not None and zarr_array is not None)
                else None,
                component=path,
                overwrite=False,
                compute=True,
                return_stored=False,
                **to_zarr_kwargs,
            )


def _handle_large_array_writing(
    image,
    arr: dask.array.Array,
    store: StoreLike,
    path: str,
    dims: Tuple[str, ...],
    dim_factors: Dict[str, int],
    chunks: Tuple[int, ...],
    sharding_kwargs: Dict,
    zarr_kwargs: Dict,
    format_kwargs: Dict,
    dimension_names_kwargs: Dict,
    use_tensorstore: bool,
    store_path: Optional[str],
    zarr_format: int,
    dimension_names: Tuple[str, ...],
    internal_chunk_shape: Optional[Tuple[int, ...]],
    shards: Optional[Tuple[int, ...]],
    progress: Optional[Union[NgffProgress, NgffProgressCallback]],
    index: int,
    nscales: int,
    **kwargs,
) -> None:
    """Handle writing large arrays by splitting them into manageable pieces."""
    shrink_factors = []
    for dim in dims:
        if dim in dim_factors:
            shrink_factors.append(dim_factors[dim])
        else:
            shrink_factors.append(1)

    # Ensure chunks are compatible with Dask's to_zarr when writing with regions.
    # The Zarr chunk size must divide evenly into the dimension size to avoid
    # PerformanceWarning and potential data loss during region writes.
    def _find_optimal_chunk_size(first_chunk, dim_size, min_divisor=16):
        """Find a chunk size that divides evenly into dim_size and is ideally divisible by min_divisor.

        The returned chunk size will:
        1. Divide evenly into dim_size (required for safe region writes)
        2. Be as close as possible to first_chunk
        3. Preferably be divisible by min_divisor for performance
        """
        # If dimension is very small, just use it directly
        if dim_size <= min_divisor:
            return dim_size

        # Start with the target chunk size
        target = first_chunk

        # First try to find a divisor of dim_size that's divisible by min_divisor
        # and close to our target
        best_chunk = dim_size  # Fallback: use full dimension
        best_distance = abs(dim_size - target)

        # Check all divisors of dim_size
        for i in range(1, int(np.sqrt(dim_size)) + 1):
            if dim_size % i == 0:
                # i and dim_size//i are both divisors
                for candidate in [i, dim_size // i]:
                    distance = abs(candidate - target)
                    # Prefer divisors that are multiples of min_divisor
                    is_multiple = candidate % min_divisor == 0

                    # Update if closer to target, with preference for multiples of min_divisor
                    if distance < best_distance or (
                        distance == best_distance
                        and is_multiple
                        and best_chunk % min_divisor != 0
                    ):
                        best_chunk = candidate
                        best_distance = distance

        return best_chunk

    # If sharding is enabled, configure it properly
    chunk_kwargs = {}
    codecs_kwargs = {}

    if sharding_kwargs and "_shard_shape" in sharding_kwargs:
        # For Zarr v3 with sharding, we need to ensure the shard shape divides evenly
        shard_shape = sharding_kwargs.pop("_shard_shape")
        internal_chunk_shape = sharding_kwargs.get(
            "chunk_shape"
        )  # This is the inner chunk shape

        # Apply _find_optimal_chunk_size to the shard shape to ensure it divides evenly
        optimized_shard_shape = tuple(
            [
                _find_optimal_chunk_size(s, arr.shape[i])
                for i, s in enumerate(shard_shape)
            ]
        )

        # Ensure internal_chunk_shape divides evenly into optimized_shard_shape
        if internal_chunk_shape is not None:
            # Adjust each internal chunk to be a divisor of the corresponding shard dimension
            adjusted_internal_chunks = []
            for shard_dim, internal_dim in zip(
                optimized_shard_shape, internal_chunk_shape
            ):
                # Find the best divisor of shard_dim that's close to internal_dim
                if shard_dim % internal_dim == 0:
                    # Already divides evenly
                    adjusted_internal_chunks.append(internal_dim)
                else:
                    # Find closest divisor
                    best_divisor = _find_optimal_chunk_size(internal_dim, shard_dim)
                    adjusted_internal_chunks.append(best_divisor)
            internal_chunk_shape = tuple(adjusted_internal_chunks)
        else:
            # No internal chunks specified, use defaults based on array chunks
            internal_chunk_shape = tuple(
                [
                    _find_optimal_chunk_size(c[0], s)
                    for c, s in zip(arr.chunks, optimized_shard_shape)
                ]
            )

        # Configure the sharding codec with proper defaults
        from zarr.codecs.sharding import ShardingCodec
        from zarr.codecs.bytes import BytesCodec
        from zarr.codecs.zstd import ZstdCodec

        # Default inner codecs for sharding
        default_codecs = [BytesCodec(), ZstdCodec()]

        # The array's chunk_shape should be the shard shape
        # The sharding codec's chunk_shape should be the internal chunk shape
        sharding_codec = ShardingCodec(
            chunk_shape=internal_chunk_shape,  # Internal chunk shape within shards
            codecs=default_codecs,
        )

        # Set up codecs with sharding
        existing_codecs = zarr_kwargs.get("codecs", [])
        if not isinstance(existing_codecs, list):
            existing_codecs = []
        codecs_kwargs["codecs"] = [sharding_codec] + existing_codecs

        # Set the array's chunk_shape to the optimized shard shape
        chunk_kwargs["chunk_shape"] = optimized_shard_shape

        # For region computation, use the optimized shard shape (actual zarr chunk shape)
        zarr_chunk_shape = optimized_shard_shape

        # Clean up remaining kwargs (remove chunk_shape since we're setting it explicitly)
        remaining_kwargs = {
            k: v
            for k, v in sharding_kwargs.items()
            if k not in ["_shard_shape", "chunk_shape"]
        }
        sharding_kwargs_clean = remaining_kwargs
    elif sharding_kwargs:
        # For Zarr v2 or other cases with sharding but no _shard_shape
        chunks = tuple(
            [
                _find_optimal_chunk_size(c[0], arr.shape[i])
                for i, c in enumerate(arr.chunks)
            ]
        )
        zarr_chunk_shape = chunks
        sharding_kwargs_clean = sharding_kwargs
    else:
        # No sharding
        chunks = tuple(
            [
                _find_optimal_chunk_size(c[0], arr.shape[i])
                for i, c in enumerate(arr.chunks)
            ]
        )
        chunk_kwargs = {"chunks": chunks}
        zarr_chunk_shape = chunks
        sharding_kwargs_clean = {}

    if format_kwargs["zarr_format"] == 2:
        if IS_ZARR_V3_PLUS:
            zarr_kwargs["dimension_separator"] = zarr_kwargs["chunk_key_encoding"][
                "separator"
            ]
            del zarr_kwargs["chunk_key_encoding"]
        else:
            zarr_kwargs["dimension_separator"] = "/"

    zarr_array = open_array(
        shape=arr.shape,
        dtype=arr.dtype,
        store=store,
        path=path,
        mode="a",
        **chunk_kwargs,
        **sharding_kwargs_clean,
        **zarr_kwargs,
        **codecs_kwargs,
        **dimension_names_kwargs,
        **format_kwargs,
    )

    shape = image.data.shape
    x_index = dims.index("x")
    y_index = dims.index("y")

    regions = _compute_write_regions(
        image, dims, arr, shape, x_index, y_index, zarr_chunk_shape, shrink_factors
    )

    for region_index, region in enumerate(regions):
        if isinstance(progress, NgffProgressCallback):
            progress.add_callback_task(
                f"[green]Writing scale {index + 1} of {nscales}, region {region_index + 1} of {len(regions)}"
            )

        arr_region = arr[region]
        arr_region = _prep_for_to_zarr(store, arr_region)
        optimized = dask.array.Array(
            dask.array.optimize(
                arr_region.__dask_graph__(), arr_region.__dask_keys__()
            ),
            arr_region.name,
            arr_region.chunks,
            meta=arr_region,
        )

        if use_tensorstore:
            _write_array_with_tensorstore(
                store_path,
                path,
                optimized,
                chunks,  # Use original array chunks, not region chunks
                shards,
                internal_chunk_shape,
                zarr_format,
                dimension_names,
                region,
                full_array_shape=arr.shape,
                create_dataset=(region_index == 0),  # Only create on first region
                **kwargs,
            )
        else:
            _write_array_direct(
                optimized,
                store,
                path,
                sharding_kwargs,
                zarr_kwargs,
                format_kwargs,
                dimension_names_kwargs,
                region,
                zarr_array,
                **kwargs,
            )


def _compute_write_regions(
    image,
    dims: Tuple[str, ...],
    arr: dask.array.Array,
    shape: Tuple[int, ...],
    x_index: int,
    y_index: int,
    chunks: Tuple[int, ...],
    shrink_factors: List[int],
) -> List[Tuple[slice, ...]]:
    """Compute the regions for writing a large array in chunks."""
    regions = []

    # If z dimension exists, handle 3D data
    if "z" in dims:
        z_index = dims.index("z")
        slice_bytes = memory_usage(image, {"z"})
        slab_slices = min(
            int(np.ceil(config.memory_target / slice_bytes)), arr.shape[z_index]
        )
        z_chunks = chunks[z_index]
        slice_planes = False
        if slab_slices < z_chunks:
            slab_slices = z_chunks
            slice_planes = True
        if slab_slices > arr.shape[z_index]:
            slab_slices = arr.shape[z_index]
        slab_slices = int(slab_slices / z_chunks) * z_chunks
        num_z_splits = int(np.ceil(shape[z_index] / slab_slices))
        while num_z_splits % shrink_factors[z_index] > 1:
            num_z_splits += 1

        # num_y_splits = 1
        # num_x_splits = 1

        for slab_index in range(num_z_splits):
            # Process individual slabs
            if slice_planes:
                regions.extend(
                    _compute_plane_regions(
                        image,
                        dims,
                        arr,
                        shape,
                        x_index,
                        y_index,
                        z_index,
                        chunks,
                        shrink_factors,
                        slab_index,
                    )
                )
            else:
                region = [slice(arr.shape[i]) for i in range(arr.ndim)]
                region[z_index] = slice(
                    slab_index * z_chunks,
                    min((slab_index + 1) * z_chunks, arr.shape[z_index]),
                )
                regions.append(tuple(region))
    else:
        # 2D data - one region covering the whole array
        regions.append(tuple([slice(arr.shape[i]) for i in range(arr.ndim)]))

    return regions


def _compute_plane_regions(
    image,
    dims: Tuple[str, ...],
    arr: dask.array.Array,
    shape: Tuple[int, ...],
    x_index: int,
    y_index: int,
    z_index: int,
    chunks: Tuple[int, ...],
    shrink_factors: List[int],
    slab_index: int,
) -> List[Tuple[slice, ...]]:
    """Compute regions for a single z-slab, dividing into planes and strips if needed."""
    plane_regions = []
    z_chunks = chunks[z_index]
    y_chunks = chunks[y_index]
    x_chunks = chunks[x_index]

    # Calculate how to divide planes
    plane_bytes = memory_usage(image, {"z", "y"})
    plane_slices = min(
        int(np.ceil(config.memory_target / plane_bytes)),
        arr.shape[y_index],
    )
    slice_strips = False
    if plane_slices < y_chunks:
        plane_slices = y_chunks
        slice_strips = True
    if plane_slices > arr.shape[y_index]:
        plane_slices = arr.shape[y_index]
    plane_slices = int(plane_slices / y_chunks) * y_chunks
    num_y_splits = int(np.ceil(shape[y_index] / plane_slices))
    while num_y_splits % shrink_factors[y_index] > 1:
        num_y_splits += 1

    if slice_strips:
        # Need to subdivide further into strips
        strip_bytes = memory_usage(image, {"z", "y", "x"})
        strip_slices = min(
            int(np.ceil(config.memory_target / strip_bytes)),
            arr.shape[x_index],
        )
        strip_slices = max(strip_slices, x_chunks)
        if strip_slices > arr.shape[x_index]:
            strip_slices = arr.shape[x_index]
        strip_slices = int(strip_slices / x_chunks) * x_chunks
        num_x_splits = int(np.ceil(shape[x_index] / strip_slices))
        while num_x_splits % shrink_factors[x_index] > 1:
            num_x_splits += 1

        for plane_index in range(num_y_splits):
            for strip_index in range(num_x_splits):
                region = [slice(arr.shape[i]) for i in range(arr.ndim)]
                region[z_index] = slice(
                    slab_index * z_chunks,
                    min((slab_index + 1) * z_chunks, arr.shape[z_index]),
                )
                region[y_index] = slice(
                    plane_index * y_chunks,
                    min((plane_index + 1) * y_chunks, arr.shape[y_index]),
                )
                region[x_index] = slice(
                    strip_index * x_chunks,
                    min((strip_index + 1) * x_chunks, arr.shape[x_index]),
                )
                plane_regions.append(tuple(region))
    else:
        # Just divide into planes
        for plane_index in range(num_y_splits):
            region = [slice(arr.shape[i]) for i in range(arr.ndim)]
            region[z_index] = slice(
                slab_index * z_chunks,
                min((slab_index + 1) * z_chunks, arr.shape[z_index]),
            )
            region[y_index] = slice(
                plane_index * y_chunks,
                min((plane_index + 1) * y_chunks, arr.shape[y_index]),
            )
            plane_regions.append(tuple(region))

    return plane_regions


def _prepare_next_scale(
    image,
    index: int,
    nscales: int,
    multiscales: Multiscales,
    store: StoreLike,
    path: str,
    progress: Optional[Union[NgffProgress, NgffProgressCallback]],
) -> Optional[object]:
    """Prepare the next scale for processing if needed."""
    # No next scale if we're at the last one
    if index >= nscales - 1:
        return None
    # Minimize task graph depth
    if multiscales.scale_factors and multiscales.method and multiscales.chunks:
        for callback in image.computed_callbacks:
            callback()
        image.computed_callbacks = []

        image.data = dask.array.from_zarr(store, component=path)

        # Fetch scale factor for this index; used directly for index 0,
        # converted to a relative factor for index > 0
        next_multiscales_factor = multiscales.scale_factors[index]

        # For subsequent levels (index > 0), compute relative scale factor
        if index > 0:
            # If scales have been passed as list of integers
            if isinstance(next_multiscales_factor, int):
                next_multiscales_factor = (
                    next_multiscales_factor // multiscales.scale_factors[index - 1]
                )
            # If scales have been passed as dict of per-dimension factors
            else:
                updated_factors = {}
                for d, f in next_multiscales_factor.items():
                    updated_factors[d] = f // multiscales.scale_factors[index - 1][d]
                next_multiscales_factor = updated_factors

        next_multiscales = to_multiscales(
            image,
            scale_factors=[
                next_multiscales_factor,
            ],
            method=multiscales.method,
            chunks=multiscales.chunks,
            progress=progress,
            cache=False,
        )
        multiscales.images[index + 1] = next_multiscales.images[1]
        return next_multiscales.images[1]
    else:
        return multiscales.images[index + 1]


def to_ngff_zarr(
    store: StoreLike,
    multiscales: Multiscales,
    version: str = "0.4",
    overwrite: bool = True,
    use_tensorstore: bool = False,
    chunk_store: Optional[StoreLike] = None,
    progress: Optional[Union[NgffProgress, NgffProgressCallback]] = None,
    chunks_per_shard: Optional[
        Union[
            int,
            Tuple[int, ...],
            Dict[str, int],
        ]
    ] = None,
    enabled_rfcs: Optional[List[int]] = None,
    **kwargs,
) -> None:
    """
    Write an image pixel array and metadata to a Zarr store with the OME-NGFF standard data model.

    :param store: Store or path to directory in file system. If the path ends with .ozx, writes an RFC-9
    compliant zipped OME-Zarr file.
    :type  store: StoreLike

    :param multiscales: Multiscales OME-NGFF image pixel data and metadata. Can be generated with ngff_zarr.to_multiscales.
    :type  multiscales: Multiscales

    :param version: OME-Zarr specification version. For .ozx files, version 0.5 is required.
    :type  version: str, optional

    :param overwrite: If True, delete any pre-existing data in `store` before creating groups.
    :type  overwrite: bool, optional

    :param use_tensorstore: If True, write array using tensorstore backend.
    :type  use_tensorstore: bool, optional

    :param chunk_store: Separate storage for chunks. If not provided, `store` will be used
        for storage of both chunks and metadata.
    :type  chunk_store: StoreLike, optional

    :param progress: Optional progress logger
    :type  progress: RichDaskProgress

    :param chunks_per_shard: Number of chunks along each axis in a shard. If None, no sharding. For .ozx files, defaults to 2 if not specified. Requires OME-Zarr version >= 0.5.
    :type  chunks_per_shard: int, tuple, or dict, optional

    :param enabled_rfcs: List of RFC numbers to enable. If RFC 4 is included, anatomical orientation metadata will be preserved in the output.
    :type  enabled_rfcs: list of int, optional

    :param **kwargs: Passed to the zarr.create_array() or zarr.creation.create() function, e.g., compression options.
    """
    # RFC-9: Handle .ozx (zipped OME-Zarr) files
    if isinstance(store, (str, Path)) and is_ozx_path(store):
        if version != "0.5":
            raise ValueError(
                "RFC-9 zipped OME-Zarr (.ozx) requires OME-Zarr version 0.5"
            )

        # Default chunks_per_shard to 2 for .ozx files if not specified
        if chunks_per_shard is None:
            chunks_per_shard = 2

        # Determine if we should use memory or disk for intermediate storage
        total_memory_usage = sum(memory_usage(img) for img in multiscales.images)
        use_memory_store = total_memory_usage <= config.memory_target

        if use_memory_store:
            # Small dataset: use memory store
            from zarr.storage import MemoryStore

            temp_store = MemoryStore()
        else:
            # Large dataset: use temporary directory store in cache
            if hasattr(zarr.storage, "DirectoryStore"):
                LocalStore = zarr.storage.DirectoryStore
            else:
                LocalStore = zarr.storage.LocalStore

            temp_dir = tempfile.mkdtemp(
                dir=config.cache_store.path
                if hasattr(config.cache_store, "path")
                else None
            )
            temp_store = LocalStore(temp_dir)

        try:
            # Write to temporary store first
            _to_ngff_zarr_impl(
                temp_store,
                multiscales,
                version=version,
                overwrite=overwrite,
                use_tensorstore=False,  # Can't use tensorstore with memory/temp stores
                chunk_store=None,
                progress=progress,
                chunks_per_shard=chunks_per_shard,
                enabled_rfcs=enabled_rfcs,
                **kwargs,
            )

            # Write temp store to .ozx file
            write_store_to_zip(temp_store, store, version=version)
        finally:
            # Clean up temporary directory if used
            if not use_memory_store:
                import shutil

                if hasattr(zarr.storage, "DirectoryStore") and isinstance(
                    temp_store, zarr.storage.DirectoryStore
                ):
                    shutil.rmtree(temp_store.dir_path(), ignore_errors=True)
                elif hasattr(zarr.storage, "LocalStore") and isinstance(
                    temp_store, zarr.storage.LocalStore
                ):
                    shutil.rmtree(temp_store.root, ignore_errors=True)

        return

    # Standard (non-.ozx) path
    _to_ngff_zarr_impl(
        store,
        multiscales,
        version=version,
        overwrite=overwrite,
        use_tensorstore=use_tensorstore,
        chunk_store=chunk_store,
        progress=progress,
        chunks_per_shard=chunks_per_shard,
        enabled_rfcs=enabled_rfcs,
        **kwargs,
    )


def _to_ngff_zarr_impl(
    store: StoreLike,
    multiscales: Multiscales,
    version: str = "0.4",
    overwrite: bool = True,
    use_tensorstore: bool = False,
    chunk_store: Optional[StoreLike] = None,
    progress: Optional[Union[NgffProgress, NgffProgressCallback]] = None,
    chunks_per_shard: Optional[
        Union[
            int,
            Tuple[int, ...],
            Dict[str, int],
        ]
    ] = None,
    enabled_rfcs: Optional[List[int]] = None,
    **kwargs,
) -> None:
    """
    Internal implementation of to_ngff_zarr without .ozx handling.
    """
    # Setup and validation
    store_path = str(store) if isinstance(store, (str, Path)) else None

    _validate_ngff_parameters(version, chunks_per_shard, use_tensorstore, store)
    metadata, dimension_names, dimension_names_kwargs = _prepare_metadata(
        multiscales, version, enabled_rfcs
    )
    metadata_dict = asdict(metadata)
    metadata_dict = _pop_metadata_optionals(metadata_dict, enabled_rfcs)
    metadata_dict["@type"] = "ngff:Image"

    # Create Zarr root
    root = _create_zarr_root(store, chunk_store, version, overwrite, metadata_dict)

    # Format parameters
    zarr_format = 2 if version == "0.4" else 3
    format_kwargs = {"zarr_format": zarr_format}
    _zarr_kwargs = zarr_kwargs.copy()

    if version == "0.4" and kwargs.get("compressors") is not None:
        raise ValueError(
            "The argument `compressors` are not supported for OME-Zarr version 0.4. (Zarr v3). Use `compression` instead."
        )

    # Process each scale level
    nscales = len(multiscales.images)
    if progress:
        progress.add_multiscales_task("[green]Writing scales", nscales)

    next_image = multiscales.images[0]
    dims = next_image.dims
    previous_dim_factors = {d: 1 for d in dims}

    for index in range(nscales):
        if progress:
            progress.update_multiscales_task_completed(index + 1)

        image = next_image
        arr = image.data
        path = metadata.datasets[index].path
        parent = str(PurePosixPath(path).parent)

        # Create parent groups if needed
        if parent not in (".", "/"):
            array_dims_group = root.create_group(parent)
            array_dims_group.attrs["_ARRAY_DIMENSIONS"] = image.dims

        # Calculate dimension factors
        if index > 0 and index < nscales - 1 and multiscales.scale_factors:
            dim_factors = _dim_scale_factors(
                dims, multiscales.scale_factors[index], previous_dim_factors
            )
        else:
            dim_factors = {d: 1 for d in dims}
        previous_dim_factors = dim_factors

        # Configure sharding if needed
        # TODO check with recent updates to zarr by Ilan whether sharding can just be configured on zarr side.
        sharding_kwargs, internal_chunk_shape, arr = _configure_sharding(
            arr, chunks_per_shard, dims, kwargs.copy()
        )

        # Get the chunks - these are now the shards if sharding is enabled
        chunks = tuple([c[0] for c in arr.chunks])

        # For TensorStore, shards are the same as chunks when sharding is enabled
        shards = chunks if chunks_per_shard is not None else None

        # Determine write method based on memory requirements
        if memory_usage(image) > config.memory_target:
            _handle_large_array_writing(
                image,
                arr,
                store,
                path,
                dims,
                dim_factors,
                chunks,
                sharding_kwargs,
                _zarr_kwargs,
                format_kwargs,
                dimension_names_kwargs,
                use_tensorstore,
                store_path,
                zarr_format,
                dimension_names,
                internal_chunk_shape,
                shards,
                progress,
                index,
                nscales,
                **kwargs,
            )
        else:
            if isinstance(progress, NgffProgressCallback):
                progress.add_callback_task(
                    f"[green]Writing scale {index + 1} of {nscales}"
                )

            # For small arrays, write in one go
            region = tuple([slice(arr.shape[i]) for i in range(arr.ndim)])
            if use_tensorstore:
                _write_array_with_tensorstore(
                    store_path,
                    path,
                    arr,
                    chunks,
                    shards,
                    internal_chunk_shape,
                    zarr_format,
                    dimension_names,
                    region,
                    full_array_shape=arr.shape,
                    create_dataset=True,  # Always create for small arrays
                    **kwargs,
                )
            else:
                _write_array_direct(
                    arr,
                    store,
                    path,
                    sharding_kwargs,
                    _zarr_kwargs,
                    format_kwargs,
                    dimension_names_kwargs,
                    None,
                    None,
                    **kwargs,
                )

        # Prepare next scale if needed
        next_image = _prepare_next_scale(
            image, index, nscales, multiscales, store, path, progress
        )

    # Clean up callbacks
    for image in multiscales.images:
        for callback in image.computed_callbacks:
            callback()
        image.computed_callbacks = []

    # Consolidate metadata
    if IS_ZARR_V3_PLUS:
        with warnings.catch_warnings():
            # Ignore consolidated metadata warning
            warnings.filterwarnings("ignore", category=UserWarning)
            zarr.consolidate_metadata(store, **format_kwargs)
    else:
        # Zarr_format is used elsewhere but for this consolidate_metadata it is not an argument in zarr v2.
        if format_kwargs.get("zarr_format"):
            del format_kwargs["zarr_format"]
        zarr.consolidate_metadata(store, **format_kwargs)
