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

import functools
import operator
from collections.abc import Iterable
from typing import Any

import click

from lsst.daf.butler.cli.opt import (
    dataset_type_option,
    options_file_option,
    register_dataset_types_option,
    repo_argument,
    transfer_dimensions_option,
    transfer_option,
)
from lsst.daf.butler.cli.utils import ButlerCommand, split_commas, unwrap

from ... import script
from ..._status import QuantumAttemptStatus, QuantumSuccessCaveats
from ...quantum_graph import aggregator
from ..opt import instrument_argument, update_output_chain_option


@click.command(short_help="Add an instrument definition to the repository", cls=ButlerCommand)
@repo_argument(required=True)
@instrument_argument(required=True, nargs=-1, help="The fully-qualified name of an Instrument subclass.")
@click.option("--update", is_flag=True)
def register_instrument(*args: Any, **kwargs: Any) -> None:
    """Add an instrument to the data repository."""
    script.register_instrument(*args, **kwargs)


@click.command(short_help="Transfer datasets from a graph to a butler.", cls=ButlerCommand)
@click.argument("graph", required=True)
@click.argument("dest", required=True)
@register_dataset_types_option()
@transfer_dimensions_option(default=False)
@update_output_chain_option()
@click.option(
    "--dry-run", is_flag=True, default=False, help="Run the transfer but do not update the destination butler"
)
@dataset_type_option(help="Subset of dataset types to transfer from graph.")
@options_file_option()
def transfer_from_graph(**kwargs: Any) -> None:
    """Transfer datasets from a quantum graph to a destination butler.

    GRAPH is a URI to the source quantum graph file.

    DEST is a URI to the Butler repository that will receive copies of the
    datasets.
    """
    number = script.transfer_from_graph(**kwargs)
    print(f"Number of datasets transferred: {number}")


@click.command(short_help="Make Zip archive from output files using graph.", cls=ButlerCommand)
@click.argument("graph", required=True)
@repo_argument(
    required=True,
    help="REPO is a URI to a butler configuration that is used to configure "
    "the datastore of the quantum-backed butler.",
)
@click.argument("dest", required=True)
@dataset_type_option(help="Dataset types to include in Zip archive.")
@options_file_option()
def zip_from_graph(**kwargs: Any) -> None:
    """Transfer datasets from a quantum graph to a Zip archive.

    GRAPH is a URI to the source quantum graph file to use when building the
    Zip archive.

    DEST is a directory to write the Zip archive.
    """
    zip = script.zip_from_graph(**kwargs)
    print(f"Zip archive written to {zip}")


@click.command(short_help="Retrieve artifacts from subset of graph.", cls=ButlerCommand)
@click.argument("graph", required=True)
@repo_argument(
    required=True,
    help="REPO is a URI to a butler configuration that is used to configure "
    "the datastore of the quantum-backed butler.",
)
@click.argument("dest", required=True)
@transfer_option()
@click.option(
    "--preserve-path/--no-preserve-path",
    is_flag=True,
    default=True,
    help="Preserve the datastore path to the artifact at the destination.",
)
@click.option(
    "--clobber/--no-clobber",
    is_flag=True,
    default=False,
    help="If clobber, overwrite files if they exist locally.",
)
@click.option(
    "--qgraph-node-id",
    callback=split_commas,
    multiple=True,
    help=unwrap(
        """Only load a specified set of nodes when graph is
        loaded from a file, nodes are identified by UUID
        values. One or more comma-separated strings are
        accepted. By default all nodes are loaded. Ignored if
        graph is not loaded from a file."""
    ),
)
@click.option(
    "--include-inputs/--no-include-inputs",
    is_flag=True,
    default=True,
    help="Whether to include input datasets in retrieval.",
)
@click.option(
    "--include-outputs/--no-include-outputs",
    is_flag=True,
    default=True,
    help="Whether to include output datasets in retrieval.",
)
@options_file_option()
def retrieve_artifacts_for_quanta(**kwargs: Any) -> None:
    """Retrieve artifacts from given quanta defined in quantum graph.

    GRAPH is a URI to the source quantum graph file to use when building the
    Zip archive.

    DEST is a directory to write the Zip archive.
    """
    artifacts = script.retrieve_artifacts_for_quanta(**kwargs)
    print(f"Written {len(artifacts)} artifacts to {kwargs['dest']}.")


_AGGREGATOR_DEFAULTS = aggregator.AggregatorConfig()


