1"""
   2Kerchunk Reference Manifest Generator
   3======================================
   4
   5Provides :class:`ReferenceGenerator`, which scans one or more GRIB2 files
   6using grib2io's :func:`build_index` infrastructure and produces a
   7`Kerchunk v1 reference manifest <https://fsspec.github.io/kerchunk/spec>`_
   8mapping Zarr chunk keys to ``[url, offset, length]`` tuples within the
   9original files.
  10
  11The manifest can be serialized to JSON or Parquet and later opened with
  12``fsspec.filesystem("reference")`` to create a virtual Zarr store that
  13reads data lazily from the original GRIB2 bytes, decoded on-the-fly by
  14:class:`grib2io.codecs.Grib2Codec`.
  15
  16Example
  17-------
  18>>> from grib2io.kerchunk import ReferenceGenerator
  19>>> gen = ReferenceGenerator("gfs.grib2")
  20>>> manifest = gen.generate()
  21>>> gen.to_json("gfs_refs.json")
  22"""
  23
  24from __future__ import annotations
  25
  26import base64
  27import json
  28import logging
  29import os
  30import re
  31from urllib.parse import urlparse
  32from typing import Any, Dict, List, Optional, Set, Union
  33
  34import numpy as np
  35
  36import grib2io
  37
  38_logger = logging.getLogger(__name__)
  39
  40
  41# ---------------------------------------------------------------------------
  42# Lazy import guards
  43# ---------------------------------------------------------------------------
  44
  45
  46def _ensure_kerchunk():
  47    """Raise ``ImportError`` if *kerchunk* is not available."""
  48    try:
  49        import kerchunk  # noqa: F401
  50    except ImportError:
  51        raise ImportError("kerchunk is required for reference generation. Install with: pip install grib2io[kerchunk]")
  52
  53
  54def _ensure_numcodecs():
  55    """Raise ``ImportError`` if *numcodecs* is not available."""
  56    try:
  57        import numcodecs  # noqa: F401
  58    except ImportError:
  59        raise ImportError("numcodecs is required for the GRIB2 codec. Install with: pip install grib2io[kerchunk]")
  60
  61
  62# ---------------------------------------------------------------------------
  63# Dimension names used for grouping (mirrors xarray_backend logic)
  64# ---------------------------------------------------------------------------
  65
  66# These are the non-geographic dimensions that can appear in GRIB2 data.
  67# The order here determines the dimension order in the Zarr array
  68# (before the trailing y, x spatial dims).
  69_ORDERED_DIM_NAMES = [
  70    "valid_time",
  71    "perturbationNumber",
  72    "duration",
  73    "percentileValue",
  74    "level",
  75]
  76
  77# These dims are always emitted (even at size 1) so that manifests from
  78# different files can be concatenated along them without shape errors.
  79_ALWAYS_INCLUDE_DIMS = frozenset({"valid_time"})
  80
  81# Lazy-loaded level name mapping (typeOfFirstFixedSurface int -> (name, source))
  82_LEVEL_NAME_MAPPING: Optional[dict] = None
  83
  84
  85def _get_level_name_mapping() -> dict:
  86    global _LEVEL_NAME_MAPPING
  87    if _LEVEL_NAME_MAPPING is None:
  88        _LEVEL_NAME_MAPPING = grib2io.tables.get_table("4.5.grib2io.level.name")
  89    return _LEVEL_NAME_MAPPING
  90
  91
  92def _level_dim_name(msg) -> str:
  93    """Return a surface-type-specific level dimension name.
  94
  95    Mirrors the xarray backend's ``swap_dims({"level": key})`` logic so
  96    variables at different surface types get distinct dimension names
  97    (e.g. ``isobaric_surface``, ``height_above_ground``) and can
  98    coexist in a single flat xarray Dataset without conflicting sizes.
  99    """
 100    toffs = getattr(msg, "typeOfFirstFixedSurface", None)
 101    if toffs is None:
 102        return "level"
 103    val = toffs.value if hasattr(toffs, "value") else int(toffs)
 104    entry = _get_level_name_mapping().get(int(val))
 105    if entry:
 106        return entry[0]  # e.g. 'isobaric_surface', 'height_above_ground'
 107    return "level"
 108
 109
 110# ---------------------------------------------------------------------------
 111# Public API
 112# ---------------------------------------------------------------------------
 113
 114
 115class ReferenceGenerator:
 116    """Generate Kerchunk v1 reference manifests from GRIB2 files.
 117
 118    Parameters
 119    ----------
 120    file_paths : str or list of str
 121        One or more GRIB2 file paths (local paths or URIs).
 122    filters : dict, optional
 123        Filter GRIB2 messages by metadata attributes.  Keys can be any
 124        ``Grib2Message`` attribute name (e.g. ``shortName``, ``leadTime``).
 125    storage_options : dict, optional
 126        Extra options passed to ``fsspec.open`` for remote URIs
 127        (e.g. ``{"anon": True}`` for public S3 buckets).
 128    """
 129
 130    def __init__(
 131        self,
 132        file_paths: Union[str, List[str]],
 133        filters: Optional[Dict[str, Any]] = None,
 134        storage_options: Optional[Dict[str, Any]] = None,
 135        max_workers: Optional[int] = None,
 136    ):
 137        _ensure_numcodecs()
 138
 139        if isinstance(file_paths, (str, os.PathLike)):
 140            file_paths = [str(file_paths)]
 141        else:
 142            file_paths = [str(p) for p in file_paths]
 143
 144        # Validate file accessibility.
 145        # Local filesystem paths must exist. URI inputs are handled by
 146        # grib2io.open/fsspec at scan time and should not be rejected here.
 147        for fp in file_paths:
 148            if _is_local_path(fp) and not os.path.isfile(fp):
 149                raise FileNotFoundError(f"GRIB2 file not found: {fp}")
 150
 151        self.file_paths = file_paths
 152        self.filters = filters or {}
 153        self.storage_options = storage_options or {}
 154        self.max_workers = max_workers
 155        self._manifest: Optional[dict] = None
 156
 157    def generate(self) -> dict:
 158        """Scan files and produce a Kerchunk v1 reference manifest.
 159
 160        Returns
 161        -------
 162        dict
 163            Kerchunk reference spec v1 dict with keys ``"version"`` and
 164            ``"refs"``.
 165        """
 166        refs: Dict[str, Any] = {}
 167
 168        # .zgroup at root
 169        refs[".zgroup"] = json.dumps({"zarr_format": 2})
 170
 171        # Collect all messages across files, keyed by variable group
 172        # group_key = (shortName, typeOfFirstFixedSurface, pdtn, typeOfSecondFixedSurface)
 173        # This ensures messages with different surface types are not mixed.
 174        all_var_messages: Dict[tuple, list] = {}
 175
 176        n_files = len(self.file_paths)
 177        use_parallel = self.max_workers != 1 and n_files > 1 and not _is_local_path(self.file_paths[0])
 178
 179        if use_parallel:
 180            import concurrent.futures
 181
 182            workers = self.max_workers or min(n_files, 8)
 183
 184            def _scan_one(file_path):
 185                file_uri = _file_uri(file_path)
 186                local_msgs: Dict[tuple, list] = {}
 187                self._scan_file(file_path, file_uri, local_msgs)
 188                return local_msgs
 189
 190            with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
 191                futures = {pool.submit(_scan_one, fp): fp for fp in self.file_paths}
 192                for future in concurrent.futures.as_completed(futures):
 193                    fp = futures[future]
 194                    try:
 195                        local_msgs = future.result()
 196                        for key, entries in local_msgs.items():
 197                            all_var_messages.setdefault(key, []).extend(entries)
 198                    except Exception as e:
 199                        raise ValueError(f"Failed to parse GRIB2 file '{fp}': {e}") from e
 200        else:
 201            for file_path in self.file_paths:
 202                file_uri = _file_uri(file_path)
 203                try:
 204                    self._scan_file(file_path, file_uri, all_var_messages)
 205                except Exception as e:
 206                    raise ValueError(f"Failed to parse GRIB2 file '{file_path}': {e}") from e
 207
 208        # For each variable group, map messages to dimensions and build refs.
 209        # Track used variable names to handle collisions (same shortName
 210        # but different surface types).
 211        # Also track level coord name -> level values so that the same surface
 212        # type but different level extents get disambiguated names (mirrors
 213        # how the xarray backend keeps each surface type in its own Dataset).
 214        used_var_names: Dict[str, int] = {}
 215        level_coord_registry: Dict[str, list] = {}  # name -> sorted level values
 216        for group_key, msg_entries in all_var_messages.items():
 217            var_name = group_key[0]  # shortName is the first element
 218            if var_name in used_var_names:
 219                used_var_names[var_name] += 1
 220                zarr_var_name = f"{var_name}_{used_var_names[var_name]}"
 221            else:
 222                used_var_names[var_name] = 0
 223                zarr_var_name = var_name
 224            self._build_variable_refs(zarr_var_name, msg_entries, refs, level_coord_registry)
 225
 226        # Build latitude/longitude coordinate arrays from the grid definition.
 227        # All messages are assumed to share the same grid (required by the
 228        # xarray backend too), so we use the first available message.
 229        if all_var_messages:
 230            first_entries = next(iter(all_var_messages.values()))
 231            rep_msg = first_entries[0].msg
 232            _build_latlon_coord_refs(rep_msg, refs)
 233
 234        self._manifest = {"version": 1, "refs": refs}
 235        return self._manifest
 236
 237    def to_json(self, output_path: str) -> None:
 238        """Serialize the manifest to a JSON file.
 239
 240        Parameters
 241        ----------
 242        output_path : str
 243            Path to the output JSON file.
 244        """
 245        if self._manifest is None:
 246            self.generate()
 247        with open(output_path, "w") as f:
 248            json.dump(self._manifest, f)
 249
 250    def to_parquet(self, output_path: str) -> None:
 251        """Serialize the manifest to a Parquet reference store.
 252
 253        Parameters
 254        ----------
 255        output_path : str
 256            Path to the output Parquet directory.
 257        """
 258        _ensure_kerchunk()
 259        if self._manifest is None:
 260            self.generate()
 261
 262        import fsspec
 263        from fsspec.implementations.reference import LazyReferenceMapper
 264
 265        fs, _ = fsspec.core.url_to_fs(output_path)
 266        out = LazyReferenceMapper.create(output_path, fs=fs, record_size=100_000, engine="pyarrow")
 267        refs = self._manifest.get("refs", self._manifest)
 268        for k in sorted(refs):
 269            out[k] = refs[k]
 270        out.flush()
 271
 272    # ------------------------------------------------------------------
 273    # Internal scanning
 274    # ------------------------------------------------------------------
 275
 276    def _build_remote_index_filtered(
 277        self,
 278        file_path: str,
 279        shortname_filter: Optional[Union[str, Set[str]]] = None,
 280        scan_storage_options: Optional[dict] = None,
 281    ):
 282        """Build a GRIB2 index for a remote file using sidecar pre-filtering.
 283
 284        Works with any combination of filters: a ``shortName`` alone,
 285        ``shortName`` plus additional filters (e.g. ``typeOfFirstFixedSurface``
 286        / ``level``), or filters without a ``shortName`` at all.
 287
 288        Instead of fetching headers for every message (~700 HTTP requests for a
 289        full GFS 0.25° file), this method:
 290
 291        1. Checks grib2io's local cache – if the full or filtered index was
 292           saved from a previous run, it loads it instantly.
 293        2. Checks for a remote grib2io ``.grib2ioidx`` sidecar (binary index)
 294           alongside the GRIB2 file — the most efficient format, containing
 295           pre-parsed section offsets and avoiding header reads entirely.
 296        3. Fetches the wgrib2 ``.idx`` text sidecar and keeps only the byte
 297           offsets whose shortName matches the filter, reducing HTTP range
 298           requests from ~700 to ~1–50.
 299        4. Saves the partial index to a filter-specific cache key so the next
 300           call for the same file+filter is also instant.
 301        5. Falls back to ``grib2io.open`` (full index, slow on first call) if
 302           no sidecar is available.
 303
 304        Returns
 305        -------
 306        tuple(dict, list)
 307            ``(index, msgs)`` where *index* is a grib2io index dict and *msgs*
 308            is a list of :class:`~grib2io.Grib2Message` objects.
 309        """
 310        import builtins
 311        import hashlib
 312        import pickle
 313
 314        import fsspec
 315
 316        from grib2io._grib2io import build_index
 317        from grib2io import msgs_from_index
 318
 319        scan_storage_options = scan_storage_options or {}
 320        cache_root = os.path.join(os.path.expanduser("~"), ".cache", "grib2io")
 321
 322        # Open the remote file to obtain its size (one lightweight HEAD/info call).
 323        # For the filtered fast-path we override cache settings: "readahead" with
 324        # a small block size means consecutive section-header reads within the same
 325        # message share one HTTP range request instead of each triggering a new one.
 326        fh_options = dict(scan_storage_options)
 327        fh_options["default_cache_type"] = "readahead"
 328        fh_options["default_block_size"] = 4096
 329        fh = fsspec.open(file_path, "rb", **fh_options).open()
 330        try:
 331            size = int(fh.info().get("size", 0) or 0)
 332        except Exception:
 333            size = 0
 334
 335        # ------------------------------------------------------------------ #
 336        # 1. Full-index cache (populated by unfiltered grib2io.open calls)    #
 337        # ------------------------------------------------------------------ #
 338        full_cache_key = hashlib.sha1((file_path + str(size)).encode("ASCII")).hexdigest()
 339        full_cache_path = os.path.join(cache_root, f"{full_cache_key}.grib2ioidx")
 340        if os.path.exists(full_cache_path):
 341            with builtins.open(full_cache_path, "rb") as cf:
 342                index = pickle.load(cf)
 343            msgs = msgs_from_index(index, filehandle=fh)
 344            fh.close()
 345            return index, msgs
 346
 347        # ------------------------------------------------------------------ #
 348        # 2. Filter-specific partial-index cache                              #
 349        # ------------------------------------------------------------------ #
 350        filter_repr = ":".join(f"{k}={v}" for k, v in sorted(self.filters.items()))
 351        filtered_cache_key = hashlib.sha1((file_path + str(size) + ":" + filter_repr).encode("ASCII")).hexdigest()
 352        filtered_cache_path = os.path.join(cache_root, f"{filtered_cache_key}.grib2ioidx")
 353        if os.path.exists(filtered_cache_path):
 354            with builtins.open(filtered_cache_path, "rb") as cf:
 355                index = pickle.load(cf)
 356            msgs = msgs_from_index(index, filehandle=fh)
 357            fh.close()
 358            return index, msgs
 359
 360        # ------------------------------------------------------------------ #
 361        # 3. Remote grib2io index sidecar (.grib2ioidx)                       #
 362        # ------------------------------------------------------------------ #
 363        # grib2io publishes its own binary index alongside the GRIB2 file.
 364        # This is the most efficient index format — it contains the full
 365        # parsed section offsets/sizes and avoids any header reads.  Check
 366        # for it before falling back to the wgrib2 text .idx sidecar.
 367        grib2io_idx_url = file_path + ".grib2ioidx"
 368        try:
 369            with fsspec.open(grib2io_idx_url, "rb", **scan_storage_options) as gf:
 370                index = pickle.load(gf)
 371            msgs = msgs_from_index(index, filehandle=fh)
 372            fh.close()
 373            # Cache locally so subsequent calls are instant.
 374            try:
 375                os.makedirs(cache_root, exist_ok=True)
 376                with builtins.open(full_cache_path, "wb") as cf:
 377                    pickle.dump(index, cf)
 378            except Exception:
 379                pass
 380            return index, msgs
 381        except Exception:
 382            pass
 383
 384        # ------------------------------------------------------------------ #
 385        # 4. wgrib2 .idx sidecar pre-filtering                                #
 386        # ------------------------------------------------------------------ #
 387        idx_url = file_path + ".idx"
 388        idx_fetch_ok = False
 389        filtered_offsets: List[int] = []
 390        try:
 391            with fsspec.open(idx_url, "r", **scan_storage_options) as idxf:
 392                filtered_offsets = _prefilter_idx_offsets(idxf, shortname_filter, self.filters)
 393            idx_fetch_ok = True
 394        except Exception:
 395            pass
 396
 397        if idx_fetch_ok:
 398            if not filtered_offsets:
 399                # shortName simply does not appear in this file.
 400                fh.close()
 401                return {}, []
 402            index = build_index(fh, offsets=filtered_offsets)
 403            msgs = msgs_from_index(index, filehandle=fh)
 404            fh.close()
 405            # Persist the partial index so the next call is instant.
 406            try:
 407                os.makedirs(cache_root, exist_ok=True)
 408                with builtins.open(filtered_cache_path, "wb") as cf:
 409                    pickle.dump(index, cf)
 410            except Exception:
 411                pass
 412            return index, msgs
 413
 414        # ------------------------------------------------------------------ #
 415        # 5. Fall back: let grib2io.open build the full index (slow on first  #
 416        #    call, but saves to grib2io's own cache for future calls).         #
 417        # ------------------------------------------------------------------ #
 418        fh.close()
 419        with grib2io.open(file_path, save_index=True, use_index=True, **scan_storage_options) as f:
 420            return f._index, list(f)
 421
 422    def _scan_file(
 423        self,
 424        file_path: str,
 425        file_uri: str,
 426        all_var_messages: Dict[str, list],
 427    ) -> None:
 428        """Scan a single GRIB2 file and collect message entries."""
 429        if _is_local_path(file_path):
 430            with grib2io.open(file_path, save_index=False, use_index=True) as f:
 431                index = f._index
 432                msgs = list(f)
 433        else:
 434            scan_storage_options = _remote_scan_storage_options(file_path, self.storage_options)
 435            shortname_filter = self.filters.get("shortName") if self.filters else None
 436            # Normalise list/tuple to a set so _prefilter_idx_offsets can do a
 437            # fast membership test; leave scalar strings as-is.
 438            if isinstance(shortname_filter, (list, tuple)):
 439                shortname_filter = set(shortname_filter)
 440            # Fast path: resolve the index from a sidecar (.grib2ioidx or the
 441            # wgrib2 .idx) instead of fetching headers for every message.  This
 442            # works whether or not a shortName filter is given: when shortName
 443            # is present it is combined with any other filters; when it is
 444            # absent, other filters (e.g. typeOfFirstFixedSurface/level) are
 445            # still applied to the .idx, and the .grib2ioidx sidecar yields the
 446            # full parsed index with no header reads at all.
 447            index, msgs = self._build_remote_index_filtered(file_path, shortname_filter, scan_storage_options)
 448
 449        n_msgs = len(msgs)
 450        for i in range(n_msgs):
 451            msg = msgs[i]
 452
 453            # Apply filters
 454            if not self._matches_filters(msg):
 455                continue
 456
 457            sec_offsets = index["sectionOffset"][i]
 458            sec_sizes = index["sectionSize"][i]
 459            bmapflag = index["bmapflag"][i]
 460
 461            # Section 7 offset and length
 462            sec7_offset = sec_offsets[7]
 463            sec7_length = sec_sizes[7]
 464
 465            # Section 5 (data representation) — always present
 466            sec5_offset = sec_offsets[5]
 467            sec5_length = sec_sizes[5]
 468
 469            # Section 6 (bitmap) — always present (6 bytes when no bitmap)
 470            sec6_length = sec_sizes[6]
 471            # Offset only needed for the old bitmap-conditional path (kept for reference)
 472            sec6_offset = sec_offsets[6] if bmapflag in {0, 254} else None
 473
 474            # Build a composite variable key that includes the surface type
 475            # to avoid grouping messages with different surface types together.
 476            # This mirrors how the xarray backend requires filtering to a
 477            # single typeOfFirstFixedSurface.
 478            var_name = str(msg.shortName)
 479            type_of_first_fixed_surface = msg.typeOfFirstFixedSurface
 480            if hasattr(type_of_first_fixed_surface, "value"):
 481                toffs_val = type_of_first_fixed_surface.value
 482            else:
 483                toffs_val = type_of_first_fixed_surface
 484
 485            # Also include typeOfGeneratingProcess and
 486            # productDefinitionTemplateNumber to disambiguate further
 487            # (same approach as xarray backend's required_uniques)
 488            pdtn = msg.productDefinitionTemplateNumber
 489            if hasattr(pdtn, "value"):
 490                pdtn_val = pdtn.value
 491            else:
 492                pdtn_val = pdtn
 493
 494            type_of_second_fixed_surface = msg.typeOfSecondFixedSurface
 495            if hasattr(type_of_second_fixed_surface, "value"):
 496                tosfs_val = type_of_second_fixed_surface.value
 497            else:
 498                tosfs_val = type_of_second_fixed_surface
 499
 500            # Group key: shortName + surface type + pdtn + second surface type
 501            group_key = (var_name, int(toffs_val), int(pdtn_val), int(tosfs_val))
 502
 503            entry = _MsgEntry(
 504                msg=msg,
 505                file_uri=file_uri,
 506                sec5_offset=sec5_offset,
 507                sec5_length=sec5_length,
 508                sec6_length=sec6_length,
 509                sec7_offset=sec7_offset,
 510                sec7_length=sec7_length,
 511                sec6_offset=sec6_offset,
 512                bmapflag=bmapflag,
 513                index_section3=index["section3"][i],
 514                index_section5=index["section5"][i],
 515            )
 516
 517            all_var_messages.setdefault(group_key, []).append(entry)
 518
 519    def _matches_filters(self, msg) -> bool:
 520        """Check if a message matches all user-supplied filters.
 521
 522        Filter values may be:
 523
 524        * **scalar** – exact equality (``{"shortName": "TMP"}``)
 525        * **list / tuple / set** – membership test
 526          (``{"level": [500, 850, 250]}``)
 527        * **slice** – inclusive range test
 528          (``{"level": slice(500, 850)}``)
 529        """
 530        for key, value in self.filters.items():
 531            msg_val = getattr(msg, key, None)
 532            if msg_val is None:
 533                return False
 534            # Unwrap Grib2Metadata wrapper objects
 535            if hasattr(msg_val, "value"):
 536                msg_val = msg_val.value
 537            if isinstance(value, slice):
 538                lo = value.start if value.start is not None else float("-inf")
 539                hi = value.stop if value.stop is not None else float("inf")
 540                try:
 541                    if not (lo <= msg_val <= hi):
 542                        return False
 543                except TypeError:
 544                    return False
 545            elif isinstance(value, (list, tuple, set)):
 546                if msg_val not in value:
 547                    return False
 548            else:
 549                if msg_val != value:
 550                    return False
 551        return True
 552
 553    # ------------------------------------------------------------------
 554    # Variable reference building
 555    # ------------------------------------------------------------------
 556
 557    def _build_variable_refs(
 558        self,
 559        var_name: str,
 560        msg_entries: list,
 561        refs: Dict[str, Any],
 562        level_coord_registry: Optional[Dict[str, list]] = None,
 563    ) -> None:
 564        """Build all Zarr refs for a single variable."""
 565        # Derive surface-type-specific level dim name (mirrors xarray_backend)
 566        level_name = _level_dim_name(msg_entries[0].msg)
 567        # Map messages to dimensions
 568        dim_mapping = _map_messages_to_dimensions(msg_entries, level_dim_name=level_name)
 569
 570        # Disambiguate the level dim name if this surface type already appears
 571        # in the manifest with different level values.  This prevents xarray
 572        # from seeing conflicting dimension sizes when variables at the same
 573        # surface type have different level counts.
 574        if level_coord_registry is not None and level_name in dim_mapping["dim_values"]:
 575            level_vals = dim_mapping["dim_values"][level_name]
 576            if level_name in level_coord_registry:
 577                if level_coord_registry[level_name] != level_vals:
 578                    # Same surface type, different level values — append suffix
 579                    suffix = 2
 580                    candidate = f"{level_name}_{suffix}"
 581                    while candidate in level_coord_registry and level_coord_registry[candidate] != level_vals:
 582                        suffix += 1
 583                        candidate = f"{level_name}_{suffix}"
 584                    level_name = candidate
 585                    dim_mapping = _map_messages_to_dimensions(msg_entries, level_dim_name=level_name)
 586            if level_name not in level_coord_registry:
 587                level_coord_registry[level_name] = dim_mapping["dim_values"].get(level_name, [])
 588
 589        dim_names = dim_mapping["dim_names"]  # ordered list of dim names
 590        dim_values = dim_mapping["dim_values"]  # dict: dim_name -> sorted unique values
 591        msg_index_map = dim_mapping["msg_index_map"]  # dict: dim_tuple -> msg_entry index
 592
 593        # Representative message for metadata
 594        rep_msg = msg_entries[0].msg
 595
 596        # Compute shape (ensure plain Python ints for JSON serialization)
 597        shape = [len(dim_values[d]) for d in dim_names] + [int(rep_msg.ny), int(rep_msg.nx)]
 598        chunks = [1] * len(dim_names) + [int(rep_msg.ny), int(rep_msg.nx)]
 599
 600        # Build .zarray
 601        codec_config = _build_codec_config(msg_entries[0])
 602        zarray = _build_zarray_metadata(rep_msg, shape, chunks, codec_config)
 603        refs[f"{var_name}/.zarray"] = json.dumps(zarray)
 604
 605        # Build .zattrs
 606        dim_labels = dim_names + ["y", "x"]
 607        zattrs = _build_zattrs(rep_msg, dim_labels)
 608        refs[f"{var_name}/.zattrs"] = json.dumps(zattrs)
 609
 610        # Build data chunk refs
 611        for dim_tuple, entry_idx in msg_index_map.items():
 612            entry = msg_entries[entry_idx]
 613            dim_indices = []
 614            for i, d in enumerate(dim_names):
 615                val = dim_tuple[i]
 616                idx = list(dim_values[d]).index(val)
 617                dim_indices.append(idx)
 618
 619            chunk_key = _build_chunk_key(var_name, dim_indices)
 620
 621            # Store a combined reference covering sections 5+6+7 so that the
 622            # codec can parse the per-chunk data representation template (sec5)
 623            # and bitmap (sec6) dynamically.  Sections 5, 6, and 7 are always
 624            # contiguous in the GRIB2 byte stream.
 625            refs[chunk_key] = [
 626                entry.file_uri,
 627                entry.sec5_offset,
 628                entry.sec5_length + entry.sec6_length + entry.sec7_length,
 629            ]
 630
 631        # Build coordinate arrays as inline base64-encoded refs
 632        for dim_name in dim_names:
 633            values = dim_values[dim_name]
 634            _build_coord_refs(dim_name, values, refs)
 635
 636    # ------------------------------------------------------------------
 637    # Manifest access
 638    # ------------------------------------------------------------------
 639
 640    @property
 641    def manifest(self) -> Optional[dict]:
 642        """The generated manifest, or ``None`` if :meth:`generate` has not
 643        been called yet."""
 644        return self._manifest
 645
 646
 647# ---------------------------------------------------------------------------
 648# Internal data class for message entries
 649# ---------------------------------------------------------------------------
 650
 651
 652class _MsgEntry:
 653    """Lightweight container for a scanned GRIB2 message."""
 654
 655    __slots__ = (
 656        "msg",
 657        "file_uri",
 658        "sec5_offset",
 659        "sec5_length",
 660        "sec6_length",
 661        "sec7_offset",
 662        "sec7_length",
 663        "sec6_offset",
 664        "bmapflag",
 665        "index_section3",
 666        "index_section5",
 667    )
 668
 669    def __init__(
 670        self,
 671        msg,
 672        file_uri: str,
 673        sec5_offset: int,
 674        sec5_length: int,
 675        sec6_length: int,
 676        sec7_offset: int,
 677        sec7_length: int,
 678        sec6_offset: Optional[int],
 679        bmapflag: int,
 680        index_section3: np.ndarray,
 681        index_section5: np.ndarray,
 682    ):
 683        self.msg = msg
 684        self.file_uri = file_uri
 685        self.sec5_offset = sec5_offset
 686        self.sec5_length = sec5_length
 687        self.sec6_length = sec6_length
 688        self.sec7_offset = sec7_offset
 689        self.sec7_length = sec7_length
 690        self.sec6_offset = sec6_offset
 691        self.bmapflag = bmapflag
 692        self.index_section3 = index_section3
 693        self.index_section5 = index_section5
 694
 695
 696# ---------------------------------------------------------------------------
 697# Helper functions
 698# ---------------------------------------------------------------------------
 699
 700
 701def _file_uri(file_path: str) -> str:
 702    """Convert a local file path to ``file://`` URI, preserving URI inputs."""
 703    if not _is_local_path(file_path):
 704        return file_path
 705    abs_path = os.path.abspath(file_path)
 706    return f"file://{abs_path}"
 707
 708
 709def _is_local_path(path: str) -> bool:
 710    """Return ``True`` if *path* looks like a local filesystem path."""
 711    parsed = urlparse(path)
 712    return parsed.scheme == ""
 713
 714
 715def _remote_scan_storage_options(file_path: str, storage_options: Dict[str, Any]) -> Dict[str, Any]:
 716    """Build tuned fsspec options for remote metadata scans.
 717
 718    These defaults reduce accidental large-block downloads while scanning
 719    message headers/indices across large remote GRIB2 objects.
 720    """
 721    tuned = {
 722        "default_fill_cache": False,
 723        "default_cache_type": "none",
 724        "default_block_size": 131072,
 725    }
 726
 727    # Public S3 GRIB2 archives are common; default to anonymous access
 728    # unless the caller explicitly requests credentialed access.
 729    if urlparse(file_path).scheme in {"s3", "s3a"} and "anon" not in storage_options:
 730        tuned["anon"] = True
 731
 732    tuned.update(storage_options)
 733    return tuned
 734
 735
 736def _value_matches(val: float, filter_val: Any) -> bool:
 737    """Return True if *val* satisfies the scalar / list / slice *filter_val*."""
 738    if isinstance(filter_val, (list, tuple, set)):
 739        return val in filter_val or val in {float(v) for v in filter_val}
 740    if isinstance(filter_val, slice):
 741        lo = float(filter_val.start) if filter_val.start is not None else float("-inf")
 742        hi = float(filter_val.stop) if filter_val.stop is not None else float("inf")
 743        return lo <= val <= hi
 744    return val == filter_val or val == float(filter_val)
 745
 746
 747# Mapping from GRIB2 Table 4.5 typeOfFirstFixedSurface to wgrib2 .idx level
 748# string prefixes for single-valued surfaces (no numeric level component).
 749_TOFS_FIXED_STRINGS: Dict[int, tuple] = {
 750    1: ("surface", "ground or water surface"),
 751    6: ("max wind",),
 752    7: ("tropopause",),
 753    8: ("top of atmosphere", "nominal top of atmosphere"),
 754    10: ("entire atmosphere",),
 755    101: ("mean sea level",),
 756}
 757
 758
 759def _idx_level_matches(level_str: str, tofs: Any, level_filter: Any) -> bool:
 760    """Return True if the wgrib2 ``.idx`` level string is consistent with filters.
 761
 762    Conservative: returns True when the level type is unrecognised so that
 763    false negatives (silently skipping a matching message) are avoided.
 764
 765    Recognised mappings:
 766
 767    * ``typeOfFirstFixedSurface=103`` (height above ground, m):
 768      ``"2 m above ground"``
 769    * ``typeOfFirstFixedSurface=100`` (isobaric surface, Pa):
 770      ``"500 mb"`` — grib2io returns Pa so ``500 mb`` → ``level=50000``.
 771      Both hPa and Pa values are tried to handle version differences.
 772    * Fixed-label surfaces (1, 6, 7, 10, 11, 101): matched by keyword.
 773    """
 774    if tofs is None and level_filter is None:
 775        return True
 776
 777    # Height above ground in metres (tofs = 103)
 778    if tofs == 103:
 779        m = re.match(r"^(\d+(?:\.\d+)?)\s+m\s+above\s+ground", level_str)
 780        if not m:
 781            return False
 782        return level_filter is None or _value_matches(float(m.group(1)), level_filter)
 783
 784    # Isobaric surface in Pa (tofs = 100); wgrib2 uses hPa ("mb")
 785    if tofs == 100:
 786        m = re.match(r"^(\d+(?:\.\d+)?)\s+mb", level_str)
 787        if not m:
 788            return False
 789        if level_filter is None:
 790            return True
 791        idx_mb = float(m.group(1))
 792        # grib2io returns Pa (500 mb → 50000); accept both Pa and hPa
 793        return _value_matches(idx_mb * 100, level_filter) or _value_matches(idx_mb, level_filter)
 794
 795    # Fixed-label surfaces with no numeric level component
 796    if tofs in _TOFS_FIXED_STRINGS:
 797        ls = level_str.lower()
 798        return any(ls.startswith(s) for s in _TOFS_FIXED_STRINGS[tofs])
 799
 800    # Unknown surface type — keep conservatively
 801    return True
 802
 803
 804def _prefilter_idx_offsets(
 805    filehandle,
 806    shortname: Optional[Union[str, Set[str]]] = None,
 807    filters: Optional[Dict[str, Any]] = None,
 808) -> List[int]:
 809    """Parse a wgrib2 ``.idx`` sidecar and return byte offsets matching filters.
 810
 811    The wgrib2 ``.idx`` line format is::
 812
 813        MSG_NUM:BYTE_OFFSET:d=YYYYMMDDCC:SHORTNAME:LEVEL:FORECAST:
 814
 815    *shortname* may be a single string, a set/list of strings, or ``None``.
 816    When ``None`` the shortName is not constrained and every message is kept
 817    unless another filter rules it out.  When *filters* is provided, the level
 818    string (``parts[4]``) is also checked against ``typeOfFirstFixedSurface``
 819    and ``level`` entries in *filters*, which can drastically reduce the number
 820    of messages passed to :func:`build_index` (e.g. from ~50 TMP pressure
 821    levels to 1 for T2M).
 822
 823    Returns an empty list if the sidecar cannot be parsed or contains no
 824    matching messages.
 825    """
 826    names: Optional[Set[str]] = None
 827    if shortname is not None:
 828        names = {shortname} if isinstance(shortname, str) else set(shortname)
 829    tofs = filters.get("typeOfFirstFixedSurface") if filters else None
 830    level_filter = filters.get("level") if filters else None
 831    offsets: List[int] = []
 832    for line in filehandle:
 833        if isinstance(line, bytes):
 834            line = line.decode("utf-8", errors="replace")
 835        parts = line.split(":")
 836        if len(parts) >= 5 and (names is None or parts[3] in names):
 837            # Level-string pre-filter: skip if we can definitively rule out a
 838            # match from the .idx level description (e.g. "500 mb" vs "2 m
 839            # above ground").  Conservative: unknown formats are kept.
 840            if not _idx_level_matches(parts[4], tofs, level_filter):
 841                continue
 842            try:
 843                offsets.append(int(parts[1]))
 844            except ValueError:
 845                continue
 846    return offsets
 847
 848
 849def _build_chunk_key(var_name: str, dim_indices: List[int]) -> str:
 850    """Construct a Zarr chunk key like ``"TMP/0.0.0"``.
 851
 852    Parameters
 853    ----------
 854    var_name : str
 855        Variable name (top-level Zarr array name).
 856    dim_indices : list of int
 857        Integer indices along each non-spatial dimension.
 858
 859    Returns
 860    -------
 861    str
 862        Zarr chunk key, e.g. ``"TMP/0.1.0.0"`` where the trailing
 863        two zeros are for the y and x spatial dimensions (always 0
 864        since each message is one full grid).
 865    """
 866    parts = [str(i) for i in dim_indices] + ["0", "0"]
 867    return f"{var_name}/{'.'.join(parts)}"
 868
 869
 870def _build_zarray_metadata(
 871    msg,
 872    shape: List[int],
 873    chunks: List[int],
 874    codec_config: dict,
 875) -> dict:
 876    """Build ``.zarray`` JSON metadata for a variable.
 877
 878    Parameters
 879    ----------
 880    msg : Grib2Message
 881        Representative message for dtype info.
 882    shape : list of int
 883        Full array shape including spatial dims.
 884    chunks : list of int
 885        Chunk shape (one message per chunk).
 886    codec_config : dict
 887        ``Grib2Codec`` configuration dict.
 888
 889    Returns
 890    -------
 891    dict
 892        Zarr ``.zarray`` metadata.
 893    """
 894    dtype = "<f4" if msg.typeOfValues == 0 else "<i4"
 895
 896    # Place the codec config in `filters` (a list) rather than `compressor`
 897    # (a single dict). VirtualiZarr v2's translator iterates the `compressor`
 898    # field directly, which unpacks a dict's keys instead of the dict itself.
 899    # Using `filters: [codec_config]` with `compressor: null` is handled
 900    # correctly by both VirtualiZarr and zarr v2/numcodecs via fsspec.
 901    return {
 902        "zarr_format": 2,
 903        "shape": shape,
 904        "chunks": chunks,
 905        "dtype": dtype,
 906        "fill_value": "NaN" if dtype == "<f4" else 0,
 907        "order": "C",
 908        "compressor": None,
 909        "filters": [codec_config],
 910    }
 911
 912
 913def _build_zattrs(msg, dim_labels: List[str]) -> dict:
 914    """Extract GRIB2 section metadata as Zarr attributes.
 915
 916    Parameters
 917    ----------
 918    msg : Grib2Message
 919        Representative message.
 920    dim_labels : list of str
 921        Ordered dimension names including ``"y"`` and ``"x"``.
 922
 923    Returns
 924    -------
 925    dict
 926        Zarr ``.zattrs`` metadata.
 927    """
 928    # Extract typeOfFirstFixedSurface - handle Grib2Metadata objects
 929    type_of_first_fixed_surface = msg.typeOfFirstFixedSurface
 930    if hasattr(type_of_first_fixed_surface, "value"):
 931        type_of_first_fixed_surface = type_of_first_fixed_surface.value
 932
 933    # Extract valueOfFirstFixedSurface
 934    value_of_first_fixed_surface = msg.valueOfFirstFixedSurface
 935    if hasattr(value_of_first_fixed_surface, "value"):
 936        value_of_first_fixed_surface = value_of_first_fixed_surface.value
 937
 938    # Extract valid_time (= refDate + leadTime = msg.validDate)
 939    vt = getattr(msg, "validDate", None)
 940    if vt is None:
 941        try:
 942            vt = msg.refDate + msg.leadTime
 943        except Exception:
 944            vt = msg.refDate
 945    if hasattr(vt, "isoformat"):
 946        valid_time_str = vt.isoformat()
 947    elif isinstance(vt, np.datetime64):
 948        valid_time_str = str(vt)
 949    else:
 950        valid_time_str = str(vt)
 951
 952    return {
 953        "_ARRAY_DIMENSIONS": dim_labels,
 954        "coordinates": "latitude longitude",
 955        "discipline": int(msg.section0[2]),
 956        "parameterCategory": int(msg.parameterCategory),
 957        "parameterNumber": int(msg.parameterNumber),
 958        "typeOfFirstFixedSurface": int(type_of_first_fixed_surface),
 959        "valueOfFirstFixedSurface": float(value_of_first_fixed_surface),
 960        "valid_time": valid_time_str,
 961        "shortName": str(msg.shortName),
 962        "fullName": str(msg.fullName),
 963        "units": str(msg.units),
 964    }
 965
 966
 967def _build_codec_config(entry: _MsgEntry) -> dict:
 968    """Build ``Grib2Codec`` configuration from a message entry.
 969
 970    Parameters
 971    ----------
 972    entry : _MsgEntry
 973        Scanned message entry with index metadata.
 974
 975    Returns
 976    -------
 977    dict
 978        Codec configuration suitable for ``Grib2Codec.from_config()``.
 979    """
 980    msg = entry.msg
 981    sec3 = entry.index_section3
 982    sec5 = entry.index_section5
 983
 984    # GDS: first 5 elements of section 3
 985    gds = [int(x) for x in sec3[:5]]
 986    # GDT: remaining elements of section 3
 987    gdt = [int(x) for x in sec3[5:]]
 988    # DRT number and template
 989    drtn = int(sec5[1])
 990    drt = [int(x) for x in sec5[2:]]
 991
 992    # Grid dimensions
 993    nx = int(msg.nx)
 994    ny = int(msg.ny)
 995
 996    # Scan mode flags
 997    scan_mode_flags = None
 998    if hasattr(msg, "scanModeFlags"):
 999        scan_mode_flags = [int(x) for x in msg.scanModeFlags]
