Import ome zarr

def _build_xy_roi_table(
    ome_zarr_image: OmeZarrContainer,
    grid_YX_shape: tuple[int, int] | None = None,
    backend: AVAILABLE_TABLE_BACKENDS = "csv",
    overwrite: bool = False,
) -> None:
    """Build a grid_ROI_table with a rectangular grid of ROIs covering the whole image.

    Args:
        ome_zarr_image: OME-Zarr image container.
        grid_YX_shape: Number of ROIs along Y and X. The image is split into a
            `grid_YX_shape[0]` by `grid_YX_shape[1]` grid of ROIs.
        backend: Table backend to use.
        overwrite: Whether to overwrite an existing `grid_ROI_table`.
    """
    if grid_YX_shape is None:
        raise ValueError("grid_YX_shape cannot be None when building grid ROI table.")
    image = ome_zarr_image.get_image()
    pixel_size = image.pixel_size
    image_shape = image.shape
    image_YX_shape = image_shape[-2:]
    roi_id = 0
    rois = []
    z_start = 0
    z_ind = image.axes_handler.get_index("z")
    z_length = image_shape[z_ind] if z_ind is not None else 1

    t_start = 0
    t_ind = image.axes_handler.get_index("t")
    t_length = image_shape[t_ind] if t_ind is not None else 1

    # `grid_YX_shape` is the number of ROIs along each axis: split the image
    # into an `n_y` by `n_x` grid of tiles of size `len_y` by `len_x` pixels
    # (the last tile along each axis is clipped to the image edge).
    n_y, n_x = grid_YX_shape
    len_y = math.ceil(image_YX_shape[0] / n_y)
    len_x = math.ceil(image_YX_shape[1] / n_x)
    for ind_y in range(n_y):
        start_y = ind_y * len_y
        y_length = min(image_YX_shape[0], start_y + len_y) - start_y
        for ind_x in range(n_x):
            start_x = ind_x * len_x
            x_length = min(image_YX_shape[1], start_x + len_x) - start_x
            roi_pixels = Roi.from_values(
                name=f"ROI_{roi_id}",
                slices={
                    "x": (start_x, x_length),
                    "y": (start_y, y_length),
                    "z": (z_start, z_length),
                    "t": (t_start, t_length),
                },
                space="pixel",
            )
            roi_world = roi_pixels.to_world(pixel_size=pixel_size)
            rois.append(roi_world)
            roi_id += 1
    table = RoiTable(rois=rois)
    ome_zarr_image.add_table(
        name="grid_ROI_table", table=table, overwrite=overwrite, backend=backend
    )

def _process_plate(
    *,
    zarr_path: str,
    ome_zarr_plate: OmeZarrPlate,
    advanced_options: AdvancedOptions,
) -> list[dict[str, Any]]:
    """For each image in the plate, create an image list update dict."""
    image_list_updates = []
    plate_name = zarr_path.rstrip("/").split("/")[-1]
    for path, image in ome_zarr_plate.get_images().items():
        row, column, _ = path.split("/")
        well_name = f"{row}{int(column):02d}"
        attributes = {
            "plate": plate_name,
            "well": well_name,
        }
        image_zarr_path = f"{zarr_path}/{path}"
        _updates = _process_single_image(
            zarr_path=image_zarr_path,
            ome_zarr_image=image,
            advanced_options=advanced_options,
            attributes=attributes,
        )
        image_list_updates.extend(_updates)
    return image_list_updates

