# Copyright © 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.
r"""
Collection of tasks.

The debusine.tasks module hierarchy hosts a collection of
:py:class:`BaseExternalTask` that are used by workers to fulfill
:py:class:`debusine.db.models.WorkRequest`\ s sent by the debusine
scheduler.

Creating a new task requires adding a new file containing a class inheriting
from the :py:class:`BaseExternalTask` or :py:class:`RunCommandTask` base class.
The name of the class must be unique among all child classes.

A child class must, at the very least, override the
:py:meth:`BaseExternalTask.execute` method.
"""

import datetime as dt
import inspect
import logging
import os
import shlex
import shutil
import signal
import subprocess
import tempfile
import traceback
from abc import ABCMeta, abstractmethod
from collections import defaultdict
from collections.abc import Collection, Generator
from contextlib import contextmanager
from pathlib import Path
from typing import (
    Any,
    AnyStr,
    BinaryIO,
    ClassVar,
    IO,
    TYPE_CHECKING,
    TextIO,
    Union,
    overload,
    override,
)

from debian.debian_support import Version

from debusine.artifacts import WorkRequestDebugLogs
from debusine.artifacts.models import (
    ArtifactCategory,
    CollectionCategory,
    TaskTypes,
    WorkRequestResults,
)
from debusine.client.debusine import Debusine
from debusine.client.models import ArtifactResponse, RelationType
from debusine.tasks.executors import (
    ExecutorInterface,
    InstanceInterface,
    executor_class,
)
from debusine.tasks.inputs import (
    BuildArchitectureInput,
    ExtraRepositoriesInput,
    Stage,
    TaskFieldContainer,
)
from debusine.tasks.models import (
    BackendType,
    BaseDynamicTaskData,
    BaseDynamicTaskDataWithExecutor,
    BaseDynamicTaskDataWithExtraRepositories,
    BaseTaskData,
    BaseTaskDataWithExecutor,
    BaseTaskDataWithExtraRepositories,
    WorkerType,
)
from debusine.utils import extract_generic_type_arguments

if TYPE_CHECKING:
    from _typeshed import OpenBinaryModeWriting, OpenTextModeWriting

    import debusine.tasks.tags as ttags
    import debusine.worker.tags as wtags


#: Fallback arhitecturer to use when a task doesn't explicitly specify one
FALLBACK_ARCHITECTURE = "amd64"


class TaskConfigError(Exception):
    """Exception raised when there is an issue with a task configuration."""

    @override
    def __init__(
        self, message: str | None, original_exception: Exception | None = None
    ) -> None:
        """
        Initialize the TaskConfigError.

        :param message: human-readable message describing the error.
        :param original_exception: the exception that triggered this error,
          if applicable. This is used to provide additional information.
        """
        super().__init__(message)
        self.original_exception = original_exception

    def add_parent_message(self, msg: str) -> None:
        """Prepend the error message with one from the containing scope."""
        if self.args[0]:
            self.args = (f"{msg}: {self.args[0]}",)
        else:
            self.args = (msg,)

    @override
    def __str__(self) -> str:
        """
        Return a string representation.

        If an original exception is present, its representation is appended
        to the message for additional context.
        """
        if self.original_exception:
            if self.args[0] is not None:
                return (
                    f"{self.args[0]} (Original exception: "
                    f"{self.original_exception})"
                )
            else:
                return str(self.original_exception)
        else:
            return str(self.args[0])


def ensure_artifact_categories(
    *,
    configuration_key: str,
    category: str,
    expected: Collection[ArtifactCategory],
) -> None:
    """
    Validate that the artifact's category is one of the expected categories.

    :param configuration_key: Optional key to identify the source of the
      artifact. Provides additional context in error messages.
    :param category: The category to validate.
    :param expected: A collection of valid artifact categories.
    """
    valid_categories = [cat.value for cat in sorted(expected)]
    if category not in valid_categories:
        raise TaskConfigError(
            f"{configuration_key}: unexpected artifact "
            f"category: '{category}'. "
            f"Valid categories: {valid_categories}"
        )


def ensure_collection_category(
    *, configuration_key: str, category: str, expected: CollectionCategory
) -> None:
    """
    Validate that the collection's category is as expected.

    :param configuration_key: Optional key to identify the source of the
      artifact. Provides additional context in error messages.
    :param category: The category to validate.
    :param expected: The expected collection category.
    """
    if category != expected:
        raise TaskConfigError(
            f"{configuration_key}: unexpected collection category: "
            f"'{category}'. Expected: '{expected}'"
        )


def analyze_external_worker_tasks() -> dict[str, Any]:
    """
    Return dictionary with metadata for each BaseExternalTask.

    Subclasses of BaseExternalTask get registered in
    BaseExternalTask._sub_tasks. Return a dictionary with the metadata of each
    of the subtasks.

    This method is executed in the worker when submitting the dynamic
    metadata.
    """
    metadata = {}

    for registry in BaseExternalTask._sub_tasks.values():
        for task_class in registry.values():
            metadata.update(task_class.analyze_worker())

    return metadata


def get_provided_external_worker_tags() -> set[str]:
    """
    Return the set of worker tags provided by all available tasks.

    This method is called on the worker to collect information about the
    worker.
    """
    tags: set[str] = set()

    for registry in BaseExternalTask._sub_tasks.values():
        for task_class in registry.values():
            tags |= task_class.get_provided_worker_tags()

    return tags


