# Copyright 2024 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.
"""
Base Workflow infrastructure.

See :ref:`explanation-workflows` for a high level explanation of concepts used
here.
"""
import logging
from abc import ABCMeta, abstractmethod
from collections.abc import Iterable
from typing import Any, TypeVar, TypedDict, cast

from django.contrib.auth.models import AnonymousUser

from debusine.artifacts.models import (
    ArtifactCategory,
    BareDataCategory,
    DebianBinaryPackage,
    DebianBinaryPackages,
    DebianUpload,
)
from debusine.client.models import LookupChildType
from debusine.db.models import (
    CollectionItem,
    TaskDatabase,
    User,
    WorkRequest,
    Workspace,
)
from debusine.db.models.work_requests import InternalTaskError
from debusine.server.collections.lookup import (
    LookupResult,
    lookup_multiple,
    lookup_single,
    reconstruct_lookup,
)
from debusine.server.tasks.base import BaseServerTask
from debusine.server.workflows.models import (
    BaseWorkflowData,
    WorkRequestWorkflowData,
)
from debusine.tasks import TaskConfigError
from debusine.tasks.models import (
    ActionUpdateCollectionWithArtifacts,
    ActionUpdateCollectionWithData,
    BaseDynamicTaskData,
    LookupMultiple,
    LookupSingle,
    TaskTypes,
)
from debusine.tasks.server import TaskDatabaseInterface

logger = logging.getLogger(__name__)


WD = TypeVar("WD", bound=BaseWorkflowData)
DTD = TypeVar("DTD", bound=BaseDynamicTaskData)


class WorkflowValidationError(Exception):
    """Raised if a workflow fails to validate its inputs."""


class WorkflowRunError(Exception):
    """Running a workflow orchestrator or callback failed."""

    def __init__(self, work_request: WorkRequest, message: str) -> None:
        """Construct the exception."""
        self.work_request = work_request
        self.message = message


class ArtifactHasNoArchitecture(Exception):
    """Raised if it's not possible to determine the artifact's architecture."""


