# 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/>.

"""Utility to facilitate testing of pipelines consisting of multiple steps."""

__all__ = ["PipelineStepTester"]

import dataclasses
import unittest
import warnings

from lsst.daf.butler import Butler, DatasetType
from lsst.pipe.base import Pipeline, PipelineGraph
from lsst.resources import ResourcePath, ResourcePathExpression


@dataclasses.dataclass
class PipelineStepTester:
    """Utility class which facilitates testing of pipelines, optionally
    consisting of multiple steps.

    Two sets will be constructed by looping over the entire pipeline or all
    named subsets within the pipeline: `pure_inputs` and `all_outputs`.

    The `pure_inputs` set consists of all inputs which must be constructed and
    provided as an input into the pipeline (i.e., they will not be generated
    by the named pipeline).

    The `all_outputs` set consists of all dataset types which are generated by
    the named pipeline, either as intermediates or as final outputs.

    These sets will be checked against user-supplied sets to ensure that the
    named pipeline may still be run without raising a missing data error.

    Attributes
    ----------
    filename : `lsst.resources.ResourcePathExpression`
        The full path or URI to the pipeline YAML.
    step_suffixes : `list` [`str`]
        A list, in the order of data reduction, of the step subsets to check.
        Must include any initial "#".
    initial_dataset_types : `list` [`tuple`]
        Dataset types which require initial registry by the butler. Each
        element must be a tuple of the type name, dimensions, storage class,
        and calibration flag, as described in the constructor for
        `~lsst.daf.butler.DatasetType`. All tuple elements are required.
    expected_inputs : `set` [`str`]
        Dataset types expected as an input into the pipeline.
    expected_outputs : `set` [`str`]
        Dataset types expected to be produced as an output by the pipeline.
    pipeline_patches : `dict` [`str`, `str`], optional
        Any config overrides to be applied to the pipeline before testing it.
        This is rarely appropriate, and should be reserved for fields that
        cannot have a file-level default. The key must have the form
        "task:subtask.field".
    """

    filename: ResourcePathExpression
    step_suffixes: list[str]
    initial_dataset_types: list[tuple[str, set[str], str, bool]]
    expected_inputs: set[str]
    expected_outputs: set[str]
    pipeline_patches: dict[str, str] = dataclasses.field(default_factory=dict)

    def register_dataset_types(self, butler: Butler) -> None:
        """Register any dataset types passed to the class constructor.

        The types registered are those specified in
        ``self.initial_dataset_types``.

        Parameters
        ----------
        butler : `lsst.daf.butler.Butler`
            The (test) butler in which to register the types.
        """
        for name, dimensions, storageClass, isCalibration in self.initial_dataset_types:
            butler.registry.registerDatasetType(
                DatasetType(
                    name,
                    dimensions,
                    storageClass=storageClass,
                    isCalibration=isCalibration,
                    universe=butler.dimensions,
                )
            )

    def run(self, butler: Butler, test_case: unittest.TestCase) -> None:
        """Run the test on all pipelines passed to the class constructor.

        Parameters
        ----------
        butler : `lsst.daf.butler.Butler`
            The butler in which to run the test. The butler must be writeable,
            and its state will be modified as part of the test.
        test_case : `unittest.TestCase`
            The test case in which this test is run.

        Raises
        ------
        AssertionError
            Raised if the pipeline requires inputs that are not in
            ``self.expected_inputs``, or fails to produce outputs that are in
            ``self.expected_outputs``.
        """
        self.register_dataset_types(butler)

        all_outputs: dict[str, DatasetType] = dict()
        pure_inputs: dict[str, str] = dict()

        file_uri = ResourcePath(self.filename, forceDirectory=False)
        for suffix in self.step_suffixes:
            step_uri = file_uri.replace(fragment=suffix.removeprefix("#"))
            step_graph = self.load_pipeline_graph(step_uri)
            step_graph.resolve(butler.registry)

            for name, _ in step_graph.iter_overall_inputs():
                if name not in all_outputs:
                    tasks = step_graph.consumers_of(name)
                    pure_inputs[name] = f"{suffix}: {', '.join(task.label for task in tasks)}"
            all_outputs.update(
                {
                    name: node.dataset_type
                    for name, node in step_graph.dataset_types.items()
                    if step_graph.producing_edge_of(name) is not None
                }
            )

            for node in step_graph.dataset_types.values():
                butler.registry.registerDatasetType(node.dataset_type)

        if not pure_inputs.keys() <= self.expected_inputs:
            missing = []
            for type_name in pure_inputs.keys() - self.expected_inputs:
                suffix = pure_inputs[type_name]
                missing.append(type_name + (f" ({suffix})" if suffix else ""))
            raise AssertionError(f"Got unexpected pure_inputs: {missing}")

        if not all_outputs.keys() >= self.expected_outputs:
            missing = list(self.expected_outputs - all_outputs.keys())
            raise AssertionError(f"Missing expected_outputs: {missing}")

    def load_pipeline_graph(self, uri: ResourcePathExpression) -> PipelineGraph:
        pipeline = Pipeline.from_uri(uri)
        for fullkey, value in self.pipeline_patches.items():
            label, key = fullkey.split(":", maxsplit=1)
            try:
                pipeline.addConfigOverride(label, key, value)
            except LookupError as e:
                warnings.warn(f"{e}, skipping configuration {fullkey}={value}", UserWarning)
        return pipeline.to_graph()