class BaseExternalTask[
    TD: BaseTaskData,
    DTD: BaseDynamicTaskData = BaseDynamicTaskData,
](TaskFieldContainer, metaclass=ABCMeta):
    r"""
    Base class for external tasks.

    A BaseExternalTask object serves two purposes: encapsulating the logic of
    what needs to be done to execute the task on an external worker (cf
    :py:meth:`configure` and :py:meth:`execute`, that are run on the worker),
    and supporting the server's task-specific logic, such as scheduling and UI
    presentation.

    Scheduling support is done in a two-step process: collating metadata from
    each worker (with the :py:meth:`analyze_worker` method that is run on a
    worker) and then, based on this metadata, see if a task is suitable (with
    :py:meth:`can_run_on` that is executed on the scheduler).

    If tag-based scheduling is in use, that is replaced by
    :py:meth`get_provided_worker_tags`, which runs on the worker and reports
    provided worker tags that need to match what tasks require.

    Concrete subclasses must implement:

    * ``run(execute_directory: Path) -> bool``: Do the main work of the
      task.

    Most concrete subclasses should also implement:

    * ``fetch_input(self, destination) -> bool``. Download the needed
      artifacts into destination. Suggestion: can use
      ``fetch_artifact(artifact, dir)`` to download them.
      (default: return True)
    * ``configure_for_execution(self, download_directory: Path) -> bool``
      (default: return True)
    * ``check_directory_for_consistency_errors(self, build_directory: Path)
      -> list[str]``
      (default: return an empty list, indicating no errors)
    * ``upload_artifacts(self, directory: Path, \*, execution_success: bool)``.
      The member variable self._source_artifacts_ids is set by
      ``fetch_input()`` and can be used to create the relations between
      uploaded artifacts and downloaded artifacts.
      (default: return True)

    Most concrete task implementations should inherit from
    :py:class:`RunCommandTask` instead.
    """

    #: Class used as the in-memory representation of task data.
    task_data_type: type[TD]
    data: TD

    #: Class used as the in-memory representation of dynamic task data.
    dynamic_task_data_type: type[DTD]
    dynamic_data: DTD | None

    #: Must be overridden by child classes to document the current version of
    #: the task's code. A task will only be scheduled on a worker if its task
    #: version is the same as the one running on the scheduler.
    TASK_VERSION: int | None = None

    #: The worker type must be suitable for the task type.  TaskTypes.WORKER
    #: requires an external worker; TaskTypes.SERVER requires a Celery
    #: worker; TaskTypes.SIGNING requires a signing worker.
    TASK_TYPE: TaskTypes = TaskTypes.WORKER

    name: ClassVar[str]
    _sub_tasks: dict[
        TaskTypes, dict[str, type["BaseExternalTask['Any', 'Any']"]]
    ] = defaultdict(dict)

    @override
    def __init_subclass__(cls, **kwargs: Any) -> None:
        """
        Register the subclass into BaseExternalTask._sub_tasks.

        Used by BaseExternalTask.class_from_name() to return the class given
        the name.
        """
        super().__init_subclass__(**kwargs)

        # The name of the task. It is computed by converting the class name
        # to lowercase.
        cls.name = getattr(cls, "TASK_NAME", cls.__name__.lower())

        # The task data types, computed by introspecting the type arguments
        # used to specialize this generic class.
        [
            cls.task_data_type,
            cls.dynamic_task_data_type,
        ] = extract_generic_type_arguments(cls, BaseExternalTask)

        if inspect.isabstract(cls):
            # Don't list abstract base classes as tasks.
            return

        registry = cls._sub_tasks[cls.TASK_TYPE]

        # The same sub-task could register twice
        # (but assert that is the *same* class, not a different
        # subtask with a name with a different capitalisation)
        if cls.name in registry and registry[cls.name] != cls:
            raise AssertionError(f'Two Tasks with the same name: {cls.name!r}')

        # Make sure SERVER and WORKER do not have conflicting task names
        match cls.TASK_TYPE:
            case TaskTypes.SERVER:
                if cls.name in cls._sub_tasks[TaskTypes.WORKER]:
                    raise AssertionError(
                        f'{cls.name!r} already registered as a Worker task'
                    )
            case TaskTypes.WORKER:
                if cls.name in cls._sub_tasks[TaskTypes.SERVER]:
                    raise AssertionError(
                        f'{cls.name!r} already registered as a Server task'
                    )

        registry[cls.name] = cls

    def __init__(
        self,
        task_data: dict[str, Any],
        dynamic_task_data: dict[str, Any] | None = None,
        *,
        wrapped: bool = False,
    ) -> None:
        """
        Initialize the task.

        :param wrapped: set to True when this task is managed by a wrapping task
        """
        super().__init__()

        #: A :py:class:`logging.Logger` instance that can be used in child
        #: classes when you override methods to implement the task.
        self.logger = logging.getLogger("debusine.tasks")

        #: Validated task data submitted through :py:meth:`configure` without
        # BaseExternalTask generic data
        self._configure(task_data, dynamic_task_data)

        self.working_directory: Path | None = None
        self.worker_environment_variables: dict[str, str] = {}

        # Task is aborted: the task does not need to be executed, and can be
        # stopped if it is already running
        self._aborted = False

        self.work_request_id: int | None = None

        # Injectable HTTP timeout from the worker configuration
        self.http_timeout: int | tuple[int, int] | None = None

        # Workspace is used when downloading and uploading artifacts.
        # When the worker instantiates the task it should set
        # self.workspace_name (see issue #186).
        self.workspace_name: str | None = None

        # fetch_input() add the downloaded artifacts. Used by
        # `BaseExternalTask._upload_work_request_debug_logs()` and maybe by
        # required method `upload_artifacts()`.
        #
        # This is distinct from get_input_artifacts_ids, which is used to
        # extract IDs from dynamic_data for use by UI views
        self._source_artifacts_ids: list[int] = []

        self.debusine: Debusine | None = None

        self.executor: ExecutorInterface | None = None
        self.executor_instance: InstanceInterface | None = None

        # Restore values of task input fields stored in dynamic data
        if not wrapped and self.dynamic_data is not None:
            self.restore_field_values(self.dynamic_data)

        self.post_init()

    def post_init(self) -> None:
        """Specific post-init code."""

    def set_working_directory(self, working_directory: Path) -> None:
        working_directory.mkdir(parents=True, exist_ok=True)
        self.working_directory = working_directory

    def set_worker_environment_variables(self, env: dict[str, str]) -> None:
        self.worker_environment_variables = env

    @property
    def log_directory(self) -> Path:
        """Return the directory storing logs of the task's execution."""
        assert self.working_directory
        return self.working_directory / "logs"

    def append_to_log_file(self, filename: str, lines: list[str]) -> None:
        """
        Open log file and write contents into it.

        :param filename: use self.open_debug_log_file(filename)
        :param lines: write contents to the logfile
        """
        with self.open_debug_log_file(filename) as file:
            file.writelines([line + "\n" for line in lines])

    @overload
    def open_debug_log_file(
        self, filename: str, *, mode: "OpenTextModeWriting" = "a"
    ) -> TextIO: ...

    @overload
    def open_debug_log_file(
        self, filename: str, *, mode: "OpenBinaryModeWriting"
    ) -> BinaryIO: ...

    def open_debug_log_file(
        self,
        filename: str,
        *,
        mode: Union["OpenTextModeWriting", "OpenBinaryModeWriting"] = "a",
    ) -> IO[Any]:
        """
        Open a temporary file and return it.

        The files are always for the same temporary directory, calling it twice
        with the same file name will open the same file.

        The caller must call .close() when finished writing.
        """
        debug_file = self.log_directory / filename
        debug_file.parent.mkdir(exist_ok=True)
        return debug_file.open(mode)

    @classmethod
    def prefix_with_task_name(cls, text: str) -> str:
        """:return: the ``text`` prefixed with the task name and a colon."""
        if cls.TASK_TYPE is TaskTypes.WORKER:
            # Worker tasks are left unprefixed for compatibility
            return f"{cls.name}:{text}"
        else:
            return f"{cls.TASK_TYPE.lower()}:{cls.name}:{text}"

    @classmethod
    def analyze_worker(cls) -> dict[str, Any]:
        """
        Return dynamic metadata about the current worker.

        This method is called on the worker to collect information about the
        worker. The information is stored as a set of key-value pairs in a
        dictionary.

        That information is then reused on the scheduler to be fed to
        :py:meth:`can_run_on` and determine if a task is suitable to be
        executed on the worker.

        Derived objects can extend the behaviour by overriding
        the method, calling ``metadata = super().analyze_worker()``,
        and then adding supplementary data in the dictionary.

        To avoid conflicts on the names of the keys used by different tasks
        you should use key names obtained with
        ``self.prefix_with_task_name(...)``.

        :return: a dictionary describing the worker.
        :rtype: dict.
        """
        version_key_name = cls.prefix_with_task_name("version")
        return {
            version_key_name: cls.TASK_VERSION,
        }

    @classmethod
    def get_provided_worker_tags(cls) -> set[str]:
        """
        Return the set of worker tags provided by this task class.

        This method is called on the worker to collect information about the
        worker.
        """
        # Prevent circular import
        import debusine.worker.tags as wtags

        return {
            wtags.TASK_PREFIX
            + f"{cls.TASK_TYPE.lower()}:{cls.name}:version:{cls.TASK_VERSION}"
        }

    def can_run_on(self, worker_metadata: dict[str, Any]) -> bool:
        """
        Check if the specified worker can run the task.

        This method shall take its decision solely based on the supplied
        ``worker_metadata`` and on the configured task data (``self.data``).

        The default implementation always returns True unless
        :py:attr:`TASK_TYPE` doesn't match the worker type or there's a
        mismatch between the :py:attr:`TASK_VERSION` on the scheduler side
        and on the worker side.

        Derived objects can implement further checks by overriding the method
        in the following way::

            if not super().can_run_on(worker_metadata):
                return False

            if ...:
                return False

            return True

        :param dict worker_metadata: The metadata collected from the worker by
            running :py:meth:`analyze_worker` on all the tasks on the worker
            under consideration.
        :return: the boolean result of the check.
        :rtype: bool.
        """
        worker_type = worker_metadata.get("system:worker_type")
        if (self.TASK_TYPE, worker_type) not in {
            (TaskTypes.WORKER, WorkerType.EXTERNAL),
            (TaskTypes.SERVER, WorkerType.CELERY),
            (TaskTypes.SIGNING, WorkerType.SIGNING),
        }:
            return False

        version_key_name = self.prefix_with_task_name("version")

        if worker_metadata.get(version_key_name) != self.TASK_VERSION:
            return False

        return True

    def compute_scheduler_tags(
        self,
    ) -> tuple["ttags.TaskProvidedTags", "wtags.TaskRequiredTags"]:
        """
        Compute task-provided and task-required tagsets.

        This is called twice, and both times with only the ``stage=PENDING``
        fields resolved:

        * Once before task configuration, to compute the early tag set used for
          the configuration lookup
        * Once after task configuration, to the final tag set used for
          scheduling
        """
        import debusine.tasks.tags as ttags
        import debusine.worker.tags as wtags

        provided = ttags.TaskProvidedTags()
        required = wtags.TaskRequiredTags()

        return provided, required

    @abstractmethod
    def get_subject(self) -> str | None:
        """
        Return the task subject used to look up task configuration.

        This method can only be called after all the task fields have been
        resolved.
        """

    def get_configuration_context(self) -> str | None:
        """
        Return the configuration context used to look up task configuration.

        This method can only be called after all the task fields have been
        resolved.
        """
        return None

    def compute_dynamic_data(self) -> DTD:
        """
        Build a dynamic task data structure for this task.

        This method can only be called after all the task fields have been
        resolved.

        :returns: the newly created dynamic task data
        """
        return self.dynamic_task_data_type()

    def _configure(
        self,
        task_data: dict[str, Any],
        dynamic_task_data: dict[str, Any] | None = None,
    ) -> None:
        """
        Configure the task with the supplied ``task_data``.

        The supplied data is validated against the pydantic data model for the
        type of task being configured. If validation fails, a TaskConfigError
        is raised. Otherwise, the `data` attribute is filled with the supplied
        `task_data`.

        :param dict task_data: The supplied data describing the task.
        :raises TaskConfigError: if the dict does not validate.
        """
        try:
            self.data = self.task_data_type(**task_data)
            self.dynamic_data = (
                None
                if dynamic_task_data is None
                else self.dynamic_task_data_type(**dynamic_task_data)
            )
        except ValueError as exc:
            raise TaskConfigError(None, original_exception=exc)

    def execute_logging_exceptions(self) -> WorkRequestResults:
        """Execute self.execute() logging any raised exceptions."""
        try:
            return self.execute()
        except Exception as exc:
            self.logger.info("Exception in Task %s", self.name, exc_info=True)
            raise exc

    def execute(self) -> WorkRequestResults:
        """
        Call the _execute() method, upload debug artifacts.

        See _execute() for more information.

        :return: result of the _execute() method.
        """  # noqa: D402
        result = self._execute()

        self._upload_work_request_debug_logs()

        return result

    def abort(self) -> None:
        """Task does not need to be executed. Once aborted cannot be changed."""
        self._aborted = True

    @property
    def aborted(self) -> bool:
        """
        Return if the task is aborted.

        Tasks cannot transition from aborted -> not-aborted.
        """
        return self._aborted

    @staticmethod
    def class_from_name(
        task_type: TaskTypes, task_name: str
    ) -> type["BaseExternalTask['Any', 'Any']"]:
        """
        Return class for :param task_name (case-insensitive).

        :param task_type: type of task to look up

        __init_subclass__() registers BaseExternalTask subclasses' into
        BaseExternalTask._sub_tasks.
        """
        if (registry := BaseExternalTask._sub_tasks.get(task_type)) is None:
            raise ValueError(f"{task_type!r} is not a registered task type")

        task_name_lowercase = task_name.lower()
        if (cls := registry.get(task_name_lowercase)) is None:
            raise ValueError(
                f"{task_name_lowercase!r} is not a registered"
                f" {task_type} task_name"
            )

        return cls

    def get_input_artifacts_ids(self) -> list[int]:
        """
        Return the list of input artifact IDs used by this task.

        This refers to the artifacts actually used by the task. If
        dynamic_data is empty, this returns the empty list.

        This is used by views to show what artifacts were used by a task.
        `_source_artifacts_ids` cannot be used for this purpose because it is
        only set during task execution.
        """
        artifact_ids: list[int] = []

        if self.dynamic_data is not None:
            artifact_ids.extend(self.dynamic_data.get_input_artifacts_ids())

        for name, field in self.inputs.items():
            if name not in self.__dict__:
                continue
            artifact_ids.extend(field.get_artifact_ids(self.__dict__[name]))

        return artifact_ids

    def configure_server_access(self, debusine: Debusine) -> None:
        """Set the object to access the server."""
        self.debusine = debusine

    @contextmanager
    def _temporary_directory(self, prefix: str) -> Generator[Path]:
        assert self.working_directory
        with tempfile.TemporaryDirectory(
            prefix=f"{prefix}-",
            dir=self.working_directory,
        ) as directory:
            yield Path(directory)

    def _log_exception(self, exc: Exception, stage: str) -> None:
        """Log an unexpected exception into the task log for stage."""
        exc_type = type(exc).__name__
        exc_traceback = traceback.format_exc()

        log_message = [
            f"Exception type: {exc_type}",
            f"Message: {exc}",
            "",
            exc_traceback,
        ]

        self.append_to_log_file(f"{stage}.log", log_message)

    def fetch_artifact(
        self, artifact_id: int, destination: Path
    ) -> ArtifactResponse:
        """
        Download artifact_id to destination.

        Add artifact_id to self._source_artifacts_ids.
        """
        if not self.debusine:
            raise AssertionError("self.debusine not set")

        if not self.workspace_name:  # Required for cross-scope artifacts
            raise AssertionError("self.workspace_name not set")

        artifact_response = self.debusine.download_artifact(
            artifact_id,
            destination,
            tarball=False,
            workspace=self.workspace_name,
        )
        self._source_artifacts_ids.append(artifact_id)
        return artifact_response

    def fetch_input(self, destination: Path) -> bool:  # noqa: ARG002, U100
        """
        Download artifacts needed by the task, update self.source_artifacts_ids.

        Task might use self.data.input to download the relevant artifacts.

        The method self.fetch_artifact(artifact, destination) might be used
        to download the relevant artifacts and update
        self.source_artifacts_ids.
        """
        return True

    def configure_for_execution(
        self,
        download_directory: Path,  # noqa: ARG002, U100
    ) -> bool:
        """
        Configure task: set variables needed for the self._cmdline().

        Called after the files are downloaded via `fetch_input()`.
        """
        return True

    def prepare_to_run(
        self,
        download_directory: Path,  # noqa: ARG002, U100
        execute_directory: Path,  # noqa: ARG002, U100
    ) -> None:
        """
        Prepare the execution environment to do the main work of the task.

        :param download_directory: Directory where required artifacts were
            previously downloaded by self.fetch_input(download_directory).
        :param execute_directory: Directory where the task is expected to
            do its job.
        """
        # Do nothing by default.

    @abstractmethod
    def run(
        self,
        execute_directory: Path,  # noqa: ARG002, U100
    ) -> WorkRequestResults:
        """
        Do the main work of the task.

        :param execute_directory: Temporary directory prepared by the worker
            and by former steps of the task (prepare_to_run notably) where the
            task is expected to do its job.
        """

    def check_directory_for_consistency_errors(
        self,
        build_directory: Path,  # noqa: ARG002, U100
    ) -> list[str]:
        """Return list of errors after doing the main work of the task."""
        return []

    def upload_artifacts(
        self,
        execute_directory: Path,
        *,
        execution_result: WorkRequestResults,  # noqa: U100
    ) -> None:
        """Upload the artifacts for the task."""
        # Do nothing by default.

    def cleanup(self) -> None:
        """Clean up after running the task."""
        # Do nothing by default.

    def _execute(self) -> WorkRequestResults:  # noqa: C901
        """
        Execute the requested task.

        The task must first have been configured. It is allowed to take
        as much time as required. This method will only be run on a worker. It
        is thus allowed to access resources local to the worker.

        It is recommended to fail early by raising a :py:exc:`TaskConfigError`
        if the parameters of the task let you anticipate that it has no chance
        of completing successfully.

        Fetch the required input, execute the command and upload artifacts.

        Flow is:

        - Call ``self.fetch_input():`` download the artifacts
        - Call ``self.configure_for_execution()``: to set any member variables
          that might be used in ``self._cmdline()``
        - Call ``self.prepare_to_run()`` to prepare the execution
          environment to do the main work of the task
        - Call ``self.run()`` to do the main work of the task
        - Call ``self.check_directory_for_consistency_errors()``. If any
          errors are returned save them into consistency.log.
        - Call ``self.upload_artifacts(exec_dir, execution_success=succeeded)``.
        - Return ``execution_result``  (the Worker will set the result)

        :return: SUCCESS to indicate success, FAILURE for a failure, ERROR for
            an internal error, SKIPPED if the task turned out to be a noop.
        :raises TaskConfigError: if the parameters of the work request are
            incompatible with the worker.
        """
        with (
            self._temporary_directory("exec") as execute_directory,
            self._temporary_directory("fetch") as download_directory,
            self.open_debug_log_file("stages.log") as stages_log,
        ):

            def log_stage(message: str) -> None:
                """Log a stage, both to the worker log and to a WR debug log."""
                self.logger.info(
                    "Work request %s: %s", self.work_request_id, message
                )
                print(
                    f"{dt.datetime.now(dt.UTC).isoformat()} {message}",
                    file=stages_log,
                )

            try:
                try:
                    log_stage("Fetching input")
                    if not self.fetch_input(download_directory):
                        return WorkRequestResults.FAILURE
                except Exception as exc:
                    self._log_exception(exc, stage="fetch_input")
                    return WorkRequestResults.FAILURE

                try:
                    log_stage("Configuring for execution")
                    if not self.configure_for_execution(download_directory):
                        return WorkRequestResults.FAILURE
                except Exception as exc:
                    self._log_exception(exc, stage="configure_for_execution")
                    return WorkRequestResults.FAILURE

                try:
                    log_stage("Preparing to run")
                    self.prepare_to_run(download_directory, execute_directory)

                    log_stage("Running")

                    execution_result = self.run(execute_directory)
                except Exception as exc:
                    self._log_exception(exc, stage="execution")
                    return WorkRequestResults.FAILURE

                try:
                    log_stage("Checking output")
                    if errors := self.check_directory_for_consistency_errors(
                        execute_directory
                    ):
                        self.append_to_log_file(
                            "consistency.log", sorted(errors)
                        )
                        return WorkRequestResults.FAILURE

                    log_stage("Uploading artifacts")
                    self.upload_artifacts(
                        execute_directory, execution_result=execution_result
                    )
                except Exception as exc:
                    self._log_exception(exc, stage="post_execution")
                    return WorkRequestResults.FAILURE
            finally:
                try:
                    log_stage("Cleaning up")
                    self.cleanup()
                except Exception as exc:
                    self._log_exception(exc, stage="post_execution")
                    execution_result = WorkRequestResults.FAILURE

        if execution_result == WorkRequestResults.SUCCESS:
            # We mainly care about stages.log for failures.  It's rather
            # boring for successes.
            (self.log_directory / "stages.log").unlink()

        return execution_result

    def _upload_work_request_debug_logs(self) -> None:
        """
        Create a WorkRequestDebugLogs artifact and upload the logs.

        The logs might exist in self.log_directory and were
        added via self.open_debug_log_file() or self.append_debug_log_file().

        For each self._source_artifacts_ids: create a relation from
        WorkRequestDebugLogs to source_artifact_id.
        """
        if not self.log_directory.exists():
            return

        work_request_debug_logs_artifact = WorkRequestDebugLogs.create(
            files=self.log_directory.glob("*")
        )

        assert self.debusine
        remote_artifact = self.debusine.upload_artifact(
            work_request_debug_logs_artifact,
            workspace=self.workspace_name,
            work_request=self.work_request_id,
        )

        for source_artifact_id in self._source_artifacts_ids:
            self.debusine.relation_create(
                remote_artifact.id,
                source_artifact_id,
                RelationType.RELATES_TO,
            )

        shutil.rmtree(self.log_directory)