class Workflow(BaseServerTask[WD, DTD], metaclass=ABCMeta):
    """
    Base class for workflow orchestrators.

    This is the base API for running :class:`WorkflowInstance` logic
    """

    TASK_TYPE = TaskTypes.WORKFLOW

    work_request: WorkRequest
    workspace: Workspace

    def __init__(self, work_request: "WorkRequest"):
        """Instantiate a Workflow with its database instance."""
        super().__init__(
            task_data=work_request.task_data,
            dynamic_task_data=work_request.dynamic_task_data,
        )
        self.set_work_request(work_request)

    @classmethod
    def from_name(cls, name: str) -> type['Workflow[Any, Any]']:
        """Instantiate a workflow by name."""
        res = super().class_from_name(TaskTypes.WORKFLOW, name)
        return cast(type[Workflow[Any, Any]], res)

    @classmethod
    def validate_template_data(cls, data: dict[str, Any]) -> None:
        """Validate WorkflowTemplate data."""
        # By default nothing is enforced

    @classmethod
    def build_workflow_data(
        cls, template_data: dict[str, Any], user_data: dict[str, Any]
    ) -> dict[str, Any]:
        """Merge template data and user-provided data."""
        # Data from the template has priority over user-provided.
        #
        # By default, merging of lists and dictionaries is not supported: only
        # toplevel keys are considered, and if a list or a dict exists both in
        # the user-provided data and in the template, the user-provided version
        # is ignored.
        #
        # Orchestrator subclasses may implement different merging strategies.
        data = user_data.copy()
        data.update(template_data)
        return data

    def validate_input(self) -> None:
        """
        Thorough validation of input data.

        This is run only once at workflow instantiation time, and can do
        slower things like database lookups to validate artifact or collection
        types.
        """
        # Do nothing by default

    def ensure_dynamic_data(self, task_database: TaskDatabaseInterface) -> bool:
        """Ensure that this workflow's dynamic task data has been computed."""
        if self.work_request.dynamic_task_data is None:
            try:
                dynamic_task_data = self.compute_dynamic_data(task_database)
            except Exception as exc:
                self.logger.warning(  # noqa: G200
                    "Workflow orchestrator failed to pre-process task data for "
                    "WorkRequest %s. Task data: %s Error: %s",
                    self.work_request.id,
                    self.work_request.task_data,
                    exc,
                )
                return False

            if dynamic_task_data is not None:
                self.work_request.dynamic_task_data = dynamic_task_data.dict()
                self.work_request.save()
                self.dynamic_data = self.dynamic_task_data_type(
                    **self.work_request.dynamic_task_data
                )

        return True

    @abstractmethod
    def populate(self) -> None:
        """
        Create the initial WorkRequest structure.

        This is called once, when the workflow first becomes runnable.
        validate_input will already have been called.

        This method is required to be idempotent: calling it multiple times
        with the same argument MUST result in the same WorkRequest structure as
        calling it once
        """

    def callback(self, work_request: "WorkRequest") -> None:
        """
        Perform an orchestration step.

        Called with the workflow callback work request (note that this is
        not the same as `self.work_request`) when the workflow node becomes
        ready to execute.

        Called with a WorkRequest of type internal/workflow to perform an
        orchestration step triggered by a workflow callback.

        This method is required to be idempotent: calling it multiple times
        with the same argument MUST result in the same WorkRequest structure as
        calling it once
        """
        # Do nothing by default

    @staticmethod
    def provides_artifact(
        work_request: WorkRequest,
        category: ArtifactCategory,
        name: str,
        *,
        data: dict[str, Any] | None = None,
        artifact_filters: dict[str, Any] | None = None,  # noqa: U100
    ) -> None:
        r"""
        Indicate work_request will provide an artifact.

        :param work_request: work request that will provide the artifact
        :param category: category of the artifact that will be provided
        :param name: name of this item in the workflow’s internal collection
        :param data: add it to the data dictionary for the event reaction
        :param artifact_filters: for the ActionUpdateCollectionWithArtifacts to
          allow workflows to add filtering
        :raise LookupError: if a key in "data" starts with promise\_

        Create an event reaction for on_creation adding a promise: this
        work request will create an artifact.

        Create an event reaction for on_success to update the collection
        with the relevant artifact.
        """
        if data is not None:
            for key, value in data.items():
                if key.startswith("promise_"):
                    raise ValueError(
                        f'Field name "{key}" starting with '
                        f'"promise_" is not allowed.'
                    )

        # work_request is part of a workflow
        assert work_request.parent is not None

        work_request.process_update_collection_with_data(
            [
                ActionUpdateCollectionWithData(
                    collection="internal@collections",
                    name_template=name,
                    category=BareDataCategory.PROMISE,
                    data={
                        "promise_work_request_id": work_request.id,
                        "promise_workflow_id": work_request.parent.id,
                        "promise_category": category,
                        **(data or {}),
                    },
                )
            ]
        )

        event_reactions = work_request.event_reactions

        artifact_filters = artifact_filters or {}

        event_reactions.on_success.append(
            ActionUpdateCollectionWithArtifacts(
                collection="internal@collections",
                name_template=name,
                variables=data,
                artifact_filters={**artifact_filters, "category": category},
            )
        )

        work_request.event_reactions = event_reactions
        work_request.save()

    @staticmethod
    def requires_artifact(
        work_request: WorkRequest, lookup: LookupSingle | LookupMultiple
    ) -> None:
        """
        Indicate that work_request requires input (lookup).

        :param work_request: for each lookup result call
          work_request.add_dependency(promise["promise_work_request_id"])
        :param lookup: resolve the lookup and iterate over the results
          (for PROMISES only)
        """

        class LookupKwargs(TypedDict):
            """Dictionary containing keyword arguments for lookup functions."""

            workspace: Workspace
            user: User | AnonymousUser
            workflow_root: WorkRequest | None

        lookup_kwargs: LookupKwargs = {
            "workspace": work_request.workspace,
            "user": work_request.created_by or AnonymousUser(),
            "workflow_root": work_request.get_workflow_root(),
        }

        if isinstance(lookup, LookupMultiple):
            results = lookup_multiple(lookup, **lookup_kwargs)
        else:
            results = [lookup_single(lookup, **lookup_kwargs)]

        for result in results:
            collection_item = result.collection_item

            if (
                collection_item is not None
                and collection_item.child_type == CollectionItem.Types.BARE
                and collection_item.category == BareDataCategory.PROMISE
            ):
                work_request_dependency = WorkRequest.objects.get(
                    id=collection_item.data["promise_work_request_id"]
                )
                work_request.add_dependency(work_request_dependency)

    def work_request_ensure_child(
        self,
        task_name: str,
        task_data: dict[str, Any],
        workflow_data: WorkRequestWorkflowData,
        task_type: TaskTypes = TaskTypes.WORKER,
    ) -> WorkRequest:
        """
        Create the child if one does not already exist.

        :return: new WorkRequest or existing work request.
        """
        try:
            return self.work_request.children.get(
                task_name=task_name,
                task_data=task_data,
                workflow_data_json=(
                    workflow_data.dict(exclude_unset=True)
                    if workflow_data
                    else {}
                ),
                task_type=task_type,
            )
        except WorkRequest.DoesNotExist:
            return self.work_request.create_child(
                task_name=task_name,
                task_data=task_data,
                workflow_data=workflow_data,
                task_type=task_type,
            )

    @staticmethod
    def lookup_result_architecture(result: LookupResult) -> str:
        """Get architecture from result of looking up an artifact."""
        if result.collection_item is not None:
            architecture = result.collection_item.data.get("architecture")
            if isinstance(architecture, str):
                return architecture
        elif result.artifact is not None:
            artifact_data = result.artifact.create_data()
            match artifact_data:
                case DebianBinaryPackages():
                    architecture = artifact_data.architecture
                case DebianBinaryPackage():
                    architecture = artifact_data.deb_fields.get("Architecture")
                case DebianUpload():
                    architecture = artifact_data.changes_fields.get(
                        "Architecture"
                    )
                case _:
                    raise ArtifactHasNoArchitecture(f"{type(artifact_data)}")

            assert isinstance(architecture, str)
            return architecture

        else:
            raise ValueError(
                "Unexpected result: must have collection_item or artifact"
            )

        raise ValueError(
            f"Cannot determine architecture for lookup result: {result}"
        )

    def filter_artifact_lookup(
        self, lookup: LookupMultiple, architectures: Iterable[str]
    ) -> LookupMultiple:
        """Filter an artifact lookup by architecture."""
        results = lookup_multiple(
            lookup,
            self.workspace,
            user=self.work_request.created_by,
            workflow_root=self.work_request.get_workflow_root(),
            expect_type=LookupChildType.ARTIFACT_OR_PROMISE,
        )
        relevant: list[LookupSingle] = []
        for result in results:
            arch_in_lookup = self.lookup_result_architecture(result)
            if arch_in_lookup in architectures:
                relevant.append(reconstruct_lookup(result))

        return LookupMultiple.parse_obj(sorted(relevant))

    def get_architectures(self, lookup: LookupMultiple) -> set[str]:
        """
        Return set with all the architectures in the artifacts from the lookup.

        The architectures are extracted from each lookup result using
        :py:meth:`lookup_result_architecture`.
        """
        results = lookup_multiple(
            lookup,
            self.workspace,
            user=self.work_request.created_by,
            workflow_root=self.work_request.get_workflow_root(),
            expect_type=LookupChildType.ARTIFACT_OR_PROMISE,
        )

        return {self.lookup_result_architecture(result) for result in results}

    def _execute(self) -> bool:
        """Unused abstract method from BaseTask: populate() is used instead."""
        raise NotImplementedError()


