"""Read and write data from/to the ANU CTLab zarr data format.

This is an optional extra module, and must be explicitly installed to be used (e.g., ``pip install anu_ctlab_io[zarr]``)."""

from collections.abc import Mapping
from pathlib import Path
from typing import Any

import dask.array as da
import numpy as np
import zarr
from ome_zarr_models.common.coordinate_transformations import ValidTransform

from anu_ctlab_io._dataset import Dataset
from anu_ctlab_io._datatype import DataType
from anu_ctlab_io._parse_history import History
from anu_ctlab_io._voxel_properties import VoxelUnit
from anu_ctlab_io.zarr._writer import OMEZarrVersion, dataset_to_zarr

__all__ = ["OMEZarrVersion", "dataset_from_zarr", "dataset_to_zarr"]


def dataset_from_zarr(path: Path, **kwargs: Any) -> Dataset:
    """Loads a :any:`Dataset` from the path to a zarr.

    This method is used by :any:`Dataset.from_path`, by preference call that constructor directly.

    :param Path: The path to the zarr to be loaded
    :param kwargs: Currently this method consumes no kwargs, but will pass provided kwargs to ``dask.Array.from_path``."""
    data = zarr.open(path, zarr_format=3)
    if isinstance(data, zarr.Array):
        return _dataset_from_zarr_array(path, **kwargs)
    elif isinstance(data, zarr.Group):
        return _dataset_from_zarr_group(path, **kwargs)
    else:
        raise TypeError(f"Unsupported zarr data type: {type(data)}") from None


def _dataset_from_zarr_array(path: Path, **kwargs: Any) -> Dataset:
    za = zarr.open_array(path, zarr_format=3)
    data = da.from_zarr(za, **kwargs)  # type: ignore[no-untyped-call]
    attrs: dict[str, Any] = dict(za.attrs)["mango"]  # type: ignore[assignment]
    dimension_names: tuple[str, ...] = za.metadata.dimension_names  # type: ignore[assignment, union-attr]
    voxel_unit = VoxelUnit.from_str(attrs["voxel_unit"])
    voxel_size = attrs["voxel_size_xyz"]
    datatype = DataType.from_basename(attrs["basename"])
    history = attrs["history"]
    dataset_id_raw = attrs.get("dataset_id")
    dataset_id = str(dataset_id_raw) if dataset_id_raw is not None else None

    return Dataset(
        data=data,
        dimension_names=dimension_names,
        datatype=datatype,
        voxel_unit=voxel_unit,
        voxel_size=voxel_size,
        history=history,
        dataset_id=dataset_id,
    )


def _dataset_from_zarr_group(path: Path, **kwargs: Any) -> Dataset:
    from ome_zarr_models.v05.image import Image

    zg = zarr.open_group(path, zarr_format=3)
    ome = Image.from_zarr(zg)
    multiscales = ome.ome_attributes.multiscales
    if len(multiscales) > 1:
        raise NotImplementedError(
            "Only single multiscale images are currently supported."
        ) from None

    # Extract the first multiscale and first dataset
    # These are guaranteed to exist by ome_zarr_metadata validation
    # Assumes the first dataset is the full resolution one
    multiscale = multiscales[0]
    dataset = multiscale.datasets[0]
    data = da.from_zarr(path, component=dataset.path, **kwargs)  # type: ignore[no-untyped-call]

    if len(multiscale.axes) != 3:
        raise ValueError(
            f"Provided zarr has {len(multiscale.axes)} axes, should have 3."
        ) from None

    dimension_names = tuple(
        [str(axis.name) for axis in multiscale.axes]
    )  # NOTE: str cast needed due to a mismatch in ome_zarr_models with the ome_zarr spec

    def _cast_unit(unit: str | Any | None) -> str | None:
        if unit is None or isinstance(unit, str):
            return unit
        raise ValueError(f"Unsupported unit type: {unit}")

    voxel_unit_list: tuple[str | None, ...] = tuple(
        [_cast_unit(axis.unit) for axis in multiscale.axes]
    )
    if not all(u == voxel_unit_list[0] for u in voxel_unit_list):
        raise ValueError(
            f"Provided zarr has differing units {voxel_unit_list}, these should all be equal."
        ) from None
    if voxel_unit_list[0] is None:
        voxel_unit = VoxelUnit.VOXEL
    else:
        voxel_unit = VoxelUnit.from_str(voxel_unit_list[0])

    # Calculate the voxel size from the transformations
    def extract_vector_scale(
        transformations: "ValidTransform | None",
    ) -> tuple[float, float, float]:
        """Extracts the scale from a list of coordinate transformations, expects VectorScale."""
        from ome_zarr_models.common.coordinate_transformations import VectorScale

        if transformations and len(transformations) > 0:
            scale_transform = transformations[0]
            if isinstance(scale_transform, VectorScale):
                scale = scale_transform.scale
                if len(scale) != 3:
                    raise ValueError(
                        f"Provided zarr has {len(scale)} scale factors provided, should have 3."
                    ) from None
                return (scale[0], scale[1], scale[2])
            else:
                raise NotImplementedError(
                    "Only vector scales are currently supported for coordinate transformations."
                ) from None
        return (1.0, 1.0, 1.0)

    voxel_size_dataset = extract_vector_scale(dataset.coordinateTransformations)
    voxel_size_root = extract_vector_scale(multiscale.coordinateTransformations)
    voxel_size = tuple(np.array(voxel_size_dataset) * np.array(voxel_size_root))

    if "mango" not in zg.attrs:
        # Handle a plain OME-Zarr dataset that has no mango attributes
        datatype = None
        history: History = {}
        dataset_id = None
    else:
        mango_attrs = zg.attrs["mango"]
        if not isinstance(mango_attrs, Mapping):
            raise TypeError(
                f'Expected "mango" attribute to be a Mapping, got {type(mango_attrs)}'
            ) from None
        basename = mango_attrs["basename"]
        if not isinstance(basename, str):
            raise TypeError(
                f'Expected mango "basename" to be a str, got {type(basename)}'
            ) from None
        datatype = DataType.from_basename(basename)
        history = mango_attrs["history"]  # type: ignore[assignment]
        dataset_id_raw = mango_attrs.get("dataset_id")
        dataset_id = str(dataset_id_raw) if dataset_id_raw is not None else None

    return Dataset(
        data=data,
        dimension_names=dimension_names,
        datatype=datatype,
        voxel_unit=voxel_unit,
        voxel_size=voxel_size,
        history=history,
        dataset_id=dataset_id,
    )
