"""Encoding utilities for zarr dataset storage."""

import xarray as xr
from zarr.codecs import BloscCodec


def from_ds(
    ds: xr.Dataset,
    store_floats_as_float32: bool = True,
    include_coords: bool = True,
    filter_existing: bool = True,
) -> dict:
    """Create compression encoding for zarr dataset storage.

    Creates Blosc compression configuration for all data variables and coordinates
    in the dataset using zstd compression with level 5.

    Args:
        ds (xr.Dataset): The xarray Dataset to create encoding for.
        store_floats_as_float32 (bool, optional): Whether to store floating point data as float32.
            Defaults to True.
        include_coords (bool, optional): Whether to include coordinates in the encoding.
            This is useful when appending to an existing store.
            Defaults to True.
        filter_existing (bool, optional): Whether to filter out variables that already have
            encoding defined. Useful when appending to existing stores. Defaults to True.

    Returns:
        dict: Encoding dictionary with compression settings for each variable.

    """
    var_names = [*ds.data_vars, *ds.coords] if include_coords else ds.data_vars
    # Filter out variables which already have encoding (e.g., when appending to existing store)
    if filter_existing:
        var_names = [var for var in var_names if not ds[var].encoding]
    encoding = {var: {"compressors": BloscCodec(cname="zstd", clevel=5)} for var in var_names}
    if store_floats_as_float32:
        for var in ds.data_vars:
            if ds[var].dtype == "float64":
                encoding[var]["dtype"] = "float32"
    return encoding


def from_da(da: xr.DataArray, store_floats_as_float32: bool = True) -> dict:
    """Create compression encoding for zarr DataArray storage.

    Creates Blosc compression configuration for the DataArray
    using zstd compression with level 5.

    Args:
        da (xr.DataArray): The xarray DataArray to create encoding for.
        store_floats_as_float32 (bool, optional): Whether to store floating point data as float32.
            Defaults to True.

    Returns:
        dict: Encoding dictionary with compression settings.

    """
    encoding = {"compressors": BloscCodec(cname="zstd", clevel=5)}
    if store_floats_as_float32 and da.dtype == "float64":
        encoding["dtype"] = "float32"
    return {da.name: encoding}