def _process_single_image(
    *,
    zarr_path: str,
    ome_zarr_image: OmeZarrContainer,
    advanced_options: AdvancedOptions,
    attributes: dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
    """Optionally generate ROI tables and update omero metadata for a single image.

    Args:
        zarr_path: Absolute path to the image Zarr group.
        ome_zarr_image: OME-Zarr image container.
        advanced_options: Advanced options for the import.
        attributes: Optional image attributes to include in the update dict.
    """
    if advanced_options.add_image_roi_table:
        table = ome_zarr_image.build_image_roi_table()
        ome_zarr_image.add_table(
            name="image_ROI_table",
            table=table,
            overwrite=advanced_options.overwrite,
            backend=advanced_options.table_backend,
        )

    if advanced_options.grid_roi_table is not None:
        grid_roi_table = advanced_options.grid_roi_table
        _build_xy_roi_table(
            ome_zarr_image=ome_zarr_image,
            grid_YX_shape=(grid_roi_table.grid_y_shape, grid_roi_table.grid_x_shape),
            overwrite=advanced_options.overwrite,
            backend=advanced_options.table_backend,
        )

    if advanced_options.update_omero_metadata:
        image = ome_zarr_image.get_image()
        channel_names = image.channel_labels
        wavelengths = image.channels_meta.channel_wavelength_ids
        ome_zarr_image.set_channel_meta(
            channel_meta=ChannelsMeta.default_init(
                labels=channel_names,
                wavelength_id=wavelengths,
            )
        )

    type_updates = {
        "is_3D": ome_zarr_image.is_3d,
        # TODO add this after v1 testing is removed
        # "is_time_series": ome_zarr_image.is_time_series,
    }
    image_list_update = {
        "zarr_url": zarr_path,
        "types": type_updates,
    }
    if attributes is not None:
        image_list_update["attributes"] = attributes
    return [image_list_update]

def _process_well(
    *,
    zarr_path: str,
    ome_zarr_well: OmeZarrWell,
    advanced_options: AdvancedOptions,
) -> list[dict[str, Any]]:
    """For each image in the well, create an image list update dict."""
    image_list_updates = []
    *_, row, column = zarr_path.rstrip("/").split("/")
    attributes = {
        "well": f"{row}{int(column):02d}",
    }
    for path in ome_zarr_well.paths():
        ome_zarr_image = ome_zarr_well.get_image(path)
        image_zarr_path = f"{zarr_path}/{path}"
        _updates = _process_single_image(
            zarr_path=image_zarr_path,
            ome_zarr_image=ome_zarr_image,
            advanced_options=advanced_options,
            attributes=attributes,
        )
        image_list_updates.extend(_updates)
    return image_list_updates

@validate_call
def import_ome_zarr(
    *,
    # Fractal parameters
    zarr_dir: str,
    # Core parameters
    zarr_name: str,
    advanced_options: AdvancedOptions = Field(default_factory=AdvancedOptions),
) -> dict[str, Any]:
    """Import a single OME-Zarr into Fractal.

    The single OME-Zarr can be a full OME-Zarr HCS plate or an individual
    OME-Zarr image. The image needs to be in the zarr_dir as specified by the
    dataset. This task registers the OME-Zarr with Fractal so it can be used
    in processing workflows, and optionally adds new ROI tables to the existing
    OME-Zarr.

    Args:
        zarr_dir: Path of the directory where the OME-Zarr is located
            (standard argument for Fractal tasks, managed by Fractal server).
        zarr_name: The OME-Zarr name, without its parent folder. The parent
            folder is provided by zarr_dir; e.g. `zarr_name="array.zarr"`,
            if the OME-Zarr path is in `/zarr_dir/array.zarr`.
        advanced_options: Advanced options for importing an OME-Zarr, including
            whether to add `image_ROI_table`/`grid_ROI_table` tables, whether
            to update Omero-channels metadata, the table backend to use, and
            whether new tables can overwrite existing ones.
    """
    zarr_path = f"{zarr_dir.rstrip('/')}/{zarr_name}"
    logger.info(f"Zarr path: {zarr_path}")

    ome_zarr = open_unknown_container(zarr_path)
    image_list_updates = []
    if isinstance(ome_zarr, OmeZarrPlate):
        image_list_updates = _process_plate(
            zarr_path=zarr_path,
            ome_zarr_plate=ome_zarr,
            advanced_options=advanced_options,
        )
    elif isinstance(ome_zarr, OmeZarrWell):
        image_list_updates = _process_well(
            zarr_path=zarr_path,
            ome_zarr_well=ome_zarr,
            advanced_options=advanced_options,
        )
    elif isinstance(ome_zarr, OmeZarrContainer):
        image_list_updates = _process_single_image(
            zarr_path=zarr_path,
            ome_zarr_image=ome_zarr,
            advanced_options=advanced_options,
        )
    else:
        raise ValueError(
            f"Unexpected OME-NGFF type for OME-Zarr at {zarr_path}: "
            f"{type(ome_zarr).__name__}"
        )

    image_list_changes = {
        "image_list_updates": image_list_updates,
    }
    return image_list_changes

def open_unknown_container(
    zarr_path: str,
) -> OmeZarrContainer | OmeZarrWell | OmeZarrPlate:
    """Detect the OME-NGFF type of the OME-Zarr, based on its root metadata.

    The OME-NGFF type can be "plate", "well" or "image". If the OME-Zarr does
    not contain valid OME-NGFF metadata, an error is raised.

    Args:
        zarr_path: Path to the OME-Zarr.

    Returns:
        OmeZarrContainer, OmeZarrWell, or OmeZarrPlate
    """
    errors = []
    try:
        ome_zarr = open_ome_zarr_container(zarr_path)
        return ome_zarr
    except Exception as e:
        errors.append(e)

    try:
        ome_zarr = open_ome_zarr_plate(zarr_path)
        return ome_zarr
    except Exception as e:
        errors.append(e)

    try:
        ome_zarr = open_ome_zarr_well(zarr_path)
        return ome_zarr
    except Exception as e:
        errors.append(e)

    base_error = f"Could not detect OME-NGFF type of OME-Zarr at {zarr_path}."
    error_messages = "\n".join([f"{type(e).__name__}: {e!s}" for e in errors])
    raise ValueError(f"{base_error}\nErrors:\n{error_messages}")