Built-in Plugins¶

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def composite(
    ds: xr.Dataset, composite_method: str = _DEFAULT_COMPOSITE_METHOD
) -> xr.Dataset:
    """Post-reproject processor that creates a temporal composite.

    Reduces the ``time`` dimension using a statistical method (median, mean, or
    max). Useful for creating cloud-free or best-pixel composites.

    Args:
        ds: Input dataset with a ``time`` dimension.
        composite_method: Statistical compositing method ('median', 'mean', 'max', or 'min').

    Returns:
        Dataset with the ``time`` dimension reduced to a single step.

    Raises:
        ValueError: If the method is unknown or ``time`` is missing.
    """
    method = composite_method

    if "time" not in ds.dims:
        raise ValueError("composite requires a 'time' dimension in the dataset.")

    if method == _MEDIAN:
        return ds.median(dim="time", keep_attrs=True)
    if method == _MEAN:
        return ds.mean(dim="time", keep_attrs=True)
    if method == _MAX:
        return ds.max(dim="time", keep_attrs=True)
    if method == _MIN:
        return ds.min(dim="time", keep_attrs=True)

    raise ValueError(
        f"composite: unknown method '{method}'. Use '{_MEDIAN}', '{_MEAN}', '{_MAX}', or '{_MIN}'."
    )

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def ndvi(ds: xr.Dataset, ndvi_nir_band: str, ndvi_red_band: str) -> xr.Dataset:
    """Post-reproject processor that computes the Normalised Difference Vegetation Index.

    NDVI is computed on co-registered pixels after reprojection so that the
    band math is physically correct. The source red and NIR bands are dropped
    and only the ``ndvi`` variable is retained.

    Args:
        ds: Input dataset containing red and NIR variables.
        ndvi_nir_band: Name of the NIR band.
        ndvi_red_band: Name of the red band.

    Returns:
        Dataset with a single ``ndvi`` variable. Source bands are removed.

    Raises:
        ValueError: If NIR/red bands are not found.
    """
    if ndvi_nir_band not in ds.data_vars:
        raise ValueError(
            f"ndvi: NIR band '{ndvi_nir_band}' not found. Available: {list(ds.data_vars)}"
        )
    if ndvi_red_band not in ds.data_vars:
        raise ValueError(
            f"ndvi: red band '{ndvi_red_band}' not found. Available: {list(ds.data_vars)}"
        )

    # Cast to float32 to prevent uint16 underflow and allow decimal division
    nir = ds[ndvi_nir_band].astype("float32")
    red = ds[ndvi_red_band].astype("float32")
    denom = nir + red
    ndvi_val = (nir - red) / denom
    ndvi_val = ndvi_val.where(denom != 0)

    result = ds.drop_vars([ndvi_nir_band, ndvi_red_band])
    result["ndvi"] = ndvi_val
    return result

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def ndwi(ds: xr.Dataset, ndwi_green_band: str, ndwi_nir_band: str) -> xr.Dataset:
    """Post-reproject processor that computes the Normalised Difference Water Index.

    NDWI is computed on co-registered pixels after reprojection so that the
    band math is physically correct. The source green and NIR bands are dropped
    and only the ``ndwi`` variable is retained. Uses the McFeeters formulation:
    ``(green - nir) / (green + nir)``.

    Args:
        ds: Input dataset containing green and NIR variables.
        ndwi_green_band: Name of the green band.
        ndwi_nir_band: Name of the NIR band.

    Returns:
        Dataset with a single ``ndwi`` variable. Source bands are removed.

    Raises:
        ValueError: If green/NIR bands are not found.
    """
    if ndwi_green_band not in ds.data_vars:
        raise ValueError(
            f"ndwi: green band '{ndwi_green_band}' not found. Available: {list(ds.data_vars)}"
        )
    if ndwi_nir_band not in ds.data_vars:
        raise ValueError(
            f"ndwi: NIR band '{ndwi_nir_band}' not found. Available: {list(ds.data_vars)}"
        )

    # Cast to float32 to prevent uint16 underflow and allow decimal division
    green = ds[ndwi_green_band].astype("float32")
    nir = ds[ndwi_nir_band].astype("float32")
    denom = green + nir
    ndwi_val = (green - nir) / denom
    ndwi_val = ndwi_val.where(denom != 0)

    result = ds.drop_vars([ndwi_green_band, ndwi_nir_band])
    result["ndwi"] = ndwi_val
    return result

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def normalize(
    ds: xr.Dataset, normalize_method: str = _DEFAULT_NORMALIZE_METHOD
) -> xr.Dataset:
    """Post-reproject processor that normalises pixel values per band.

    Supports min-max scaling and z-score normalisation. NaN pixels are ignored
    when computing statistics.

    Args:
        ds: Input dataset.
        normalize_method: Scaling method ('minmax' or 'zscore').

    Returns:
        Dataset with normalised variables.

    Raises:
        ValueError: If the method is unknown.
    """
    method = normalize_method

    if method not in (_MINMAX, _ZSCORE):
        raise ValueError(
            f"normalize: unknown method '{method}'. Use 'minmax' or 'zscore'."
        )

    normalized = ds.copy()
    for var in normalized.data_vars:
        da = normalized[var]
        if method == _MINMAX:
            vmin = da.min(skipna=True)
            vmax = da.max(skipna=True)
            denom = vmax - vmin
            denom = denom.where(denom != 0, 1)
            normalized[var] = (da - vmin) / denom
        else:  # _ZSCORE
            mean = da.mean(skipna=True)
            std = da.std(skipna=True)
            std = std.where(std != 0, 1)
            normalized[var] = (da - mean) / std

    return normalized

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def qa_mask(ds: xr.Dataset, qa_band: str, qa_mask_bits: list[int]) -> xr.Dataset:
    """Pre-reproject processor that masks cloudy/invalid pixels using a QA band.

    Pixels where the QA band matches any of the specified bit-masks are set to
    NaN in all *other* data variables. The QA band itself is dropped after
    masking.

    Args:
        ds: Input dataset containing a QA variable.
        qa_band: Name of the QA variable.
        qa_mask_bits: List of bit indices to mask.

    Returns:
        Dataset with masked pixels set to NaN. The QA band is removed.

    Raises:
        ValueError: If required params are missing or the QA band does not exist.
    """
    if qa_band not in ds.data_vars:
        raise ValueError(
            f"qa_mask: QA band '{qa_band}' not found in dataset. Available: {list(ds.data_vars)}"
        )

    qa = ds[qa_band]
    mask = np.zeros(qa.shape, dtype=bool)
    qa_arr = qa.values
    for bit in qa_mask_bits:
        mask |= ((qa_arr >> bit) & 1).astype(bool)

    masked = ds.drop_vars(qa_band)
    for var in masked.data_vars:
        masked[var] = masked[var].where(~mask)

    return masked

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def select_bands(ds: xr.Dataset, bands: list[str]) -> xr.Dataset:
    """Pre-reproject processor that keeps only specified data variables.

    Dropping bands before reprojection reduces the data volume that the
    expensive reproject step must process.

    Args:
        ds: Input dataset.
        bands: List of band names to keep.

    Returns:
        A new dataset containing only the requested variables.

    Raises:
        ValueError: If requested bands are not found.
    """
    keep = [str(b) for b in bands]
    missing = [b for b in keep if b not in ds.data_vars]
    if missing:
        raise ValueError(
            f"select_bands: requested bands not found in dataset: {missing}"
        )

    return ds[keep]

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def read_odc_stac(
    task: ExtractionTask,
    gdal_env: dict[str, Any] | None = None,
    **kwargs: Any,
) -> xr.Dataset:
    """Load STAC assets using ``odc.stac.load``.

    Reconstructs :class:`pystac.Item` objects from ``task.stac_items`` and
    returns a dataset tagged with temporal bounds in ``ds.attrs``.

    Args:
        task: The extraction task containing STAC items and read context.
        gdal_env: Optional GDAL configuration options to merge with odc-stac's
            cloud defaults (e.g. ``{"GDAL_HTTP_MAX_RETRY": "3"}``). These are
            forwarded to ``odc.loader.configure_rio`` and become the process-wide
            default for rasterio sessions used by this reader.
        **kwargs: Keyword arguments forwarded to ``odc.stac.load``. AEREO
            injects sensible defaults for ``chunks``, ``bands``, and ``bbox``
            only when they are not already provided.

    Returns:
        xr.Dataset (potentially dask-backed) in the native CRS of the
        STAC items.

    Raises:
        ImportError: If ``odc-stac`` is not installed.
        ValueError: If ``task.stac_items`` is empty.
    """
    if odc_load is None:  # pragma: no cover
        raise ImportError(
            "odc-stac is required for read_odc_stac. "
            "Install it with: pip install 'aereo[stac]'"
        )

    items = task.stac_items
    if not items:
        raise ValueError(
            "read_odc_stac requires at least one STAC item in task.stac_items. "
            "Ensure the search plugin (e.g. search_stac) stores full STAC "
            "item dictionaries in the assets."
        )

    params = dict(kwargs)

    if "chunks" not in params:
        params["chunks"] = {}

    if "bands" not in params and "channel_id" in task.assets.columns:
        params["bands"] = list(task.assets["channel_id"].unique())

    if task.bbox is not None and "bbox" not in params:
        params["bbox"] = task.bbox

    _ensure_rio_configured(gdal_env)

    ds: xr.Dataset = odc_load(items, **params)

    infer_dataset_time_bounds(ds)
    # remove time dimension, we are working only with a single time slice per task
    if "time" in ds.dims:
        ds = ds.isel(time=0).drop_vars("time", errors="ignore")

    return ds

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def reproject_odc(
    ds: xr.Dataset,
    geobox: Any | None = None,
    crs: str | None = None,
    resolution: float | None = None,
    resampling: str = "nearest",
    fill_value: Any = None,
    dtype: Any = None,
) -> xr.Dataset:
    """Reproject *ds* using ``odc.geo.xr.reproject``.

    Exactly one of *geobox* or both *crs* and *resolution* should be provided.
    When *geobox* is supplied (the normal case in ``reproject_mode="grid"``),
    the dataset is warped to that geobox. Otherwise it is warped to the
    requested CRS and resolution using the dataset's own extent.

    Args:
        ds: Input dataset.
        geobox: Target ``odc.geo.GeoBox`` (optional).
        crs: Target CRS string (optional).
        resolution: Target resolution in metres (optional).
        resampling: Resampling method (e.g. 'nearest', 'bilinear').
        fill_value: Optional fill value for out-of-bounds pixels.
        dtype: Optional output dtype.

    Returns:
        Reprojected ``xr.Dataset``.
    """
    from odc.geo.geobox import GeoBox  # type: ignore[reportMissingTypeStubs]
    from odc.geo.xr import xr_reproject  # type: ignore[reportAttributeAccessIssue]

    kwargs: dict[str, Any] = {"resampling": resampling}
    if fill_value is not None:
        kwargs["fill_value"] = fill_value
    if dtype is not None:
        kwargs["dtype"] = dtype

    if geobox is not None:
        return xr_reproject(ds, geobox, **kwargs)

    if crs is None or resolution is None:
        raise ValueError(
            "reproject_odc requires either 'geobox' or both 'crs' and 'resolution'."
        )

    target_geobox = GeoBox.from_bbox(
        ds.rio.bounds(),
        crs=crs,
        resolution=resolution,
    )
    return xr_reproject(ds, target_geobox, **kwargs)

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def reproject_swath(
    ds: xr.Dataset,
    geobox: Any | None = None,
    crs: str | None = None,
    resolution: float | None = None,
    buffer: float = 0.05,
    max_distance: float = 3_000.0,
    fill_value: Any = np.nan,
    mask_invalid: bool = True,
    nprocs: int = 1,
) -> xr.Dataset:
    """Reproject a swath dataset using pyresample's nearest-neighbour resampler.

    The input dataset must contain ``lons`` and ``lats`` variables (or
    coordinates) that define the swath geometry for every pixel.

    This is the AEREO builtins equivalent of pyresample-based swath
    reprojection. It builds a ``pyresample.SwathDefinition`` from the source
    lon/lat arrays and a ``pyresample.AreaDefinition`` from the target geobox,
    then calls ``pyresample.kd_tree.resample_nearest``.

    By default, NaN and infinite source pixels are excluded from the source
    geometry so that VIIRS bow-tie gaps do not pollute the output. Set
    ``mask_invalid=False`` to get plain pyresample nearest-neighbour behaviour,
    where a target pixel whose nearest source is NaN receives NaN.

    Args:
        ds: Input swath dataset with ``lons`` and ``lats`` (or ``longitude``
            and ``latitude``).
        geobox: Target ``odc.geo.GeoBox`` (optional). Used in grid mode.
        crs: Target CRS string (optional). Used in raw mode with ``resolution``.
        resolution: Target resolution in metres (optional). Used in raw mode.
        buffer: Deprecated. Kept for backward compatibility; no longer used.
        max_distance: Maximum source-to-target distance in metres before a
            target pixel is filled with ``fill_value`` (passed as
            ``radius_of_influence`` to pyresample).
        fill_value: Value for out-of-bounds / distant pixels.
        mask_invalid: If True, exclude NaN/inf source pixels from the search.
        nprocs: Number of processor cores for pyresample (default 1).

    Returns:
        Reprojected ``xr.Dataset`` on a regular ``y``/``x`` grid.
    """
    _ = buffer  # noqa: F841
    from pyresample import AreaDefinition, SwathDefinition  # type: ignore[reportMissingTypeStubs]
    from pyresample.kd_tree import resample_nearest  # type: ignore[reportMissingTypeStubs]

    if not (
        ("lons" in ds and "lats" in ds) or ("longitude" in ds and "latitude" in ds)
    ):
        raise ValueError(
            "Input dataset must contain 'lons' and 'lats' variables or coordinates."
        )

    lons_var = "lons" if "lons" in ds else "longitude"
    lats_var = "lats" if "lats" in ds else "latitude"
    lons = _as_numpy(ds[lons_var])
    lats = _as_numpy(ds[lats_var])

    if lons.ndim != 2 or lats.ndim != 2:
        raise ValueError("reproject_swath expects 2-D 'lons' and 'lats' arrays.")

    target_geobox = _resolve_target_geobox(lons, lats, geobox, crs, resolution)
    x_coords, y_coords, target_lons, target_lats = _target_grid_from_geobox(
        target_geobox
    )

    area_def = AreaDefinition(
        area_id="aereo_pyresample_target",
        description="AEREO pyresample target grid",
        proj_id="aereo",
        projection=str(target_geobox.crs),
        width=target_geobox.shape.x,
        height=target_geobox.shape.y,
        area_extent=(
            float(target_geobox.boundingbox.left),
            float(target_geobox.boundingbox.bottom),
            float(target_geobox.boundingbox.right),
            float(target_geobox.boundingbox.top),
        ),
    )

    swath_shape = lons.shape
    target_shape = (target_geobox.shape.y, target_geobox.shape.x)
    valid_coords = np.isfinite(lons) & np.isfinite(lats)

    data_vars: dict[str, xr.DataArray] = {}
    skip_vars = {lons_var, lats_var}

    for name, da in ds.data_vars.items():
        var_name = str(name)
        if var_name in skip_vars:
            continue
        if da.shape[-2:] != swath_shape:
            continue

        arr = _as_numpy(da)
        if np.isnan(fill_value) and not np.issubdtype(arr.dtype, np.floating):
            arr = arr.astype(np.float64)

        extra_shape = arr.shape[:-2]
        n_extra = int(np.prod(extra_shape, dtype=np.int64)) if extra_shape else 1

        # pyresample expects (n_pixels, n_channels) and returns (y, x, n_channels).
        flat = arr.reshape((n_extra, -1)).T

        if mask_invalid and np.issubdtype(arr.dtype, np.floating):
            data_valid = np.isfinite(arr)
            if data_valid.ndim > 2:
                data_valid = data_valid.all(axis=tuple(range(data_valid.ndim - 2)))
            combined_valid = valid_coords & data_valid
        else:
            combined_valid = None

        if combined_valid is not None:
            if not combined_valid.any():
                # No valid pixels for this variable; emit a filled grid.
                remapped = np.full(
                    (*extra_shape, *target_shape), fill_value, dtype=arr.dtype
                )
                data_vars[var_name] = xr.DataArray(
                    remapped,
                    dims=tuple(str(d) for d in da.dims[:-2]) + ("y", "x"),
                    coords={"y": y_coords, "x": x_coords},
                    attrs=da.attrs,
                    name=var_name,
                )
                continue

            flat = flat[combined_valid.ravel(), :]
            swath_def = SwathDefinition(
                lons=lons.ravel()[combined_valid.ravel()],
                lats=lats.ravel()[combined_valid.ravel()],
            )
        else:
            swath_def = SwathDefinition(lons=lons.ravel(), lats=lats.ravel())

        resampled = resample_nearest(
            swath_def,
            flat,
            area_def,
            radius_of_influence=max_distance,
            fill_value=fill_value,
            nprocs=nprocs,
        )

        resampled_arr = np.asarray(resampled)
        remapped = np.moveaxis(resampled_arr.reshape((*target_shape, n_extra)), -1, 0)
        remapped = remapped.reshape((*extra_shape, *target_shape))

        data_vars[var_name] = xr.DataArray(
            remapped,
            dims=tuple(str(d) for d in da.dims[:-2]) + ("y", "x"),
            coords={"y": y_coords, "x": x_coords},
            attrs=da.attrs,
            name=var_name,
        )

    if not data_vars:
        raise ValueError("No data variables found matching the swath shape.")

    out = xr.Dataset(data_vars, attrs=ds.attrs)
    out = out.rio.write_crs(str(target_geobox.crs))
    if fill_value is not None and not (
        isinstance(fill_value, float) and np.isnan(fill_value)
    ):
        for name in out.data_vars:
            out[str(name)].rio.write_nodata(fill_value, inplace=True)
    return out

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def search_earthaccess(
    collections: Mapping[str, Sequence[str]] | Sequence[str] | None,
    intersects: BaseGeometry | dict[str, Any] | str | Path | None,
    start_datetime: datetime | None,
    end_datetime: datetime | None,
    search_params: dict[str, Any] | None = None,
) -> GeoDataFrame[AssetSchema]:
    """Search NASA Earthdata using the earthaccess library.

    Args:
        collections: Mapping of collection -> asset keys, or list of collections.
        intersects: AOI geometry as a Shapely object, GeoJSON dict, or path.
        start_datetime: Optional start of temporal window.
        end_datetime: Optional end of temporal window.
        search_params: Extra arguments forwarded to ``earthaccess.search_data``.

    Returns:
        A GeoDataFrame of matched assets.

    Raises:
        ImportError: If the ``earthaccess`` library is not installed.
    """
    try:
        import earthaccess  # type: ignore
    except ImportError as e:
        raise ImportError(
            "The 'earthaccess' library is required to use search_earthaccess. "
            "Please install it (e.g., 'pip install earthaccess')."
        ) from e

    collections, _ = build_collection_asset_filters(collections)
    if not collections:
        return empty_asset_result()

    kwargs: dict[str, Any] = {"short_name": collections}

    if start_datetime is not None and end_datetime is not None:
        kwargs["temporal"] = (
            start_datetime.strftime("%Y-%m-%d %H:%M:%S"),
            end_datetime.strftime("%Y-%m-%d %H:%M:%S"),
        )

    normalized_intersects = normalize_geometry_input(intersects)
    if normalized_intersects is not None:
        bounds = getattr(normalized_intersects, "bounds", None)
        if bounds is not None:
            kwargs["bounding_box"] = bounds

    kwargs.update(search_params or {})

    try:
        granules = earthaccess.search_data(**kwargs)
    except Exception as e:
        logger.error("earthaccess search failed", error=str(e), **kwargs)
        return empty_asset_result()

    if not granules:
        return empty_asset_result()

    rows = []
    for g in granules:
        row = _process_granule(
            g, collections, cast(BaseGeometry | None, normalized_intersects)
        )
        if row is not None:
            rows.append(row)

    if not rows:
        return empty_asset_result()

    gdf = gpd.GeoDataFrame(rows, geometry="geometry")
    gdf["start_time"] = pd.to_datetime(gdf["start_time"])
    gdf["end_time"] = pd.to_datetime(gdf["end_time"])

    return cast(GeoDataFrame, AssetSchema.validate(gdf))

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def search_stac(
    collections: Mapping[str, Sequence[str]] | Sequence[str] | None,
    intersects: BaseGeometry | dict[str, Any] | str | Path | None,
    start_datetime: datetime | None,
    end_datetime: datetime | None,
    stac_api_url: str,
    pystac_open_params: dict[str, Any] | None = None,
    search_params: dict[str, Any] | None = None,
) -> GeoDataFrame[AssetSchema]:
    """Search a generic STAC API and return assets as a GeoDataFrame.

    Args:
        collections: Mapping of collection -> asset keys, or list of collections.
        intersects: AOI geometry as a Shapely object, GeoJSON dict, or path.
        start_datetime: Optional start of temporal window.
        end_datetime: Optional end of temporal window.
        stac_api_url: URL of the STAC API catalog.
        pystac_open_params: Extra arguments forwarded to ``pystac_client.Client.open``.
        search_params: Extra arguments forwarded to ``client.search``.

    Returns:
        A GeoDataFrame of matched assets.

    Raises:
        ValueError: If connection to the STAC API fails or the search query fails.
    """
    collections, collection_asset_filters = build_collection_asset_filters(collections)

    time_range = None
    q_start = None
    q_end = None
    if start_datetime and end_datetime:
        q_start = start_datetime.astimezone(timezone.utc)
        q_end = end_datetime.astimezone(timezone.utc)
        time_range = f"{q_start.strftime(TIME_FORMAT)}/{q_end.strftime(TIME_FORMAT)}"
    elif start_datetime:
        q_start = start_datetime.astimezone(timezone.utc)
        time_range = f"{q_start.strftime(TIME_FORMAT)}/.."
    elif end_datetime:
        q_end = end_datetime.astimezone(timezone.utc)
        time_range = f"../{q_end.strftime(TIME_FORMAT)}"

    client_kwargs: dict[str, Any] = dict(pystac_open_params or {})
    client_kwargs.pop("url", None)
    if "headers" in client_kwargs and isinstance(client_kwargs["headers"], dict):
        client_kwargs["headers"] = {
            str(k): str(v) for k, v in client_kwargs["headers"].items()
        }

    from pystac_client import Client

    try:
        client = Client.open(stac_api_url, **client_kwargs)
    except Exception as e:
        logger.error("Failed to connect to STAC API", url=stac_api_url, error=str(e))
        raise ValueError(f"Failed to connect to STAC API at {stac_api_url}: {e}") from e

    searchkwargs: dict[str, Any] = {}
    if collections:
        searchkwargs["collections"] = collections
    if time_range:
        searchkwargs["datetime"] = time_range

    normalized_intersects = normalize_geometry_input(intersects)
    if normalized_intersects is not None:
        searchkwargs["intersects"] = normalized_intersects.__geo_interface__

    searchkwargs.update(search_params or {})

    try:
        search_req = client.search(**searchkwargs)
        items = list(search_req.items())
    except Exception as e:
        logger.error("STAC search query failed", error=str(e))
        raise ValueError(f"STAC search query failed: {e}") from e

    if not items:
        return empty_asset_result()

    return _parse_stac_items_to_assets(
        items=items,
        collection_asset_filters=collection_asset_filters,
        q_start=q_start,
        q_end=q_end,
    )

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def build_grouped_tasks(
    search_results: GeoDataFrame[AssetSchema],
    job: ExtractionJob,
    cells_per_task: int | None = None,
    buffer_m: float = 0.0,
) -> Sequence[ExtractionTask]:
    """Build extraction tasks from *search_results* using *job* configuration.

    Assets are grouped by ``start_time`` and native ``crs``. For each group,
    the effective AOI is the intersection of the asset footprints with the
    job's ``target_aoi`` (or the asset footprints alone when no target AOI is
    configured). The AOI is tiled with the MajorTOM grid using ``job.grid_dist``
    and consecutive cells are batched into tasks of at most ``cells_per_task``
    cells.

    Args:
        search_results: GeoDataFrame of assets from the search phase.
        job: Parent ``ExtractionJob`` supplying extraction configuration.
        cells_per_task: Maximum number of MajorTOM grid cells per task.
            Defaults to ``None``, which places all grid cells for a group into
            a single task.
        buffer_m: Optional padding in metres around each chunk of grid cells.
            A value such as ``job.grid_dist * 0.1`` is useful when cropping
            assets to ensure edge pixels are included.

    Returns:
        A sequence of ``ExtractionTask`` objects ready for execution.

    Raises:
        ValueError: If ``job.output_uri`` is empty or ``cells_per_task`` is
            zero.
    """
    output_uri = job.output_uri
    if not output_uri:
        raise ValueError("ExtractionJob.output_uri must be a non-empty string.")

    if cells_per_task == 0:
        raise ValueError(
            "cells_per_task must be a positive integer, None, or negative to use all cells."
        )

    if search_results.empty:
        return []

    has_crs = "crs" in search_results.columns
    if has_crs and bool(search_results["crs"].isna().any()):
        raise ValueError(
            "assets['crs'] contains null values. "
            "Either populate crs for all assets or omit the column entirely."
        )

    if not has_crs:
        warn(
            "assets has no 'crs' column; assuming all assets share the same "
            "native CRS. Mixed-CRS assets in one task may fail or produce "
            "incorrect results.",
            UserWarning,
            stacklevel=2,
        )

    group_keys = ["start_time", "crs"] if has_crs else ["start_time"]
    target_aoi = job.effective_target_aoi

    tasks: list[ExtractionTask] = []
    for keys, group in search_results.groupby(group_keys):
        if has_crs:
            start_time, crs = keys  # type: ignore[misc]
        else:
            start_time, crs = keys, None  # type: ignore[assignment]

        group = cast(GeoDataFrame[AssetSchema], group.copy())
        group_aoi = _resolve_group_aoi(group, target_aoi)

        if group_aoi is None or group_aoi.is_empty:
            # No usable geometry/aoi to tile; keep the original job so the task
            # can still be executed with the full asset set.
            task_id = _task_id(job.name, start_time, crs, len(tasks))
            tasks.append(
                ExtractionTask(
                    id=task_id,
                    assets=group,
                    job=job,
                )
            )
            continue

        cells = build_grid_cells(aoi=group_aoi, grid_dist=job.grid_dist)
        chunk_size = (
            len(cells)
            if cells_per_task is None or cells_per_task < 0
            else cells_per_task
        )

        for chunk_index, chunk in enumerate(_chunks(cells, chunk_size)):
            chunk_bounds = cells_bounds(chunk, buffer_m=buffer_m)
            chunk_aoi = shapely.geometry.box(*chunk_bounds)

            task_id = _task_id(
                job.name,
                start_time,
                crs,
                len(tasks),
            )
            tasks.append(
                ExtractionTask(
                    id=task_id,
                    assets=group,
                    job=job,
                    aoi=chunk_aoi,
                    grid_cells=chunk,
                    task_context={
                        "chunk_index": chunk_index,
                    },
                )
            )

    return tasks

@validate_call(config=ConfigDict(arbitrary_types_allowed=True))
def write_geotiff(
    ds: xr.Dataset,
    path: str | Path,
    **kwargs: Any,
) -> str:
    """Write *ds* to a GeoTIFF at *path*.

    All variables are combined into a single multi-band raster. The caller
    (the AEREO orchestrator) is responsible for constructing *path*, splitting
    time dimensions, and building the artifact catalog.

    Args:
        ds: The xarray.Dataset to write. Must not contain a ``time`` dimension;
            the orchestrator calls this function once per time slice.
        path: Destination path to write.
        **kwargs: Keyword arguments forwarded to ``da.rio.to_raster()``.

    Returns:
        The path that was written.

    Raises:
        ValueError: If the dataset contains no data variables or has a
            ``time`` dimension.
    """
    import rioxarray  # noqa: F401

    path = Path(path)
    if "time" in ds.dims:
        raise ValueError(
            "write_geotiff does not accept datasets with a 'time' dimension; "
            "the orchestrator must split time slices before calling the writer."
        )

    combined_da = _dataset_to_raster_bands(ds)
    combined_da.rio.to_raster(str(path), **kwargs)
    return str(path)