@click.command(short_help="Scan for the outputs of an active or completed quantum graph.", cls=ButlerCommand)
@click.argument("predicted_graph", required=True)
@repo_argument(required=True, help="Path or alias for the butler repository.")
@click.option(
    "-o",
    "--output",
    "output_path",
    default=_AGGREGATOR_DEFAULTS.output_path,
    help=(
        "Path to the output provenance quantum graph.  THIS OPTION IS FOR "
        "DEVELOPMENT AND DEBUGGING ONLY.  IT MAY BE REMOVED IN THE FUTURE."
    ),
)
@click.option(
    "--processes",
    "-j",
    "n_processes",
    default=_AGGREGATOR_DEFAULTS.n_processes,
    type=click.IntRange(min=1),
    help="Number of processes to use.",
)
@click.option(
    "--incomplete/--complete",
    "incomplete",
    default=_AGGREGATOR_DEFAULTS.incomplete,
    help="Whether execution has completed (and failures cannot be retried).",
)
@click.option(
    "--dry-run",
    is_flag=True,
    default=_AGGREGATOR_DEFAULTS.dry_run,
    help="Do not actually perform any central database ingests.",
)
@click.option(
    "--interactive-status/--no-interactive-status",
    "interactive_status",
    default=_AGGREGATOR_DEFAULTS.interactive_status,
    help="Use progress bars for status reporting instead of periodic logging.",
)
@click.option(
    "--log-status-interval",
    type=int,
    default=_AGGREGATOR_DEFAULTS.log_status_interval,
    help="Interval (in seconds) between periodic logger status updates.",
)
@click.option(
    "--register-dataset-types/--no-register-dataset-types",
    default=_AGGREGATOR_DEFAULTS.register_dataset_types,
    help="Register output dataset types.",
)
@click.option(
    "--update-output-chain/--no-update-output-chain",
    default=_AGGREGATOR_DEFAULTS.update_output_chain,
    help="Prepend the output RUN collection to the output CHAINED collection.",
)
@click.option(
    "--worker-log-dir",
    type=str,
    default=_AGGREGATOR_DEFAULTS.worker_log_dir,
    help="Path to a directory (POSIX only) for parallel worker logs.",
)
@click.option(
    "--worker-log-level",
    type=str,
    default=_AGGREGATOR_DEFAULTS.worker_log_level,
    help="Log level for worker processes/threads (use DEBUG for per-quantum messages).",
)
@click.option(
    "--zstd-level",
    type=int,
    default=_AGGREGATOR_DEFAULTS.zstd_level,
    help="Compression level for the provenance quantum graph file.",
)
@click.option(
    "--zstd-dict-size",
    type=int,
    default=_AGGREGATOR_DEFAULTS.zstd_dict_size,
    help="Size (in bytes) of the ZStandard compression dictionary.",
)
@click.option(
    "--zstd-dict-n-inputs",
    type=int,
    default=_AGGREGATOR_DEFAULTS.zstd_dict_n_inputs,
    help=("Number of samples of each type to include in ZStandard compression dictionary training."),
)
@click.option(
    "--mock-storage-classes/--no-mock-storage-classes",
    default=_AGGREGATOR_DEFAULTS.mock_storage_classes,
    help="Enable support for storage classes created by the lsst.pipe.base.tests.mocks package.",
)
@click.option(
    "--promise-ingest-graph/--no-promise-ingest-graph",
    default=_AGGREGATOR_DEFAULTS.promise_ingest_graph,
    help=(
        "Promise to run 'butler ingest-graph' later, allowing aggregate-graph "
        "to skip metadata/log/config ingestion for now."
    ),
)
def aggregate_graph(predicted_graph: str, repo: str, **kwargs: Any) -> None:
    """Scan for quantum graph's outputs to gather provenance, ingest datasets
    into the central butler repository, and delete datasets that are no
    longer needed.
    """
    # It'd be nice to allow to the user to provide a path to an
    # AggregatorConfig JSON file for options that weren't provided, but Click
    # 8.1 fundamentally cannot handle flag options that default to None rather
    # than True or False (i.e. so they fall back to the config value when not
    # set).  It's not clear whether Click 8.2.x has actually fixed this; Click
    # 8.2.0 tried but caused new problems.

    config = aggregator.AggregatorConfig(**kwargs)
    try:
        aggregator.aggregate_graph(predicted_graph, repo, config)
    except aggregator.FatalWorkerError as err:
        # When this exception is raised, we'll have already logged the relevant
        # traceback from a separate worker.
        raise click.ClickException(str(err)) from None