class BaseTaskWithExecutor[
    TDE: BaseTaskDataWithExecutor,
    DTDE: BaseDynamicTaskDataWithExecutor,
](BaseExternalTask[TDE, DTDE], metaclass=ABCMeta):
    r"""
    Base for tasks with executor capabilities.

    Concrete subclasses must implement ``fetch_input()``,
    ``configure_for_execution()``, ``run()``,
    ``check_directory_for_consistency_errors()``, and
    ``upload_artifacts()``, as documented by :py:class:`BaseExternalTask`.
    """

    DEFAULT_BACKEND = BackendType.UNSHARE

    #: Task build architecture
    build_architecture = BuildArchitectureInput(
        stage=Stage.PENDING,
        verbose_name="Build architecture",
        help_text="System architecture for the build environment",
    )

    @override
    def can_run_on(self, worker_metadata: dict[str, Any]) -> bool:
        # Also require a matching build architecture
        if not super().can_run_on(worker_metadata):
            return False

        if self.build_architecture not in worker_metadata.get(
            "system:architectures", []
        ):
            return False

        return True

    def _prepare_executor(self) -> None:
        """
        Prepare the executor.

        * Set self.executor to the new executor with self.backend and the
          appropriate environment ID (which must have been looked up by the
          task's compute_dynamic_data method)
        * Download the image
        """
        assert self.debusine
        assert self.dynamic_data
        assert self.dynamic_data.environment_id
        assert self.workspace_name
        assert self.working_directory
        self.executor = executor_class(self.backend)(
            self.debusine,
            self.dynamic_data.environment_id,
            workspace=self.workspace_name,
            working_directory=self.working_directory,
        )
        self.executor.download_image()

    def _prepare_executor_instance(self) -> None:
        """
        Create and start an executor instance.

        If self.executor is None: call self._prepare_executor()

        Set self.executor_instance to the new executor instance, starts
        the instance.
        """
        if self.executor is None:
            self._prepare_executor()

        assert self.executor
        self.executor_instance = self.executor.create()
        self.executor_instance.start()

    @property
    def backend(self) -> BackendType:
        """Return the backend name to use."""
        backend = self.data.backend
        if backend == BackendType.AUTO:
            backend = self.DEFAULT_BACKEND
        return backend

    @override
    def compute_scheduler_tags(
        self,
    ) -> tuple["ttags.TaskProvidedTags", "wtags.TaskRequiredTags"]:
        # Prevent circular import
        import debusine.worker.tags as wtags

        provided, required = super().compute_scheduler_tags()
        required.add(
            wtags.ProvenanceRequired.SYSTEM,
            {
                wtags.EXECUTOR_PREFIX + self.backend,
                wtags.BUILD_ARCH_PREFIX + self.build_architecture,
            },
        )
        return provided, required

    def get_package_version(self, package_name: str) -> Version:
        """
        Get the installed version of a package in the executor.

        :raises subprocess.CalledProcessError: if the package is not
          installed (also logged to ``install.log``).
        """
        assert self.executor_instance is not None
        cmd = ["dpkg-query", "--show", "--showformat=${Version}", package_name]
        try:
            return Version(
                self.executor_instance.run(
                    cmd, check=True, run_as_root=True, text=True
                ).stdout
            )
        except subprocess.CalledProcessError as e:
            with self.open_debug_log_file("install.log", mode="a") as log:
                log.write(
                    f"Failed to determine available {package_name} version.\n"
                    f"command: {shlex.join(cmd)}\n"
                    f"exitcode: {e.returncode}\n"
                    f"stdout: {e.stdout}\n"
                    f"stderr: {e.stderr}\n"
                )
                raise

    @override
    def prepare_to_run(
        self, download_directory: Path, execute_directory: Path
    ) -> None:
        """Copy the download and execution directories into the executor."""
        if self.executor_instance:
            self.executor_instance.create_user()
            non_root_user = self.executor_instance.non_root_user
            self.executor_instance.mkdir(
                download_directory.parent,
                user=non_root_user,
                group=non_root_user,
                parents=True,
            )
            self.executor_instance.directory_push(
                download_directory,
                download_directory.parent,
                user=non_root_user,
                group=non_root_user,
            )
            self.executor_instance.mkdir(
                execute_directory,
                user=non_root_user,
                group=non_root_user,
                parents=True,
            )

    def run_executor_command(
        self,
        cmd: list[str],
        log_filename: str,
        run_as_root: bool = False,
        check: bool = True,
    ) -> None:
        """Run cmd within the executor, logging the output to log_name."""
        assert self.executor_instance
        with self.open_debug_log_file(log_filename, mode="ab") as cmd_log:
            msg = f"Executing: {shlex.join(cmd)}\n"
            cmd_log.write(msg.encode())
            cmd_log.flush()
            p = self.executor_instance.run(
                cmd,
                run_as_root=run_as_root,
                stdout=cmd_log,
                stderr=cmd_log,
                check=check,
            )
            msg = f"Execution completed (exit code {p.returncode})\n"
            cmd_log.write(msg.encode())

    @override
    def cleanup(self) -> None:
        """
        Clean up after running the task.

        Some tasks use the executor in upload_artifacts, so we clean up the
        executor here rather than in run().
        """
        if self.executor_instance and self.executor_instance.is_started():
            self.executor_instance.stop()