def orchestrate_workflow(work_request: WorkRequest) -> None:
    """
    Orchestrate a workflow in whatever way is appropriate.

    For a workflow callback, run `callback` and mark the work request as
    completed.  For a workflow, run `populate` and unblock workflow
    children, but leave the workflow running until all its children have
    finished.  For any other work request, raise an error.
    """
    try:
        match (work_request.task_type, work_request.task_name):
            case (TaskTypes.INTERNAL, "workflow") | (TaskTypes.WORKFLOW, _):
                try:
                    orchestrator = work_request.get_task()
                except InternalTaskError as exc:
                    raise WorkflowRunError(work_request, str(exc)) from exc
                except TaskConfigError as exc:
                    raise WorkflowRunError(
                        work_request, f"failed to configure: {exc}"
                    ) from exc
                else:
                    assert isinstance(orchestrator, Workflow)

            case _:
                raise WorkflowRunError(
                    work_request, "does not have a workflow orchestrator"
                )

        if not orchestrator.ensure_dynamic_data(
            TaskDatabase(orchestrator.work_request)
        ):
            raise WorkflowRunError(
                work_request, "failed to compute dynamic data"
            )

        match (work_request.task_type, work_request.task_name):
            case (TaskTypes.INTERNAL, "workflow"):
                try:
                    orchestrator.callback(work_request)
                except Exception as exc:
                    raise WorkflowRunError(
                        work_request, f"orchestrator failed: {exc}"
                    ) from exc

                work_request.mark_completed(WorkRequest.Results.SUCCESS)

            case (TaskTypes.WORKFLOW, _):
                try:
                    orchestrator.populate()
                except Exception as exc:
                    raise WorkflowRunError(
                        work_request, f"orchestrator failed: {exc}"
                    ) from exc

                orchestrator.work_request.unblock_workflow_children()

                # The workflow is left running until all its children have
                # finished.

            case _:  # pragma: no cover
                # Already checked above.
                raise AssertionError(
                    f"Unexpected work request: "
                    f"{work_request.task_type}/{work_request.task_name}"
                )

    except WorkflowRunError as exc:
        logger.warning(
            "Error running work request %s/%s (%s): %s",
            exc.work_request.task_type,
            exc.work_request.task_name,
            exc.work_request.id,
            exc.message,
        )
        work_request.mark_completed(WorkRequest.Results.ERROR)
        raise