@click.command(
    short_help="Ingest a provenance quantum graph into a butler.",
    cls=ButlerCommand,
)
@repo_argument(required=True, help="Path or alias for the butler repository.")
@click.argument("provenance_graph", required=False)
@transfer_option(default="move")
@click.option("--batch-size", default=10000, help="How many datasets to process in each transaction.")
@click.option(
    "--output-run",
    default=None,
    help=(
        "Name of the output RUN collection.  Must be provided if the provenance graph is not"
        " provided (so the graph can be found in the butler)."
    ),
)
def ingest_graph(
    *,
    repo: str,
    provenance_graph: str | None,
    transfer: str | None,
    batch_size: int,
    output_run: str | None,
) -> None:
    """Ingest a provenance graph into a butler repository."""
    from ...quantum_graph.ingest_graph import ingest_graph as ingest_graph_py

    ingest_graph_py(repo, provenance_graph, transfer=transfer, batch_size=batch_size, output_run=output_run)


@click.command(
    short_help="Print and write provenance reports.",
    cls=ButlerCommand,
)
@click.argument("repo_or_qg")
@click.argument("collection", required=False, default=None)
@click.option(
    "--state",
    multiple=True,
    type=click.Choice(QuantumAttemptStatus),
    help=(
        "Additional quantum state to include in the status report and data ID tables "
        "(FAILED, ABORTED, and ABORTED_SUCCESS are included by default)."
    ),
)
@click.option(
    "--no-state",
    multiple=True,
    type=str,
    metavar="STATE",
    help="Quantum state to drop from in status report and data ID tables (same options as --state).",
)
@click.option(
    "--status-report",
    default=None,
    metavar="URI",
    help="File or URI (.json) for a detailed report (with data IDs) on quanta with certain states.",
)
@click.option(
    "--quantum-table/--no-quantum-table",
    default=True,
    help="Whether to print summary of quantum status counts to STDOUT.",
)
@click.option(
    "--exception-table/--no-exception-table",
    default=True,
    help="Whether to print summary of exception type counts STDOUT.",
)
@click.option(
    "--caveat",
    multiple=True,
    type=click.Choice(QuantumSuccessCaveats),
    help=(
        "Include successful quanta in the status report if they have this caveat. "
        "May be passed multiple times; any matching caveat is included. "
        "Passing this option implicitly adds '--state SUCCESSFUL'."
    ),
)
@click.option(
    "--data-id-table-dir",
    default=None,
    metavar="URI",
    help=(
        "Directory (may be a URI) for a tree of data ID tables for each "
        "task label, status, and exception type combination in the status report."
    ),
)
def provenance_report(
    *,
    repo_or_qg: str,
    collection: str | None,
    state: Iterable[QuantumAttemptStatus],
    no_state: Iterable[str],
    status_report: str | None,
    quantum_table: bool = False,
    exception_table: bool = False,
    caveat: Iterable[QuantumSuccessCaveats],
    data_id_table_dir: str | None,
) -> None:
    """Read a provenance quantum graph from a butler or file and use it to
    generate reports.

    REPO_OR_QG is a path or alias for the butler repository (if reading an
    ingested graph, as indicated by passing COLLECTION), or the path to a
    provenance quantum graph file.
    """
    from ...quantum_graph import ProvenanceQuantumGraph

    states = set(state)
    states.add(QuantumAttemptStatus.FAILED)
    states.add(QuantumAttemptStatus.ABORTED)
    states.add(QuantumAttemptStatus.ABORTED_SUCCESS)
    for state_name in no_state:
        states.discard(QuantumAttemptStatus.__members__[state_name])
    with_caveats: QuantumSuccessCaveats | None = None
    if caveat:
        states.add(QuantumAttemptStatus.SUCCESSFUL)
        with_caveats = functools.reduce(
            operator.__or__,
            caveat,
            QuantumSuccessCaveats.NO_CAVEATS,
        )
    with ProvenanceQuantumGraph.from_args(repo_or_qg, collection=collection, datasets=()) as (graph, _):
        graph.make_many_reports(
            status_report_file=status_report,
            states=states,
            print_quantum_table=quantum_table,
            print_exception_table=exception_table,
            with_caveats=with_caveats,
            data_id_table_dir=data_id_table_dir,
        )