1000
1001    # Bitmap info
1002    bitmap_flag = int(entry.bmapflag)
1003    bitmap_offset = None
1004    bitmap_length = None
1005    if bitmap_flag in {0, 254} and entry.sec6_offset is not None:
1006        bitmap_offset = int(entry.sec6_offset)
1007        bitmap_length = int(entry.sec6_length)
1008
1009    # Number of data points and packed values
1010    number_of_data_points = int(msg.numberOfDataPoints)
1011    number_of_packed_values = int(msg.numberOfPackedValues)
1012
1013    # Type of values
1014    type_of_values = int(msg.typeOfValues) if hasattr(msg, "typeOfValues") else 0
1015
1016    # Emit a Zarr v2/v3-compatible codec config dict for VirtualiZarr compatibility.
1017    config = {
1018        "id": "grib2io",
1019        "drtn": drtn,
1020        "drt": drt,
1021        "gdtn": int(msg.gdtn),
1022        "gdt": gdt,
1023        "gds": gds,
1024        "nx": nx,
1025        "ny": ny,
1026        "bitmap_flag": bitmap_flag,
1027        "bitmap_offset": bitmap_offset,
1028        "bitmap_length": bitmap_length,
1029        "scan_mode_flags": scan_mode_flags,
1030        "type_of_values": type_of_values,
1031        "number_of_data_points": number_of_data_points,
1032        "number_of_packed_values": number_of_packed_values,
1033    }
1034    # For VirtualiZarr, the 'compressor' field must be a dict, not a string or id.
1035    return config
1036
1037
1038def _get_dim_value(msg, dim_name: str) -> Any:
1039    """Extract a dimension coordinate value from a message.
1040
1041    Parameters
1042    ----------
1043    msg : Grib2Message
1044        The GRIB2 message.
1045    dim_name : str
1046        Dimension name.
1047
1048    Returns
1049    -------
1050    Any
1051        The coordinate value, converted to a hashable/sortable type.
1052    """
1053    if dim_name == "level":
1054        # Use the tuple (valueOfFirstFixedSurface, valueOfSecondFixedSurface)
1055        # as the level identifier, matching xarray_backend logic
1056        v1 = msg.valueOfFirstFixedSurface
1057        v2 = msg.valueOfSecondFixedSurface
1058        return (float(v1), float(v2))
1059    elif dim_name == "valid_time":
1060        # valid_time = refDate + leadTime (i.e. msg.validDate)
1061        vt = getattr(msg, "validDate", None)
1062        if vt is None:
1063            rd = msg.refDate
1064            lt = msg.leadTime
1065            try:
1066                vt = rd + lt
1067            except Exception:
1068                vt = rd
1069        if hasattr(vt, "isoformat"):
1070            return vt.isoformat()
1071        if isinstance(vt, np.datetime64):
1072            return str(vt)
1073        return str(vt)
1074    elif dim_name == "duration":
1075        d = msg.duration
1076        if hasattr(d, "total_seconds"):
1077            return d.total_seconds()
1078        return str(d)
1079    else:
1080        val = getattr(msg, dim_name, None)
1081        if hasattr(val, "value"):
1082            val = val.value
1083        return val
1084
1085
1086def _map_messages_to_dimensions(
1087    msg_entries: List[_MsgEntry],
1088    level_dim_name: str = "level",
1089) -> dict:
1090    """Group messages by variable and map to dimension indices.
1091
1092    This mirrors the logic in ``parse_grib_index()`` from the xarray
1093    backend: for each message, extract the values of potential dimension
1094    coordinates (level, leadTime, refDate, perturbationNumber, etc.),
1095    determine which dimensions have more than one unique value, and
1096    build a mapping from dimension-value tuples to message indices.
1097
1098    Parameters
1099    ----------
1100    msg_entries : list of _MsgEntry
1101        All message entries for a single variable.
1102
1103    Returns
1104    -------
1105    dict
1106        Dictionary with keys:
1107        - ``"dim_names"``: ordered list of active dimension names
1108        - ``"dim_values"``: dict mapping dim name to sorted unique values
1109        - ``"msg_index_map"``: dict mapping dim-value tuple to entry index
1110    """
1111    # Build ordered dim names, substituting the surface-type-specific level name
1112    ordered_dims = [level_dim_name if d == "level" else d for d in _ORDERED_DIM_NAMES]
1113
1114    # Collect dimension values for each message
1115    all_dim_vals: Dict[str, list] = {d: [] for d in ordered_dims}
1116
1117    for entry in msg_entries:
1118        msg = entry.msg
1119        for dim_name in ordered_dims:
1120            # Map the (possibly renamed) level dim back to "level" for _get_dim_value
1121            orig_name = "level" if dim_name == level_dim_name else dim_name
1122            try:
1123                val = _get_dim_value(msg, orig_name)
1124                all_dim_vals[dim_name].append(val)
1125            except (AttributeError, TypeError):
1126                all_dim_vals[dim_name].append(None)
1127
1128    # Determine which dimensions are active (have >1 unique value)
1129    active_dims = []
1130    dim_values: Dict[str, list] = {}
1131
1132    for dim_name in ordered_dims:
1133        vals = all_dim_vals[dim_name]
1134        # Filter out None values
1135        non_none = [v for v in vals if v is not None]
1136        if not non_none:
1137            continue
1138        unique_vals = sorted(set(non_none))
1139        if len(unique_vals) > 1:
1140            active_dims.append(dim_name)
1141            dim_values[dim_name] = unique_vals
1142        elif len(unique_vals) == 1:
1143            # Always emit valid_time and the level dim (so multi-file concat
1144            # can grow those axes).  All other dims are optional: only emit
1145            # them when they actually vary within this variable group.
1146            if dim_name in _ALWAYS_INCLUDE_DIMS or dim_name == level_dim_name:
1147                active_dims.append(dim_name)
1148                dim_values[dim_name] = unique_vals
1149
1150    # If no dimensions are active at all, fall back to a single valid_time
1151    if not active_dims:
1152        msg = msg_entries[0].msg
1153        vt = _get_dim_value(msg, "valid_time")
1154        active_dims = ["valid_time"]
1155        dim_values = {"valid_time": [vt]}
1156
1157    # Remap back so callers can use dim_names as coordinate keys directly
1158    # (level_dim_name is already baked in via ordered_dims)
1159
1160    # Build the mapping from dimension-value tuples to entry indices
1161    msg_index_map: Dict[tuple, int] = {}
1162    for idx, entry in enumerate(msg_entries):
1163        msg = entry.msg
1164        dim_tuple = tuple(_get_dim_value(msg, "level" if d == level_dim_name else d) for d in active_dims)
1165        msg_index_map[dim_tuple] = idx
1166
1167    return {
1168        "dim_names": active_dims,
1169        "dim_values": dim_values,
1170        "msg_index_map": msg_index_map,
1171    }
1172
1173
1174def _build_coord_refs(
1175    dim_name: str,
1176    values: list,
1177    refs: Dict[str, Any],
1178) -> None:
1179    """Build inline base64-encoded coordinate array refs.
1180
1181    Parameters
1182    ----------
1183    dim_name : str
1184        Coordinate/dimension name.
1185    values : list
1186        Sorted unique coordinate values.
1187    refs : dict
1188        The refs dict to populate.
1189    """
1190    if values and isinstance(values[0], tuple):
1191        # Level-type coordinate: values are (v1, v2) tuples; use v1
1192        coord_values = np.array(
1193            [v[0] if isinstance(v, tuple) else float(v) for v in values],
1194            dtype=np.float64,
1195        )
1196    elif dim_name == "valid_time":
1197        # Store as int64 nanoseconds since epoch so xarray decodes as datetime64
1198        ns_vals = [int(np.datetime64(v, "ns").astype(np.int64)) for v in values]
1199        coord_values = np.array(ns_vals, dtype=np.int64)
1200    elif dim_name == "duration":
1201        # Store as float seconds
1202        coord_values = np.array(values, dtype=np.float64)
1203    elif dim_name == "perturbationNumber":
1204        coord_values = np.array(values, dtype=np.int32)
1205    elif dim_name == "percentileValue":
1206        coord_values = np.array(values, dtype=np.float64)
1207    else:
1208        coord_values = np.array(values, dtype=np.float64)
1209
1210    # Encode as base64
1211    raw_bytes = coord_values.tobytes()
1212    b64_data = base64.b64encode(raw_bytes).decode("ascii")
1213
1214    # .zarray for the coordinate
1215    coord_zarray = {
1216        "zarr_format": 2,
1217        "shape": [len(values)],
1218        "chunks": [len(values)],
1219        "dtype": coord_values.dtype.str,
1220        "fill_value": None if coord_values.dtype.kind in {"U", "S"} else 0,
1221        "order": "C",
1222        "compressor": None,
1223        "filters": None,
1224    }
1225    refs[f"{dim_name}/.zarray"] = json.dumps(coord_zarray)
1226
1227    # .zattrs for the coordinate
1228    coord_zattrs: dict = {"_ARRAY_DIMENSIONS": [dim_name]}
1229    if dim_name == "valid_time":
1230        # CF-compliant time metadata so xarray decodes int64 ns as datetime64
1231        coord_zattrs["units"] = "nanoseconds since 1970-01-01T00:00:00"
1232        coord_zattrs["calendar"] = "proleptic_gregorian"
1233    refs[f"{dim_name}/.zattrs"] = json.dumps(coord_zattrs)
1234
1235    # Inline data chunk
1236    refs[f"{dim_name}/0"] = "base64:" + b64_data
1237
1238
1239def _build_latlon_coord_refs(msg, refs: Dict[str, Any]) -> None:
1240    """Build inline latitude/longitude 2-D coordinate arrays from the grid.
1241
1242    Calls ``msg.latlons()`` to compute the full (ny, nx) grids and encodes
1243    them as base64 inline Zarr refs, matching the xarray backend's behaviour.
1244    """
1245    try:
1246        lats, lons = msg.latlons()
1247    except Exception:
1248        return
1249
1250    ny, nx = int(msg.ny), int(msg.nx)
1251
1252    for name, data, attrs in [
1253        (
1254            "latitude",
1255            lats.astype(np.float64),
1256            {
1257                "_ARRAY_DIMENSIONS": ["y", "x"],
1258                "standard_name": "latitude",
1259                "units": "degrees_north",
1260            },
1261        ),
1262        (
1263            "longitude",
1264            lons.astype(np.float64),
1265            {
1266                "_ARRAY_DIMENSIONS": ["y", "x"],
1267                "standard_name": "longitude",
1268                "units": "degrees_east",
1269            },
1270        ),
1271    ]:
1272        # Skip if already present (e.g. from a prior variable group)
1273        if f"{name}/.zarray" in refs:
1274            continue
1275
1276        zarray = {
1277            "zarr_format": 2,
1278            "shape": [ny, nx],
1279            "chunks": [ny, nx],
1280            "dtype": "<f8",
1281            "fill_value": None,
1282            "order": "C",
1283            "compressor": None,
1284            "filters": None,
1285        }
1286        refs[f"{name}/.zarray"] = json.dumps(zarray)
1287        refs[f"{name}/.zattrs"] = json.dumps(attrs)
1288        refs[f"{name}/0.0"] = "base64:" + base64.b64encode(data.tobytes()).decode("ascii")