class RunCommandTask[
    TD: BaseTaskData,
    DTD: BaseDynamicTaskData = BaseDynamicTaskData,
](BaseExternalTask[TD, DTD], metaclass=ABCMeta):
    r"""
    A :py:class:`BaseExternalTask` that executes commands and uploads artifacts.

    Concrete subclasses must implement:

    * ``_cmdline(self) -> list[str]``
    * ``task_result(self, returncode: Optional[int], execute_directory: Path)
      -> WorkRequestResults`` (defaults to SUCCESS)

    They must also implement ``configure_for_execution()``,
    ``fetch_input()``, ``check_directory_for_consistency_errors()``, and
    ``upload_artifacts()``, as documented by :py:class:`BaseExternalTask`.
    (They do not need to implement ``run()``, but may do so if they need to
    run multiple commands rather than just one.)

    Use ``self.append_to_log_file()`` / ``self.open_debug_log_file()`` to
    provide information for the user (it will be available to the user as an
    artifact).

    Command execution uses process groups to make sure that the command and
    possible spawned commands are finished, and cancels the execution of the
    command if ``BaseExternalTask.aborted()`` is True. If the task also
    inherits from BaseTaskWithExecutor, then the command is run inside
    an executor instance.

    Optionally: _cmdline_as_root() and _cmd_env() may be implemented, to
    customize behaviour.

    See the main entry point ``BaseExternalTask._execute()`` for details of the
    flow.
    """

    # If CAPTURE_OUTPUT_FILENAME is not None: self.run_via_executor()
    # creates a file in the cwd of the command to save the stdout. The file
    # is available in self.upload_artifacts().
    CAPTURE_OUTPUT_FILENAME: str | None = None

    CMD_LOG_SEPARATOR = "--------------------"
    CMD_LOG_FILENAME = "cmd-output.log"

    @override
    def run(self, execute_directory: Path) -> WorkRequestResults:
        """
        Run the command returned by ``self._cmdline()``.

        The command is executed with ``execute_directory`` as the current
        directory. That directory is a dedicated temporary directory setup
        by the worker for use by the task.

        When RunCommandTask is combined with BaseTaskWithExecutor, then
        the command is executed inside the executor instance.

        .. note::

          If the member variable CAPTURE_OUTPUT_FILENAME is set:
          create a file with its name with the stdout of the command. Otherwise,
          the stdout of the command is saved in self.CMD_LOG_FILENAME together
          with stderr's output.
        """
        cmd = self._cmdline()
        self.logger.info("Executing: %s", " ".join(cmd))

        returncode = self.run_cmd(
            cmd,
            execute_directory,
            env=self._cmd_env(),
            run_as_root=self._cmdline_as_root(),
            capture_stdout_filename=self.CAPTURE_OUTPUT_FILENAME,
        )

        self.logger.info("%s exited with code %s", cmd[0], returncode)

        return self.task_result(returncode, execute_directory)

    def task_result(
        self,
        returncode: int | None,
        execute_directory: Path,  # noqa: ARG002, U100
    ) -> WorkRequestResults:
        """
        Sub-tasks can evaluate if the task was a success or failure.

        By default, return SUCCESS if returncode is 0, FAILURE otherwise.
        Sub-classes can re-implement it.

        :param returncode: return code of the command, or None if aborted
        :param execute_directory: directory with the output of the task
        :return: WorkRequestResults.
        """
        if returncode == 0:
            return WorkRequestResults.SUCCESS
        return WorkRequestResults.FAILURE

    @abstractmethod
    def _cmdline(self) -> list[str]:
        """Return the command to execute, as a list of program arguments."""

    def _cmd_env(self) -> dict[str, str]:
        """Return the (minimal) environment to execute the command under."""
        env = {
            "PATH": os.environ["PATH"],
        }
        env.update(self.worker_environment_variables)
        return env

    @staticmethod
    def _cmdline_as_root() -> bool:
        """Run _cmdline() as root."""
        return False

    @staticmethod
    def _write_utf8(file: BinaryIO, text: str) -> None:
        file.write(text.encode("utf-8", errors="replace") + b"\n")

    def _write_popen_result(
        self, file: BinaryIO, p: subprocess.Popen[AnyStr]
    ) -> None:
        self._write_utf8(file, f"\naborted: {self.aborted}")
        self._write_utf8(file, f"returncode: {p.returncode}")

    def run_cmd(
        self,
        cmd: list[str],
        working_directory: Path,
        *,
        env: dict[str, str] | None = None,
        run_as_root: bool = False,
        capture_stdout_filename: str | None = None,
    ) -> int | None:
        """
        Run cmd in working_directory. Create self.CMD_OUTPUT_FILE log file.

        If BaseExternalTask.aborted == True terminates the process.

        :param cmd: command to execute with its arguments.
        :param working_directory: working directory where the command
          is executed.
        :param run_as_root: if True, run the command as root. Otherwise,
          the command runs as the worker's user
        :param capture_stdout_filename: for some commands the output of the
          command is the output of stdout (e.g. lintian) and not a set of files
          generated by the command (e.g. sbuild). If capture_stdout is not None,
          save the stdout into this file. The caller can then use it.
        :return: returncode of the process or None if aborted
        """
        if env is None:
            env = self._cmd_env()

        with self.open_debug_log_file(
            self.CMD_LOG_FILENAME, mode="ab"
        ) as cmd_log:
            self._write_utf8(cmd_log, f"cmd: {shlex.join(cmd)}")

            out_file: BinaryIO
            if capture_stdout_filename:
                self._write_utf8(
                    cmd_log,
                    "output (contains stderr only, stdout was captured):",
                )
                capture_stdout = working_directory / capture_stdout_filename
                out_file = capture_stdout.open(mode="wb")
            else:
                self._write_utf8(
                    cmd_log, "output (contains stdout and stderr):"
                )
                out_file = cmd_log

            try:
                return self._run_cmd(
                    cmd,
                    working_directory,
                    env=env,
                    run_as_root=run_as_root,
                    cmd_log=cmd_log,
                    out_file=out_file,
                )
            finally:
                file_names = "\n".join(
                    str(file.relative_to(working_directory))
                    for file in sorted(working_directory.rglob("*"))
                )

                self._write_utf8(cmd_log, "\nFiles in working directory:")
                self._write_utf8(cmd_log, file_names)

                self._write_utf8(cmd_log, self.CMD_LOG_SEPARATOR)
                if capture_stdout_filename:
                    cmd_log.write(capture_stdout.read_bytes() + b"\n")
                    self._write_utf8(cmd_log, self.CMD_LOG_SEPARATOR)

                out_file.close()

    def _run_cmd(
        self,
        cmd: list[str],
        working_directory: Path,
        *,
        env: dict[str, str],
        run_as_root: bool,
        cmd_log: BinaryIO,
        out_file: BinaryIO,
    ) -> int | None:
        """
        Execute cmd.

        :param run_as_root: if using an executor: run command as root or user.
          If not using an executor and run_as_root is True: raise ValueError().
        :param cmd_log: save the command log (parameters, stderr, return code,
          stdout)
        :param out_file: save the command stdout (might be the same as
          cmd_log or a different file)

        :return: returncode of the process or None if aborted
        """
        # Need to flush or subprocess.Popen() overwrites part of it
        cmd_log.flush()
        out_file.flush()

        run_kwargs: dict[str, Any] = {
            "cwd": working_directory,
            "env": env,
            "stderr": cmd_log.fileno(),
            "stdout": out_file.fileno(),
        }

        if self.executor_instance:
            return self.executor_instance.run(
                cmd, run_as_root=run_as_root, **run_kwargs
            ).returncode

        if run_as_root:
            raise ValueError("run_as_root requires an executor")

        assert env is not None
        with self._temporary_directory("tmp") as tmpdir:
            tmpdir.chmod(0o755)  # Some tasks will root containers inside here
            env["TMPDIR"] = str(tmpdir)

            p = subprocess.Popen(cmd, start_new_session=True, **run_kwargs)

            process_group = os.getpgid(p.pid)

            while True:
                if self.aborted:
                    break

                try:
                    self._wait_popen(p, timeout=1)
                    break
                except subprocess.TimeoutExpired:
                    pass

        if self.aborted:
            self.logger.debug("Task (cmd: %s PID %s) aborted", cmd, p.pid)
            try:
                if not self._send_signal_pid(p.pid, signal.SIGTERM):
                    # _send_signal_pid failed probably because cmd finished
                    # after aborting and before sending the signal
                    #
                    # p.poll() to read the returncode and avoid leaving cmd
                    # as zombie
                    p.poll()

                    # Kill possible processes launched by cmd
                    self._send_signal_group(process_group, signal.SIGKILL)
                    self.logger.debug("Could not send SIGTERM to %s", p.pid)

                    self._write_popen_result(cmd_log, p)
                    return None

                # _wait_popen with a timeout=5 to leave 5 seconds of grace
                # for the cmd to finish after sending SIGTERM
                self._wait_popen(p, timeout=5)
            except subprocess.TimeoutExpired:
                # SIGTERM was sent and 5 seconds later cmd
                # was still running. A SIGKILL to the process group will
                # be sent
                self.logger.debug(
                    "Task PID %s not finished after SIGTERM", p.pid
                )

            # debusine sends a SIGKILL if:
            # - SIGTERM was sent to cmd AND cmd was running 5 seconds later:
            #   SIGTERM was not enough so SIGKILL to the group is needed
            # - SIGTERM was sent to cmd AND cmd finished: SIGKILL to the
            #   group to make sure that there are not processes spawned
            #   by cmd running
            # (note that a cmd could launch processes in a new group
            # could be left running)
            self._send_signal_group(process_group, signal.SIGKILL)
            self.logger.debug("Sent SIGKILL to process group %s", process_group)

            # p.poll() to set p.returncode and avoid leaving cmd
            # as a zombie process.
            # But cmd might be left as a zombie process: if cmd was in a
            # non-interruptable kernel call p.returncode will be None even
            # after p.poll() and it will be left as a zombie process
            # (until debusine worker dies and the zombie is adopted by
            # init and waited on by init). If this happened there we might be a
            # ResourceWarning from Popen.__del__:
            # "subprocess %s is still running"
            #
            # A solution would e to wait (p.waitpid()) that the
            # process finished dying. This is implemented in the unit test
            # to avoid the warning but not implemented here to not delay
            # the possible shut down of debusine worker
            p.poll()
            self._write_popen_result(cmd_log, p)
            self.logger.debug("Returncode for PID %s: %s", p.pid, p.returncode)

            return None
        else:
            # The cmd has finished. The cmd might have spawned
            # other processes. debusine will kill any alive processes.
            #
            # If they existed they should have been finished by cmd:
            # run_cmd() should not leave processes behind.
            #
            # Since the parent died they are adopted by init and on
            # killing them they are not zombie.
            # (cmd might have spawned new processes in a different process
            # group: if this is the case they will be left running)
            self._send_signal_group(process_group, signal.SIGKILL)

        self._write_popen_result(cmd_log, p)

        return p.returncode

    def _wait_popen(
        self, popen: subprocess.Popen[AnyStr], timeout: float
    ) -> int:
        return popen.wait(timeout)

    @staticmethod
    def _send_signal_pid(pid: int, signal: signal.Signals) -> bool:
        try:
            os.kill(pid, signal)
        except ProcessLookupError:
            return False

        return True

    @staticmethod
    def _send_signal_group(process_group: int, signal: signal.Signals) -> None:
        """Send signal to the process group."""
        try:
            os.killpg(process_group, signal)
        except ProcessLookupError:
            pass


