Storage¶

class Writer(abc.ABC, Loggable):
    """Abstract base class for storage backend writers.

    Writers are long-lived singletons, created once when the application
    starts and reused across every acquisition.  The class-level registry
    is keyed by a **composite key** ``(name, mimetype)`` where:

    - ``mimetype`` identifies the storage format and determines which
      ``Writer`` subclass handles serialisation (e.g.
      ``"application/x-zarr"`` maps to ``ZarrWriter``).
    - ``name`` is a **store group name** that distinguishes multiple
      independent stores that share the same format.  Use ``"default"``
      for the common single-store case.  This is *not* a device name —
      many devices may contribute sources to the same writer by
      referencing the same ``(name, mimetype)`` key.

    Example keys:

        ("default", "application/x-zarr")  # the normal single-store case
        ("live", "application/x-zarr")     # a separate live-preview store
        ("default", "application/x-hdf5")  # a different format entirely

    URI setting is controlled by the Presenter layer
    via [`set_uri`][redsun.storage.Writer.set_uri].
    Subclasses must override ``set_uri`` to perform any backend-specific path
    translation without disturbing the registry.

    Call order per acquisition:

    1. ``set_uri(uri)``
        - called by a presenter before the plan
    2. ``register(source_name, data_key, dtype, shape, capacity)``
        - called by each device in its own ``prepare()``
        - returns a [`FrameSink`][redsun.storage.FrameSink]
    3. ``kickoff()``
        - opens the backend
    4. ``sink.write(frame)``
        - push frames (thread-safe)
    5. ``sink.close()``
        - signals completion for this source

    Subclasses must implement the following methods and properties:

    - [`mimetype`][redsun.storage.Writer.mimetype] — MIME type string
    - [`kickoff`][redsun.storage.Writer.kickoff]
        - open the backend
        - must call ``super().kickoff()`` to set ``is_open`` and enforce the URI guard
    - [`_on_prepare`][redsun.storage.Writer._on_prepare] (private)
        - backend-specific setup after a source is registered (e.g. pre-declare Zarr array
          dimensions); called at the end of [`prepare`][redsun.storage.Writer.prepare]
    - [`_write_frame`][redsun.storage.Writer._write_frame]
        - write one frame to the backend
    - [`_finalize`][redsun.storage.Writer._finalize]
        - close the backend when all sources are done
    - [`_class_mimetype`][redsun.storage.Writer._class_mimetype] (class method)
        - return the MIME type string for this subclass; used for registry keys

    Parameters
    ----------
    name : str
        Store group name.  See class docstring for semantics.
    """

    _registry: dict[_WriterKey, "Writer"] = {}
    _registry_lock: th.Lock = th.Lock()

    def __init__(self, name: str) -> None:
        self._name = name
        self._uri: str = ""
        self._lock = th.Lock()
        self._is_open = False
        self._sources: dict[str, SourceInfo] = {}
        self._metadata: dict[str, dict[str, Any]] = {}

    @classmethod
    def get(cls: type[_W], name: str = "default") -> _W:
        """Return the singleton writer for *(name, cls.mimetype)*.

        Creates a new instance on first call; returns the existing one
        on every subsequent call for the same ``(name, mimetype)`` pair.
        The URI is *not* set here — call
        [`set_uri`][redsun.storage.Writer.set_uri] separately.

        This method is normally not called directly by devices or
        application code.  Use
        [`make_writer`][redsun.storage.make_writer] instead, which
        resolves the correct subclass from the mimetype string.

        Parameters
        ----------
        name : str
            Store group name.  Defaults to ``"default"``.

        Returns
        -------
        Writer
            Singleton instance for ``(name, cls.mimetype)``.

        Raises
        ------
        TypeError
            If the existing registry entry is not an instance of ``cls``.
        """
        # cls.mimetype is an abstract property; each concrete subclass
        # provides a fixed string (e.g. ZarrWriter.mimetype == "application/x-zarr").
        # We access it via the class directly to avoid needing an instance.
        key: _WriterKey = (name, cls._class_mimetype())
        with cls._registry_lock:
            if key not in cls._registry:
                cls._registry[key] = cls(name)
            instance = cls._registry[key]
        if not isinstance(instance, cls):
            raise TypeError(
                f"Registry entry for {key!r} is {type(instance).__name__!r}, "
                f"expected {cls.__name__!r}"
            )
        return instance

    @classmethod
    def release(cls, name: str = "default") -> None:
        """Remove the registry entry for *(name, cls.mimetype)*.

        Called by a presenter at application shutdown.
        Devices should not call this directly.

        Parameters
        ----------
        name : str
            Store group name.  Defaults to ``"default"``.
        """
        key: _WriterKey = (name, cls._class_mimetype())
        with cls._registry_lock:
            cls._registry.pop(key, None)

    @classmethod
    @abc.abstractmethod
    def _class_mimetype(cls) -> str:
        """Return the MIME type string for this subclass.

        Used by [`get`][redsun.storage.Writer.get] and
        [`release`][redsun.storage.Writer.release] to build the registry
        key before any instance exists.  Must return the same value as
        the instance property [`mimetype`][redsun.storage.Writer.mimetype].

        Subclasses implement this as a one-liner::

            @classmethod
            def _class_mimetype(cls) -> str:
                return "application/x-zarr"
        """
        ...

    @property
    def is_open(self) -> bool:
        """Return whether the backend is currently open."""
        return self._is_open

    @property
    def uri(self) -> str:
        """The current store URI.

        Empty string until
        [`set_uri`][redsun.storage.Writer.set_uri] has been called.
        """
        return self._uri

    @property
    def mimetype(self) -> str:
        """The MIME type string for this backend.

        Must return the same value as
        [`_class_mimetype`][redsun.storage.Writer._class_mimetype].
        """
        return self._class_mimetype()

    def set_uri(self, uri: str) -> None:
        """Update the store URI for the next acquisition.

        Called by a presenter before each acquisition and
        whenever the user changes the output directory.  The writer
        must not be open when this is called.

        Subclasses should override this to perform any backend-specific
        path translation (e.g. rebuilding ``StreamSettings.store_path``
        for ``ZarrWriter``) and must call ``super().set_uri(uri)``.

        Parameters
        ----------
        uri : str
            New store URI (e.g. ``"file:///data/2026_02_25/scan_00001"``).

        Raises
        ------
        RuntimeError
            If the writer is currently open.
        """
        if self._is_open:
            raise RuntimeError(
                f"Cannot change URI on writer ({self._name!r}, {self.mimetype!r}) "
                "while it is open."
            )
        self._uri = uri
        self.logger.debug(f"URI updated to {uri!r}")

    def prepare(
        self,
        name: str,
        data_key: str,
        dtype: np.dtype[np.generic],
        shape: tuple[int, ...],
        capacity: int = 0,
    ) -> FrameSink:
        """Register a data source and return a ready ``FrameSink``.

        Called by each device inside its own ``prepare()`` method, once
        per acquisition.  Replaces the former two-step
        ``update_source`` + ``prepare`` sequence.

        Safe to call multiple times on the same source name — counters
        are reset on each call, making it suitable for repeated
        acquisitions without recreating the writer.

        Parameters
        ----------
        name : str
            Source name, typically the device name.  Multiple devices
            may register distinct source names on the same writer.
        data_key : str
            Bluesky data key used in stream documents.
        dtype : np.dtype[np.generic]
            NumPy data type of the frames.
        shape : tuple[int, ...]
            Shape of each individual frame.
        capacity : int
            Maximum number of frames to accept.  ``0`` means unlimited.

        Returns
        -------
        FrameSink
            Bound sink ready to accept frames via ``sink.write(frame)``.

        Raises
        ------
        RuntimeError
            If the writer is currently open.
        """
        if self._is_open:
            raise RuntimeError(
                f"Cannot register source {name!r} on writer "
                f"({self._name!r}, {self.mimetype!r}) while it is open."
            )
        self._sources[name] = SourceInfo(
            name=name,
            dtype=dtype,
            shape=shape,
            data_key=data_key,
            capacity=capacity,
        )
        self.logger.debug(
            f"Registered source {name!r} — shape={shape}, capacity={capacity}"
        )
        self._on_prepare(name)
        return FrameSink(self, name)

    @abc.abstractmethod
    def _on_prepare(self, name: str) -> None:
        """Backend-specific hook called after a source is registered.

        Invoked at the end of [`prepare`][redsun.storage.Writer.prepare]
        once ``self._sources[name]`` is fully populated.  Subclasses use
        this to pre-declare backend structures — for example, ``ZarrWriter``
        builds its ``ArraySettings`` here from the source dtype and shape.

        Parameters
        ----------
        name : str
            The source name just registered.
        """
        ...

    def get_indices_written(self, name: str | None = None) -> int:
        """Return the number of frames written for a source.

        Parameters
        ----------
        name : str | None
            Source name.  If ``None``, returns the minimum across all
            registered sources (useful for synchronisation checks).

        Raises
        ------
        KeyError
            If *name* is not registered.
        """
        if name is None:
            if not self._sources:
                return 0
            return min(s.frames_written for s in self._sources.values())
        if name not in self._sources.keys():
            raise KeyError(f"Unknown source {name!r}")
        return self._sources[name].frames_written

    def reset_collection_state(self, name: str) -> None:
        """Reset the stream-document counters for *name*.

        Parameters
        ----------
        name : str
            Source name to reset.
        """
        source = self._sources[name]
        source.collection_counter = 0
        source.stream_resource_uid = str(uuid.uuid4())

    @abc.abstractmethod
    def kickoff(self) -> None:
        """Open the storage backend for a new acquisition.

        Subclasses must call ``super().kickoff()`` to set
        [`is_open`][redsun.storage.Writer.is_open] and to enforce the
        URI guard.  Subsequent calls while already open must be no-ops.

        Raises
        ------
        RuntimeError
            If [`uri`][redsun.storage.Writer.uri] has not been set yet.
        """
        if not self._uri:
            clear_metadata()
            raise RuntimeError(
                f"Writer ({self._name!r}, {self.mimetype!r}) has no URI. "
                "A presenter must call set_uri() before kickoff()."
            )
        if not self._is_open:
            self._metadata = snapshot_metadata()
            self._is_open = True

    def complete(self, name: str) -> None:
        """Mark acquisition complete for source *name*.

        Called automatically by [`FrameSink.close`][redsun.storage.FrameSink.close].
        When the last registered source calls ``complete``, the backend
        is finalised and ``is_open`` is reset.  The writer instance
        remains in the registry and is ready for the next acquisition.

        Parameters
        ----------
        name : str
            Source name signalling completion.
        """
        self._sources[name].completed = True
        if all(s.completed for s in self._sources.values()):
            self._finalize()
            self._is_open = False
            clear_metadata()
            self.logger.debug("All sources complete; backend finalised.")

    def clear_sources(self) -> None:
        """Remove all registered sources.

        A presenter in charge of monitoring
        writing progress should take care to call this
        after each plan is finished.
        """
        self._sources.clear()
        self.logger.debug("Source cache reset.")

    @abc.abstractmethod
    def _write_frame(self, name: str, frame: npt.NDArray[np.generic]) -> None:
        """Write one frame to the backend.

        Called by [`FrameSink.write`][redsun.storage.FrameSink.write]
        under the writer lock.

        Parameters
        ----------
        name : str
            Source name.
        frame : npt.NDArray[np.generic]
            Frame data to write.
        """
        ...

    @abc.abstractmethod
    def _finalize(self) -> None:
        """Close the backend after all sources have completed."""
        ...

    def collect_stream_docs(
        self,
        name: str,
        indices_written: int,
    ) -> Iterator[StreamAsset]:
        """Yield ``StreamResource`` and ``StreamDatum`` documents for *name*.

        Parameters
        ----------
        name : str
            Source name.
        indices_written : int
            Number of frames to report in this call.

        Yields
        ------
        StreamAsset
            Tuples of ``("stream_resource", doc)`` or
            ``("stream_datum", doc)``.

        Raises
        ------
        KeyError
            If *name* is not registered.
        """
        if name not in self._sources:
            raise KeyError(f"Unknown source {name!r}")

        source = self._sources[name]

        if indices_written == 0:
            return

        frames_to_report = min(indices_written, source.frames_written)

        if source.collection_counter >= frames_to_report:
            return

        if source.collection_counter == 0:
            stream_resource: StreamResource = {
                "data_key": source.data_key,
                "mimetype": self.mimetype,
                "parameters": {"array_name": source.name},
                "uid": source.stream_resource_uid,
                "uri": self.uri,
            }
            yield ("stream_resource", stream_resource)

        stream_datum: StreamDatum = {
            "descriptor": "",
            "indices": {"start": source.collection_counter, "stop": frames_to_report},
            "seq_nums": {"start": 0, "stop": 0},
            "stream_resource": source.stream_resource_uid,
            "uid": f"{source.stream_resource_uid}/{source.collection_counter}",
        }
        yield ("stream_datum", stream_datum)

        source.collection_counter = frames_to_report

