# This file is part of pipe_base.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (http://www.lsst.org).
# See the COPYRIGHT file at the top-level directory of this distribution
# for details of code ownership.
#
# This software is dual licensed under the GNU General Public License and also
# under a 3-clause BSD license. Recipients may choose which of these licenses
# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
# respectively.  If you choose the GPL option then the following text applies
# (but note that there is still no warranty even if you opt for BSD instead):
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""The base class for the QuantumGraph-generation algorithm and various
helper classes.
"""

from __future__ import annotations

__all__ = (
    "QuantumGraphBuilder",
    "ExistingDatasets",
    "QuantumGraphBuilderError",
    "OutputExistsError",
    "PrerequisiteMissingError",
)

import dataclasses
from abc import ABC, abstractmethod
from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, final

from lsst.daf.butler import (
    Butler,
    CollectionType,
    DataCoordinate,
    DatasetRef,
    DatasetType,
    DimensionUniverse,
    NamedKeyDict,
    NamedKeyMapping,
    Quantum,
)
from lsst.daf.butler.registry import MissingCollectionError, MissingDatasetTypeError
from lsst.utils.logging import LsstLogAdapter, getLogger
from lsst.utils.timer import timeMethod

from . import automatic_connection_constants as acc
from ._status import NoWorkFound
from ._task_metadata import TaskMetadata
from .connections import AdjustQuantumHelper
from .graph import QuantumGraph
from .pipeline_graph import PipelineGraph, TaskNode
from .prerequisite_helpers import PrerequisiteInfo, SkyPixBoundsBuilder, TimespanBuilder
from .quantum_graph_skeleton import (
    DatasetKey,
    PrerequisiteDatasetKey,
    QuantumGraphSkeleton,
    QuantumKey,
    TaskInitKey,
)

if TYPE_CHECKING:
    from .pipeline import TaskDef


class QuantumGraphBuilderError(Exception):
    """Base class for exceptions generated by QuantumGraphBuilder."""

    pass


class GraphBuilderError(QuantumGraphBuilderError):
    """Backwards-compatibility near-alias for QuantumGraphBuilderError."""

    pass


# Inherit from backwards-compatibility alias for backwards-compatibility.
class OutputExistsError(GraphBuilderError):
    """Exception generated when output datasets already exist."""

    pass


# Inherit from backwards-compatibility alias for backwards-compatibility.
class PrerequisiteMissingError(GraphBuilderError):
    """Exception generated when a prerequisite dataset does not exist."""

    pass


class InitInputMissingError(QuantumGraphBuilderError):
    """Exception generated when an init-input dataset does not exist."""

    pass


class QuantumGraphBuilder(ABC):
    """An abstract base class for building `QuantumGraph` objects from a
    pipeline.

    Parameters
    ----------
    pipeline_graph : `.pipeline_graph.PipelineGraph`
        Pipeline to build a `QuantumGraph` from, as a graph.  Will be resolved
        in-place with the given butler (any existing resolution is ignored).
    butler : `lsst.daf.butler.Butler`
        Client for the data repository.  Should be read-only.
    input_collections : `~collections.abc.Sequence` [ `str` ], optional
        Collections to search for overall-input datasets.  If not provided,
        ``butler.collections`` is used (and must not be empty).
    output_run : `str`, optional
        Output `~lsst.daf.butler.CollectionType.RUN` collection.  If not
        provided, ``butler.run`` is used (and must not be `None`).
    skip_existing_in : `~collections.abc.Sequence` [ `str` ], optional
        Collections to search for outputs that already exist for the purpose of
        skipping quanta that have already been run.
    clobber : `bool`, optional
        Whether to raise if predicted outputs already exist in ``output_run``
        (not including those quanta that would be skipped because they've
        already been run).  This never actually clobbers outputs; it just
        informs the graph generation algorithm whether execution will run with
        clobbering enabled.  This is ignored if ``output_run`` does not exist.

    Notes
    -----
    Constructing a `QuantumGraphBuilder` will run queries for existing datasets
    with empty data IDs (including but not limited to init inputs and outputs),
    in addition to resolving the given pipeline graph and testing for existence
    of the ``output`` run collection.

    The `build` method splits the pipeline graph into independent subgraphs,
    then calls the abstract method `process_subgraph` on each, to allow
    concrete implementations to populate the rough graph structure (the
    `~quantum_graph_skeleton.QuantumGraphSkeleton` class) and search for
    existing datasets (further populating the builder's `existing_datasets`
    struct).  The `build` method then:

    - assembles `lsst.daf.butler.Quantum` instances from all data IDs in the
      skeleton;
    - looks for existing outputs found in ``skip_existing_in`` to see if any
      quanta should be skipped;
    - calls `PipelineTaskConnections.adjustQuantum` on all quanta, adjusting
      downstream quanta appropriately when preliminary predicted outputs are
      rejected (pruning nodes that will not have the inputs they need to run);
    - attaches datastore records and registry dataset types to the graph.

    In addition to implementing `process_subgraph`, derived classes are
    generally expected to add new construction keyword-only arguments to
    control the data IDs of the quantum graph, while forwarding all of the
    arguments defined in the base class to `super`.
    """

    def __init__(
        self,
        pipeline_graph: PipelineGraph,
        butler: Butler,
        *,
        input_collections: Sequence[str] | None = None,
        output_run: str | None = None,
        skip_existing_in: Sequence[str] = (),
        clobber: bool = False,
    ):
        self.log = getLogger(__name__)
        self.metadata = TaskMetadata()
        self._pipeline_graph = pipeline_graph
        self.butler = butler
        if input_collections is None:
            input_collections = butler.collections
        if not input_collections:
            raise ValueError("No input collections provided.")
        self.input_collections = input_collections
        if output_run is None:
            output_run = butler.run
        if not output_run:
            raise ValueError("No output RUN collection provided.")
        self.output_run = output_run
        self.skip_existing_in = skip_existing_in
        self.empty_data_id = DataCoordinate.make_empty(butler.dimensions)
        self.clobber = clobber
        # See whether the output run already exists.
        self.output_run_exists = False
        try:
            if self.butler.registry.getCollectionType(self.output_run) is not CollectionType.RUN:
                raise RuntimeError(f"{self.output_run!r} is not a RUN collection.")
            self.output_run_exists = True
        except MissingCollectionError:
            # If the run doesn't exist we never need to clobber.  This is not
            # an error so you can run with clobber=True the first time you
            # attempt some processing as well as all subsequent times, instead
            # of forcing the user to make the first attempt different.
            self.clobber = False
        # We need to know whether the skip_existing_in collection sequence
        # starts with the output run collection, as an optimization to avoid
        # queries later.
        try:
            skip_existing_in_flat = self.butler.registry.queryCollections(
                self.skip_existing_in, flattenChains=True
            )
        except MissingCollectionError:
            skip_existing_in_flat = []
        if not skip_existing_in_flat:
            self.skip_existing_in = []
        if self.skip_existing_in and self.output_run_exists:
            self.skip_existing_starts_with_output_run = self.output_run == skip_existing_in_flat[0]
        else:
            self.skip_existing_starts_with_output_run = False
        self.existing_datasets = ExistingDatasets()
        try:
            packages_storage_class = butler.get_dataset_type(acc.PACKAGES_INIT_OUTPUT_NAME).storageClass_name
        except MissingDatasetTypeError:
            packages_storage_class = acc.PACKAGES_INIT_OUTPUT_STORAGE_CLASS
        self._global_init_output_types = {
            acc.PACKAGES_INIT_OUTPUT_NAME: DatasetType(
                acc.PACKAGES_INIT_OUTPUT_NAME,
                self.universe.empty,
                packages_storage_class,
            )
        }
        with self.butler.registry.caching_context():
            self._pipeline_graph.resolve(self.butler.registry)
            self._find_empty_dimension_datasets()
            self.prerequisite_info = {
                task_node.label: PrerequisiteInfo(task_node, self._pipeline_graph)
                for task_node in pipeline_graph.tasks.values()
            }

    log: LsstLogAdapter
    """Logger to use for all quantum-graph generation messages.

    General and per-task status messages should be logged at `~logging.INFO`
    level or higher, per-dataset-type status messages should be logged at
    `~lsst.utils.logging.VERBOSE` or higher, and per-data-ID status messages
    should be logged at `logging.DEBUG` or higher.
    """

    metadata: TaskMetadata
    """Metadata to store in the QuantumGraph.

    The `TaskMetadata` class is used here primarily in order to enable
    resource-usage collection with the `lsst.utils.timer.timeMethod` decorator.
    """

    butler: Butler
    """Client for the data repository.

    Should be read-only.
    """

    input_collections: Sequence[str]
    """Collections to search for overall-input datasets.
    """

    output_run: str
    """Output `~lsst.daf.butler.CollectionType.RUN` collection.
    """

    skip_existing_in: Sequence[str]
    """Collections to search for outputs that already exist for the purpose
    of skipping quanta that have already been run.
    """

    clobber: bool
    """Whether to raise if predicted outputs already exist in ``output_run``

    This never actually clobbers outputs; it just informs the graph generation
    algorithm whether execution will run with clobbering enabled.  This is
    always `False` if `output_run_exists` is `False`.
    """

    empty_data_id: DataCoordinate
    """An empty data ID in the data repository's dimension universe.
    """

    output_run_exists: bool
    """Whether the output run exists in the data repository already.
    """

    skip_existing_starts_with_output_run: bool
    """Whether the `skip_existing_in` sequence begins with `output_run`.

    If this is true, any dataset found in `output_run` can be used to
    short-circuit queries in `skip_existing_in`.
    """

    existing_datasets: ExistingDatasets
    """Struct holding datasets that have already been found in the data
    repository.

    This is updated in-place as the `QuantumGraph` generation algorithm
    proceeds.
    """

    prerequisite_info: Mapping[str, PrerequisiteInfo]
    """Helper objects for finding prerequisite inputs, organized by task label.

    Subclasses that find prerequisites should remove the
    covered `~prerequisite_helpers.PrerequisiteFinder` objects from this
    attribute.
    """

    @property
    def universe(self) -> DimensionUniverse:
        """Definitions of all data dimensions."""
        return self.butler.dimensions

    @final
    @timeMethod
    def build(self, metadata: Mapping[str, Any] | None = None) -> QuantumGraph:
        """Build the quantum graph.

        Parameters
        ----------
        metadata : `~collections.abc.Mapping`, optional
            Flexible metadata to add to the quantum graph.

        Returns
        -------
        quantum_graph : `QuantumGraph`
            DAG describing processing to be performed.

        Notes
        -----
        External code is expected to construct a `QuantumGraphBuilder` and then
        call this method exactly once.  See class documentation for details on
        what it does.
        """
        with self.butler.registry.caching_context():
            full_skeleton = QuantumGraphSkeleton(self._pipeline_graph.tasks)
            subgraphs = list(self._pipeline_graph.split_independent())
            for i, subgraph in enumerate(subgraphs):
                self.log.info(
                    "Processing pipeline subgraph %d of %d with %d task(s).",
                    i + 1,
                    len(subgraphs),
                    len(subgraph.tasks),
                )
                self.log.verbose("Subgraph tasks: [%s]", ", ".join(label for label in subgraph.tasks))
                subgraph_skeleton = self.process_subgraph(subgraph)
                full_skeleton.update(subgraph_skeleton)
            # Loop over tasks.  The pipeline graph must be topologically
            # sorted, so a quantum is only processed after any quantum that
            # provides its inputs has been processed.
            for task_node in self._pipeline_graph.tasks.values():
                self._resolve_task_quanta(task_node, full_skeleton)
            # Add global init-outputs to the skeleton.
            for dataset_type in self._global_init_output_types.values():
                dataset_key = full_skeleton.add_dataset_node(
                    dataset_type.name, self.empty_data_id, is_global_init_output=True
                )
                ref = self.existing_datasets.outputs_in_the_way.get(dataset_key)
                if ref is None:
                    ref = DatasetRef(dataset_type, self.empty_data_id, run=self.output_run)
                full_skeleton[dataset_key]["ref"] = ref
            # Remove dataset nodes with no edges that are not global init
            # outputs, which are generally overall-inputs whose original quanta
            # end up skipped or with no work to do (we can't remove these along
            # with the quanta because no quantum knows if its the only
            # consumer).
            full_skeleton.remove_orphan_datasets()
            self._attach_datastore_records(full_skeleton)
            # TODO initialize most metadata here instead of in ctrl_mpexec.
            if metadata is None:
                metadata = {}
            return self._construct_quantum_graph(full_skeleton, metadata)

    @abstractmethod
    def process_subgraph(self, subgraph: PipelineGraph) -> QuantumGraphSkeleton:
        """Build the rough structure for an independent subset of the
        `QuantumGraph` and query for relevant existing datasets.

        Parameters
        ----------
        subgraph : `.pipeline_graph.PipelineGraph`
            Subset of the pipeline graph that should be processed by this call.
            This is always resolved and topologically sorted.  It should not be
            modified.

        Returns
        -------
        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
            Class representing an initial quantum graph. See
            `quantum_graph_skeleton.QuantumGraphSkeleton` docs for details.
            After this is returned, the object may be modified in-place in
            unspecified ways.

        Notes
        -----
        In addition to returning a
        `quantum_graph_skeleton.QuantumGraphSkeleton`, this method should
        populate the `existing_datasets` structure by querying for all relevant
        datasets with non-empty data IDs (those with empty data IDs will
        already be present).  In particular:

        - `~ExistingDatasets.inputs` must always be populated with all
          overall-input datasets (but not prerequisites), by querying
          `input_collections`;
        - `~ExistingDatasets.outputs_for_skip` must be populated with any
          intermediate our output datasets present in `skip_existing_in` (it
          can be ignored if `skip_existing_in` is empty);
        - `~ExistingDatasets.outputs_in_the_way` must be populated with any
          intermediate or output datasets present in `output_run`, if
          `output_run_exists` (it can be ignored if `output_run_exists` is
          `False`).  Note that the presence of such datasets is not
          automatically an error, even if `clobber is `False`, as these may be
          quanta that will be skipped.
        - `~ExistingDatasets.inputs` must be populated with all
          prerequisite-input datasets that were included in the skeleton, by
          querying `input_collections` (not all prerequisite inputs need to be
          included in the skeleton, but the base class can only use per-quantum
          queries to find them, and that can be slow when there are many
          quanta).

        Dataset types should never be components and should always use the
        "common" storage class definition in `pipeline_graph.DatasetTypeNode`
        (which is the data repository definition when the dataset type is
        registered).
        """
        raise NotImplementedError()

    @final
    @timeMethod
    def _resolve_task_quanta(self, task_node: TaskNode, skeleton: QuantumGraphSkeleton) -> None:
        """Process the quanta for one task in a skeleton graph to skip those
        that have already completed and adjust those that request it.

        Parameters
        ----------
        task_node : `pipeline_graph.TaskNode`
            Node for this task in the pipeline graph.
        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
            Preliminary quantum graph, to be modified in-place.

        Notes
        -----
        This method modifies ``skeleton`` in-place in several ways:

        - It adds a "ref" attribute to dataset nodes, using the contents of
          `existing_datasets`.  This ensures producing and consuming tasks
          start from the same `DatasetRef`.
        - It adds "inputs", "outputs", and "init_inputs" attributes to the
          quantum nodes, holding the same `NamedValueMapping` objects needed to
          construct an actual `Quantum` instances.
        - It removes quantum nodes that are to be skipped because their outputs
          already exist in `skip_existing_in`.  It also removes their outputs
          from `ExistingDatasets.outputs_in_the_way`.
        - It adds prerequisite dataset nodes and edges that connect them to the
          quanta that consume them.
        - It removes quantum nodes whose
          `~PipelineTaskConnections.adjustQuantum` calls raise `NoWorkFound` or
          predict no outputs;
        - It removes the nodes of output datasets that are "adjusted away".
        - It removes the edges of input datasets that are "adjusted away".

        The difference between how adjusted inputs and outputs are handled
        reflects the fact that many quanta can share the same input, but only
        one produces each output.  This can lead to the graph having
        superfluous isolated nodes after processing is complete, but these
        should only be removed after all the quanta from all tasks have been
        processed.
        """
        # Extract the helper object for the prerequisite inputs of this task,
        # and tell it to prepare to construct skypix bounds and timespans for
        # each quantum (these will automatically do nothing if nothing needs
        # those bounds).
        task_prerequisite_info = self.prerequisite_info[task_node.label]
        task_prerequisite_info.update_bounds()
        # Loop over all quanta for this task, remembering the ones we've
        # gotten rid of.
        skipped_quanta = []
        no_work_quanta = []
        for quantum_key in skeleton.get_quanta(task_node.label):
            if self._skip_quantum_if_metadata_exists(task_node, quantum_key, skeleton):
                skipped_quanta.append(quantum_key)
                continue
            quantum_data_id = skeleton[quantum_key]["data_id"]
            skypix_bounds_builder = task_prerequisite_info.bounds.make_skypix_bounds_builder(quantum_data_id)
            timespan_builder = task_prerequisite_info.bounds.make_timespan_builder(quantum_data_id)
            adjusted_outputs = self._gather_quantum_outputs(
                task_node, quantum_key, skeleton, skypix_bounds_builder, timespan_builder
            )
            adjusted_inputs = self._gather_quantum_inputs(
                task_node,
                quantum_key,
                skeleton,
                task_prerequisite_info,
                skypix_bounds_builder,
                timespan_builder,
            )
            # Give the task's Connections class an opportunity to remove
            # some inputs, or complain if they are unacceptable.  This will
            # raise if one of the check conditions is not met, which is the
            # intended behavior.
            helper = AdjustQuantumHelper(inputs=adjusted_inputs, outputs=adjusted_outputs)
            try:
                helper.adjust_in_place(
                    task_node._get_imported_data().connections, task_node.label, quantum_data_id
                )
            except NoWorkFound as err:
                # Do not generate this quantum; it would not produce any
                # outputs.  Remove it and all of the outputs it might have
                # produced from the skeleton.
                try:
                    _, connection_name, _ = err.args
                    details = f"not enough datasets for connection {connection_name}."
                except ValueError:
                    details = str(err)
                self.log.debug(
                    "No work found for quantum %s of task %s: %s",
                    quantum_key.data_id_values,
                    quantum_key.task_label,
                    details,
                )
                no_work_quanta.append(quantum_key)
                continue
            if helper.outputs_adjusted:
                if not any(adjusted_refs for adjusted_refs in helper.outputs.values()):
                    # No outputs also means we don't generate this quantum.
                    self.log.debug(
                        "No outputs predicted for quantum %s of task %s.",
                        quantum_key.data_id_values,
                        quantum_key.task_label,
                    )
                    no_work_quanta.append(quantum_key)
                    continue
                # Remove output nodes that were not retained by
                # adjustQuantum.
                skeleton.remove_dataset_nodes(
                    self._find_removed(skeleton.iter_outputs_of(quantum_key), helper.outputs)
                )
            if helper.inputs_adjusted:
                if not any(bool(adjusted_refs) for adjusted_refs in helper.inputs.values()):
                    raise QuantumGraphBuilderError(
                        f"adjustQuantum implementation for {task_node.label}@{quantum_key.data_id_values} "
                        "returned outputs but no inputs."
                    )
                # Remove input dataset edges that were not retained by
                # adjustQuantum.  We can't remove the input dataset nodes
                # because some other quantum might still want them.
                skeleton.remove_input_edges(
                    quantum_key, self._find_removed(skeleton.iter_inputs_of(quantum_key), helper.inputs)
                )
            # Save the adjusted inputs and outputs to the quantum node's
            # state so we don't have to regenerate those data structures
            # from the graph.
            skeleton[quantum_key]["inputs"] = helper.inputs
            skeleton[quantum_key]["outputs"] = helper.outputs
        for no_work_quantum in no_work_quanta:
            skeleton.remove_quantum_node(no_work_quantum, remove_outputs=True)
        for skipped_quantum in skipped_quanta:
            skeleton.remove_quantum_node(skipped_quantum, remove_outputs=False)
        remaining_quanta = skeleton.get_quanta(task_node.label)
        self._resolve_task_init(task_node, skeleton, bool(skipped_quanta))
        message_terms = []
        if no_work_quanta:
            message_terms.append(f"{len(no_work_quanta)} had no work to do")
        if skipped_quanta:
            message_terms.append(f"{len(no_work_quanta)} previously succeeded")
        message_parenthetical = f" ({', '.join(message_terms)})" if message_terms else ""
        if remaining_quanta:
            self.log.info(
                "Generated %s for task %s%s.",
                _quantum_or_quanta(len(remaining_quanta)),
                task_node.label,
                message_parenthetical,
            )
        else:
            self.log.info(
                "Dropping task %s because no quanta remain%s.", task_node.label, message_parenthetical
            )
            skeleton.remove_task(task_node.label)

    def _skip_quantum_if_metadata_exists(
        self, task_node: TaskNode, quantum_key: QuantumKey, skeleton: QuantumGraphSkeleton
    ) -> bool:
        """Identify and drop quanta that should be skipped because their
        metadata datasets already exist.

        Parameters
        ----------
        task_node : `pipeline_graph.TaskNode`
            Node for this task in the pipeline graph.
        quantum_key : `QuantumKey`
            Identifier for this quantum in the graph.
        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
            Preliminary quantum graph, to be modified in-place.

        Returns
        -------
        skipped : `bool`
            `True` if the quantum is being skipped and has been removed from
            the graph, `False` otherwise.

        Notes
        -----
        If the metadata dataset for this quantum exists in
        `ExistingDatasets.outputs_for_skip`, the quantum will be skipped. This
        causes the quantum node to be removed from the graph.  Dataset nodes
        that were previously the outputs of this quantum will have their "ref"
        attribute set from `ExistingDatasets.outputs_for_skip`, or will be
        removed if there is no such dataset there.  Any output dataset in
        `ExistingDatasets.outputs_in_the_way` will be removed.
        """
        metadata_dataset_key = DatasetKey(
            task_node.metadata_output.parent_dataset_type_name, quantum_key.data_id_values
        )
        if metadata_dataset_key in self.existing_datasets.outputs_for_skip:
            # This quantum's metadata is already present in the the
            # skip_existing_in collections; we'll skip it.  But the presence of
            # the metadata dataset doesn't guarantee that all of the other
            # outputs we predicted are present; we have to check.
            for output_dataset_key in list(skeleton.iter_outputs_of(quantum_key)):
                if (
                    output_ref := self.existing_datasets.outputs_for_skip.get(output_dataset_key)
                ) is not None:
                    # Populate the skeleton graph's node attributes
                    # with the existing DatasetRef, just like a
                    # predicted output of a non-skipped quantum.
                    skeleton[output_dataset_key]["ref"] = output_ref
                else:
                    # Remove this dataset from the skeleton graph,
                    # because the quantum that would have produced it
                    # is being skipped and it doesn't already exist.
                    skeleton.remove_dataset_nodes([output_dataset_key])
                # If this dataset was "in the way" (i.e. already in the
                # output run), it isn't anymore.
                self.existing_datasets.outputs_in_the_way.pop(output_dataset_key, None)
            # Removing the quantum node from the graph will happen outside this
            # function.
            return True
        return False

    @final
    def _gather_quantum_outputs(
        self,
        task_node: TaskNode,
        quantum_key: QuantumKey,
        skeleton: QuantumGraphSkeleton,
        skypix_bounds_builder: SkyPixBoundsBuilder,
        timespan_builder: TimespanBuilder,
    ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
        """Collect outputs or generate datasets for a preliminary quantum and
        put them in the form used by `~lsst.daf.butler.Quantum` and
        `~PipelineTaskConnections.adjustQuantum`.

        Parameters
        ----------
        task_node : `pipeline_graph.TaskNode`
            Node for this task in the pipeline graph.
        quantum_key : `QuantumKey`
            Identifier for this quantum in the graph.
        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
            Preliminary quantum graph, to be modified in-place.
        skypix_bounds_builder : `~prerequisite_helpers.SkyPixBoundsBuilder`
            An object that accumulates the appropriate spatial bounds for a
            quantum.
        timespan_builder : `~prerequisite_helpers.TimespanBuilder`
            An object that accumulates the appropriate timespan for a quantum.

        Returns
        -------
        outputs : `~lsst.daf.butler.NamedKeyDict` [ \
                `~lsst.daf.butler.DatasetType`, `list` [ \
                `~lsst.daf.butler.DatasetRef` ] ]
            All outputs to the task, using the storage class and components
            defined by the task's own connections.

        Notes
        -----
        This first looks for outputs already present in the `output_run` by
        looking in `ExistingDatasets.outputs_in_the_way`; if it finds something
        and `clobber` is `True`, it uses that ref (it's not ideal that both the
        original dataset and its replacement will have the same UUID, but we
        don't have space in the quantum graph for two UUIDs, and we need the
        datastore records of the original there).  If `clobber` is `False`,
        `RuntimeError` is raised.  If there is no output already present, a new
        one with a random UUID is generated.  In all cases the "ref" attribute
        of the dataset node in the skeleton is set.
        """
        outputs_by_type: dict[str, list[DatasetRef]] = {}
        dataset_key: DatasetKey
        for dataset_key in skeleton.iter_outputs_of(quantum_key):
            dataset_data_id = skeleton[dataset_key]["data_id"]
            dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
            if (ref := self.existing_datasets.outputs_in_the_way.get(dataset_key)) is None:
                ref = DatasetRef(dataset_type_node.dataset_type, dataset_data_id, run=self.output_run)
            elif not self.clobber:
                # We intentionally raise here, before running adjustQuantum,
                # because it'd be weird if we left an old potential output of a
                # task sitting there in the output collection, just because the
                # task happened to not actually produce it.
                raise OutputExistsError(
                    f"Potential output dataset {ref} already exists in the output run "
                    f"{self.output_run}, but clobbering outputs was not expected to be necessary."
                )
            skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
            timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
            skeleton[dataset_key]["ref"] = ref
            outputs_by_type.setdefault(dataset_key.parent_dataset_type_name, []).append(ref)
        adapted_outputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
        for write_edge in task_node.iter_all_outputs():
            dataset_type_node = self._pipeline_graph.dataset_types[write_edge.parent_dataset_type_name]
            edge_dataset_type = write_edge.adapt_dataset_type(dataset_type_node.dataset_type)
            adapted_outputs[edge_dataset_type] = [
                write_edge.adapt_dataset_ref(ref)
                for ref in sorted(outputs_by_type.get(write_edge.parent_dataset_type_name, []))
            ]
        return adapted_outputs

    @final
    def _gather_quantum_inputs(
        self,
        task_node: TaskNode,
        quantum_key: QuantumKey,
        skeleton: QuantumGraphSkeleton,
        task_prerequisite_info: PrerequisiteInfo,
        skypix_bounds_builder: SkyPixBoundsBuilder,
        timespan_builder: TimespanBuilder,
    ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
        """Collect input datasets for a preliminary quantum and put them in the
        form used by `~lsst.daf.butler.Quantum` and
        `~PipelineTaskConnections.adjustQuantum`.

        Parameters
        ----------
        task_node : `pipeline_graph.TaskNode`
            Node for this task in the pipeline graph.
        quantum_key : `QuantumKey`
            Identifier for this quantum in the graph.
        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
            Preliminary quantum graph, to be modified in-place.
        skypix_bounds_builder : `~prerequisite_helpers.SkyPixBoundsBuilder`
            An object that accumulates the appropriate spatial bounds for a
            quantum.
        timespan_builder : `~prerequisite_helpers.TimespanBuilder`
            An object that accumulates the appropriate timespan for a quantum.

        Returns
        -------
        inputs : `~lsst.daf.butler.NamedKeyDict` [ \
                `~lsst.daf.butler.DatasetType`, `list` [ \
                `~lsst.daf.butler.DatasetRef` ] ]
            All regular and prerequisite inputs to the task, using the storage
            class and components defined by the task's own connections.

        Notes
        -----
        On return, the dataset nodes that represent inputs to this quantum will
        either have their "ref" attribute set (using the common dataset type,
        not the task-specific one) or will be removed from the graph.

        For regular inputs, usually an existing "ref" (corresponding to an
        output of another quantum) will be found and left unchanged.  When
        there is no existing "ref" attribute, `ExistingDatasets.inputs` is
        searched next; if there is nothing there, the input will be removed.

        Prerequisite inputs are always queried for directly here (delegating to
        `_find_prerequisite_inputs`).  They are never produced by other tasks,
        and cannot in general be queried for in advance when
        `ExistingDatasets.inputs` is populated.
        """
        quantum_data_id = skeleton[quantum_key]["data_id"]
        inputs_by_type: dict[str, set[DatasetRef]] = {}
        dataset_key: DatasetKey | PrerequisiteDatasetKey
        # Process inputs already present in the skeleton - this should include
        # all regular inputs (including intermediates) and may include some
        # prerequisites.
        for dataset_key in list(skeleton.iter_inputs_of(quantum_key)):
            if (ref := skeleton[dataset_key].get("ref")) is None:
                # This dataset is an overall input - if it was an intermediate,
                # we would have already either removed the node or set the
                # "ref" attribute when processing its producing quantum - and
                # this is the first time we're trying to resolve it.
                if (ref := self.existing_datasets.inputs.get(dataset_key)) is None:
                    # It also doesn't exist in the input collections, so we
                    # remove its node in the skeleton graph (so other consumers
                    # won't have to check for it).
                    skeleton.remove_dataset_nodes([dataset_key])
                    continue
                skeleton[dataset_key]["ref"] = ref
            inputs_by_type.setdefault(dataset_key.parent_dataset_type_name, set()).add(ref)
            skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
            timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
        # Query for any prerequisites not handled by process_subgraph.  Note
        # that these were not already in the skeleton graph, so we add them
        # now.
        skypix_bounds = skypix_bounds_builder.finish()
        timespan = timespan_builder.finish()
        for finder in task_prerequisite_info.finders.values():
            inputs_for_type = inputs_by_type.setdefault(finder.dataset_type_node.name, set())
            dataset_keys = []
            for ref in finder.find(
                self.butler, self.input_collections, quantum_data_id, skypix_bounds, timespan
            ):
                dataset_key = skeleton.add_prerequisite_node(ref.datasetType.name, ref=ref)
                dataset_keys.append(dataset_key)
                inputs_for_type.add(ref)
            skeleton.add_input_edges(quantum_key, dataset_keys)
        adapted_inputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
        for read_edge in task_node.iter_all_inputs():
            dataset_type_node = self._pipeline_graph.dataset_types[read_edge.parent_dataset_type_name]
            edge_dataset_type = read_edge.adapt_dataset_type(dataset_type_node.dataset_type)
            if (current_dataset_type := adapted_inputs.keys().get(edge_dataset_type.name)) is None:
                adapted_inputs[edge_dataset_type] = [
                    read_edge.adapt_dataset_ref(ref)
                    for ref in sorted(inputs_by_type.get(read_edge.parent_dataset_type_name, frozenset()))
                ]
            elif current_dataset_type != edge_dataset_type:
                raise NotImplementedError(
                    f"Task {task_node.label!r} has {edge_dataset_type.name!r} as an input via "
                    "two different connections, with two different storage class overrides. "
                    "This is not yet supported due to limitations in the Quantum data structure."
                )
            # If neither the `if` nor the `elif` above match, it means
            # multiple input connections have exactly the same dataset
            # type, and hence nothing to do after the first one.
        return adapted_inputs

    @final
    def _resolve_task_init(
        self, task_node: TaskNode, skeleton: QuantumGraphSkeleton, has_skipped_quanta: bool
    ) -> None:
        """Add init-input and init-output dataset nodes and edges for a task to
        the skeleton.

        Parameters
        ----------
        task_node : `pipeline_graph.TaskNode`
            Pipeline graph description of the task.
        skeleton : `QuantumGraphSkeleton`
            In-progress quantum graph data structure to update in-place.
        has_skipped_quanta : `bool`
            Whether any of this task's quanta were skipped because they had
            already succeeded.
        """
        quanta = skeleton.get_quanta(task_node.label)
        task_init_key = TaskInitKey(task_node.label)
        if quanta:
            adapted_inputs: NamedKeyDict[DatasetType, DatasetRef] = NamedKeyDict()
            # Process init-inputs.
            input_keys: list[DatasetKey] = []
            for read_edge in task_node.init.iter_all_inputs():
                dataset_key = skeleton.add_dataset_node(
                    read_edge.parent_dataset_type_name, self.empty_data_id
                )
                skeleton.add_input_edge(task_init_key, dataset_key)
                if (ref := skeleton[dataset_key].get("ref")) is None:
                    try:
                        ref = self.existing_datasets.inputs[dataset_key]
                    except KeyError:
                        raise InitInputMissingError(
                            f"Overall init-input dataset {read_edge.parent_dataset_type_name!r} "
                            f"needed by task {task_node.label!r} not found in input collection(s) "
                            f"{self.input_collections}."
                        ) from None
                    skeleton[dataset_key]["ref"] = ref
                for quantum_key in skeleton.get_quanta(task_node.label):
                    skeleton.add_input_edge(quantum_key, dataset_key)
                input_keys.append(dataset_key)
                adapted_ref = read_edge.adapt_dataset_ref(ref)
                adapted_inputs[adapted_ref.datasetType] = adapted_ref
            # Save the quantum-adapted init inputs to each quantum, and add
            # skeleton edges connecting the init inputs to each quantum.
            for quantum_key in skeleton.get_quanta(task_node.label):
                skeleton[quantum_key]["init_inputs"] = adapted_inputs
            # Process init-outputs.
            adapted_outputs: NamedKeyDict[DatasetType, DatasetRef] = NamedKeyDict()
            for write_edge in task_node.init.iter_all_outputs():
                dataset_key = skeleton.add_dataset_node(
                    write_edge.parent_dataset_type_name, self.empty_data_id
                )
                if (ref := self.existing_datasets.outputs_in_the_way.get(dataset_key)) is None:
                    ref = DatasetRef(
                        self._pipeline_graph.dataset_types[write_edge.parent_dataset_type_name].dataset_type,
                        self.empty_data_id,
                        run=self.output_run,
                    )
                skeleton[dataset_key]["ref"] = ref
                skeleton.add_output_edge(task_init_key, dataset_key)
                adapted_ref = write_edge.adapt_dataset_ref(ref)
                adapted_outputs[adapted_ref.datasetType] = adapted_ref
            skeleton[task_init_key]["inputs"] = adapted_inputs
            skeleton[task_init_key]["outputs"] = adapted_outputs
        elif has_skipped_quanta:
            # No quanta remain for this task, but at least one quantum was
            # skipped because its outputs were present in the skip_existing_in
            # collections.  This means all init outputs should be present in
            # the skip_existing_in collections, too, and we need to put those
            # refs in the graph.
            for write_edge in task_node.init.iter_all_outputs():
                dataset_key = skeleton.add_dataset_node(
                    write_edge.parent_dataset_type_name, self.empty_data_id
                )
                if (ref := self.existing_datasets.outputs_for_skip.get(dataset_key)) is None:
                    raise InitInputMissingError(
                        f"Init-output dataset {write_edge.parent_dataset_type_name!r} of skipped task "
                        f"{task_node.label!r} not found in skip-existing-in collection(s) "
                        f"{self.skip_existing_in}."
                    ) from None
                skeleton[dataset_key]["ref"] = ref
                # If this dataset was "in the way" (i.e. already in the output
                # run), it isn't anymore.
                self.existing_datasets.outputs_in_the_way.pop(dataset_key, None)
        # No quanta remain in this task, but none were skipped; this means
        # they all got pruned because of NoWorkFound conditions.  This
        # dooms all downstream quanta to the same fate, so we don't bother
        # doing anything with the task's init-outputs, since nothing is
        # going to consume them.

    @final
    @timeMethod
    def _find_empty_dimension_datasets(self) -> None:
        """Query for all dataset types with no dimensions, updating
        `existing_datasets` in-place.

        This includes but is not limited to init inputs and init outputs.
        """
        _, dataset_type_nodes = self._pipeline_graph.group_by_dimensions()[self.universe.empty.as_group()]
        dataset_types = [node.dataset_type for node in dataset_type_nodes.values()]
        dataset_types.extend(self._global_init_output_types.values())
        for dataset_type in dataset_types:
            key = DatasetKey(dataset_type.name, self.empty_data_id.required_values)
            if (
                self._pipeline_graph.producer_of(dataset_type.name) is None
                and dataset_type.name not in self._global_init_output_types
            ):
                # Dataset type is an overall input; we always need to try to
                # find these.
                try:
                    ref = self.butler.find_dataset(dataset_type.name, collections=self.input_collections)
                except MissingDatasetTypeError:
                    ref = None
                if ref is not None:
                    self.existing_datasets.inputs[key] = ref
            elif self.skip_existing_in:
                # Dataset type is an intermediate or output; need to find these
                # if only they're from previously executed quanta that we might
                # skip...
                try:
                    ref = self.butler.find_dataset(dataset_type.name, collections=self.skip_existing_in)
                except MissingDatasetTypeError:
                    ref = None
                if ref is not None:
                    self.existing_datasets.outputs_for_skip[key] = ref
                    if ref.run == self.output_run:
                        self.existing_datasets.outputs_in_the_way[key] = ref
            if self.output_run_exists and not self.skip_existing_starts_with_output_run:
                # ...or if they're in the way and would need to be clobbered
                # (and we haven't already found them in the previous block).
                try:
                    ref = self.butler.find_dataset(dataset_type.name, collections=[self.output_run])
                except MissingDatasetTypeError:
                    ref = None
                if ref is not None:
                    self.existing_datasets.outputs_in_the_way[key] = ref

    @final
    @timeMethod
    def _attach_datastore_records(self, skeleton: QuantumGraphSkeleton) -> None:
        """Add datastore records for all overall inputs to a preliminary
        quantum graph.

        Parameters
        ----------
        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
            Preliminary quantum graph to update in place.

        Notes
        -----
        On return, all quantum nodes in the skeleton graph will have a
        "datastore_records" attribute that is a mapping from datastore name
        to `lsst.daf.butler.DatastoreRecordData`, as used by
        `lsst.daf.butler.Quantum`.
        """
        overall_inputs = skeleton.extract_overall_inputs()
        exported_records = self.butler._datastore.export_records(overall_inputs.values())
        for quantum_key in skeleton.iter_all_quanta():
            quantum_records = {}
            input_ids = {
                ref.id
                for dataset_key in skeleton.iter_inputs_of(quantum_key)
                if (ref := overall_inputs.get(dataset_key)) is not None
            }
            if input_ids:
                for datastore_name, records in exported_records.items():
                    matching_records = records.subset(input_ids)
                    if matching_records is not None:
                        quantum_records[datastore_name] = matching_records
            skeleton[quantum_key]["datastore_records"] = quantum_records

    @final
    @timeMethod
    def _construct_quantum_graph(
        self, skeleton: QuantumGraphSkeleton, metadata: Mapping[str, Any]
    ) -> QuantumGraph:
        """Construct a `QuantumGraph` object from the contents of a
        fully-processed `quantum_graph_skeleton.QuantumGraphSkeleton`.

        Parameters
        ----------
        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
            Preliminary quantum graph.  Must have "init_inputs", "inputs", and
            "outputs" attributes on all quantum nodes, as added by
            `_resolve_task_quanta`, as well as a "datastore_records" attribute
            as added by `_attach_datastore_records`.
        metadata : `Mapping`
            Flexible metadata to add to the graph.

        Returns
        -------
        quantum_graph : `QuantumGraph`
            DAG describing processing to be performed.
        """
        quanta: dict[TaskDef, set[Quantum]] = {}
        init_inputs: dict[TaskDef, Iterable[DatasetRef]] = {}
        init_outputs: dict[TaskDef, Iterable[DatasetRef]] = {}
        for task_def in self._pipeline_graph._iter_task_defs():
            if not skeleton.has_task(task_def.label):
                continue
            task_node = self._pipeline_graph.tasks[task_def.label]
            task_init_key = skeleton.get_task_init_node(task_def.label)
            init_inputs[task_def] = skeleton[task_init_key]["inputs"].values()
            init_outputs[task_def] = skeleton[task_init_key]["outputs"].values()
            quanta_for_task: set[Quantum] = set()
            for quantum_key in skeleton.get_quanta(task_node.label):
                node_state = skeleton[quantum_key]
                quanta_for_task.add(
                    Quantum(
                        taskName=task_node.task_class_name,
                        taskClass=task_node.task_class,
                        dataId=node_state["data_id"],
                        initInputs=node_state["init_inputs"],
                        inputs=node_state["inputs"],
                        outputs=node_state["outputs"],
                        datastore_records=node_state.get("datastore_records"),
                    )
                )
            quanta[task_def] = quanta_for_task

        registry_dataset_types: list[DatasetType] = [
            node.dataset_type for node in self._pipeline_graph.dataset_types.values()
        ]

        all_metadata = self.metadata.to_dict()
        all_metadata.update(metadata)
        return QuantumGraph(
            quanta,
            metadata=all_metadata,
            universe=self.universe,
            initInputs=init_inputs,
            initOutputs=init_outputs,
            globalInitOutputs=[skeleton[key]["ref"] for key in skeleton.global_init_outputs],
            registryDatasetTypes=registry_dataset_types,
        )

    @staticmethod
    @final
    def _find_removed(
        original: Iterable[DatasetKey | PrerequisiteDatasetKey],
        adjusted: NamedKeyMapping[DatasetType, Sequence[DatasetRef]],
    ) -> set[DatasetKey | PrerequisiteDatasetKey]:
        """Identify skeleton-graph dataset nodes that have been removed by
        `~PipelineTaskConnections.adjustQuantum`.

        Parameters
        ----------
        original : `~collections.abc.Iterable` [ `DatasetKey` or \
                `PrerequisiteDatasetKey` ]
            Identifiers for the dataset nodes that were the original neighbors
            (inputs or outputs) of a quantum.
        adjusted : `~lsst.daf.butler.NamedKeyMapping` [ \
                `~lsst.daf.butler.DatasetType`, \
                `~collections.abc.Sequence` [ `lsst.daf.butler.DatasetType` ] ]
            Adjusted neighbors, in the form used by `lsst.daf.butler.Quantum`.

        Returns
        -------
        removed : `set` [ `DatasetKey` ]
            Datasets in ``original`` that have no counterpart in ``adjusted``.
        """
        result = set(original)
        for dataset_type, kept_refs in adjusted.items():
            parent_dataset_type_name, _ = DatasetType.splitDatasetTypeName(dataset_type.name)
            for kept_ref in kept_refs:
                # We don't know if this was a DatasetKey or a
                # PrerequisiteDatasetKey; just try both.
                result.discard(DatasetKey(parent_dataset_type_name, kept_ref.dataId.required_values))
                result.discard(PrerequisiteDatasetKey(parent_dataset_type_name, kept_ref.id.bytes))
        return result


@dataclasses.dataclass(eq=False, order=False)
class ExistingDatasets:
    """Struct that holds the results of dataset queries for
    `QuantumGraphBuilder`.
    """

    inputs: dict[DatasetKey | PrerequisiteDatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
    """Overall-input datasets found in `QuantumGraphBuilder.input_collections`.

    This may include prerequisite inputs.  It does include init-inputs.
    It does not include intermediates.
    """

    outputs_for_skip: dict[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
    """Output datasets found in `QuantumGraphBuilder.skip_existing_in`.

    It is unspecified whether this contains include init-outputs; there is
    no concept of skipping at the init stage, so this is not expected to
    matter.
    """

    outputs_in_the_way: dict[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
    """Output datasets found in `QuantumGraphBuilder.output_run`.

    This includes regular outputs and init-outputs.
    """


def _quantum_or_quanta(n: int) -> str:
    """Correctly pluralize 'quantum' if needed."""
    return f"{n} quanta" if n != 1 else "1 quantum"