class ExtraRepositoryMixin[
    TDER: BaseTaskDataWithExtraRepositories,
    DTDER: BaseDynamicTaskDataWithExtraRepositories,
](BaseExternalTask[TDER, DTDER], metaclass=ABCMeta):
    """Methods for configuring external APT repositories."""

    data: TDER

    #: Extra repositories configured in task data
    extra_repositories = ExtraRepositoriesInput()

    # Deb822 style sources:
    extra_repository_sources: list[Path]
    # One-line sources:
    extra_repository_lists: list[Path]
    # GPG keys (for apt < 2.3.10):
    extra_repository_keys: list[Path]

    @override
    def post_init(self) -> None:
        super().post_init()
        self.extra_repository_sources = []
        self.extra_repository_lists = []
        self.extra_repository_keys = []

    def supports_inline_signed_by(self, codename: str) -> bool:
        """Determine if codename supports inline signatures."""
        # apt < 2.3.10 has no support for keys inline in Signed-By
        return codename not in (
            "stretch",
            "buster",
            "bullseye",
            "xenial",
            "bionic",
            "focal",
        )

    def write_extra_repository_config(
        self, codename: str, destination: Path
    ) -> None:
        """
        Write ``extra_repositories`` config to files in destination.

        ``extra_repository_keys`` will be populated with keys to install
        into /etc/apt/keyrings, if needed.

        ``extra_repository_sources`` will be populated with deb822 sources
        files.
        """
        assert self.dynamic_data is not None

        inline_signed_by = self.supports_inline_signed_by(codename)
        build_arch = self.data.build_architecture
        filtered_repos = [
            repo
            for repo in (self.dynamic_data.extra_repositories or [])
            if repo.architectures is None or (build_arch in repo.architectures)
        ]
        for i, extra_repo in enumerate(filtered_repos):
            signed_by = None
            if not inline_signed_by:
                if extra_repo.signing_key:
                    path = destination / f"extra_apt_key_{i}.asc"
                    path.write_text(extra_repo.signing_key + "\n")
                    self.extra_repository_keys.append(path)
                    signed_by = f"/etc/apt/keyrings/{path.name}"

            path = destination / f"extra_repository_{i}.sources"
            path.write_text(
                extra_repo.as_deb822_source(signed_by_filename=signed_by)
            )
            self.extra_repository_sources.append(path)