class FrameSink:
    """Device-facing handle for pushing frames to a storage backend.

    Returned by [`Writer.register`][redsun.storage.Writer.register].
    Devices write frames by calling [`write`][redsun.storage.FrameSink.write];
    the sink routes each frame to the backend and updates the frame counter
    atomically.

    Calling [`close`][redsun.storage.FrameSink.close] signals that no more
    frames will arrive from this source and triggers backend finalisation
    once all sources are complete.

    Parameters
    ----------
    writer : Writer
        The writer that owns this sink.
    name : str
        Source name this sink is bound to.
    """

    def __init__(self, writer: Writer, name: str) -> None:
        self._writer = writer
        self._name = name

    def write(self, frame: npt.NDArray[np.generic]) -> None:
        """Push ``frame`` to the storage backend.

        Thread-safe.

        Parameters
        ----------
        frame : npt.NDArray[np.generic]
            Array data to write. dtype and shape must match those declared
            in [`Writer.register`][redsun.storage.Writer.register].
        """
        with self._writer._lock:
            self._writer._write_frame(self._name, frame)
            self._writer._sources[self._name].frames_written += 1

    def close(self) -> None:
        """Signal that no more frames will be written from this sink.

        Delegates to [`Writer.complete`][redsun.storage.Writer.complete].
        """
        self._writer.complete(self._name)

@dataclass
class PathInfo:
    """Where and how a storage backend should write data for one device.

    !!! note

        For local file storage, the `store_uri` must be converted to
        a concrete filesystem path before use. This is the responsibility
        of the backend writer.
    """

    store_uri: str
    """URI of the storage root."""
    array_key: str
    """Key (array name) within the store for this device's data."""

    capacity: int = 0
    """Capacity in frames.  ``0`` means unlimited."""

    mimetype_hint: str = "application/x-zarr"
    """MIME type hint for the backend.  Consumers may use this to select the correct reader."""

    extra: dict[str, Any] = field(default_factory=dict)
    """Optional backend-specific metadata (e.g. OME-Zarr axis labels, physical units).  Base writers ignore this field."""

@runtime_checkable
class FilenameProvider(Protocol):
    """Callable that produces a filename stem for a given key."""

    def __call__(self, key: str | None = None, group: str | None = None) -> str:
        """Return a filename stem for *key* and *group*.

        Parameters
        ----------
        key : str | None
            Discriminator passed by the caller — typically a plan name
            (when called from a presenter) or a device name (when called
            from a writer).  Implementations may ignore it if the filename
            is key-agnostic.
        group : str | None
            Writer group name (e.g. ``'default'``).  When provided,
            implementations should embed it in the filename so that
            outputs from different writer groups do not collide.

        Returns
        -------
        str
            A filename stem without extension.
        """
        ...