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

"""Server-side task interface."""

import inspect
import logging
import tempfile
from abc import ABCMeta, abstractmethod
from collections import defaultdict
from pathlib import Path
from typing import (
    Any,
    BinaryIO,
    ClassVar,
    IO,
    Literal,
    TYPE_CHECKING,
    TextIO,
    Union,
    overload,
    override,
)

import debusine.tasks.tags as ttags
import debusine.worker.tags as wtags
from debusine.artifacts import WorkRequestDebugLogs
from debusine.artifacts.models import TaskTypes, WorkRequestResults
from debusine.db.models import (
    Artifact,
    ArtifactRelation,
    Collection,
    WorkRequest,
    Workspace,
)
from debusine.db.tasks.inputs import TaskConfigurationCollection
from debusine.db.tasks.resolver import DBFieldResolver
from debusine.server.tags import (
    make_task_provided_tagset,
    make_task_required_tagset,
)
from debusine.tasks import BaseExternalTask, TaskConfigError, inputs
from debusine.tasks.models import (
    ActionRecordInTaskHistory,
    BaseDynamicTaskData,
    BaseTaskData,
    EventReaction,
)
from debusine.utils import extract_generic_type_arguments

if TYPE_CHECKING:
    from _typeshed import OpenBinaryModeWriting, OpenTextModeWriting


class DBTask[TD: BaseTaskData, DTD: BaseDynamicTaskData = BaseDynamicTaskData](
    inputs.TaskFieldContainer, metaclass=ABCMeta
):
    """
    Base class for server-side task-specific logic.

    A DBTask object supports Debusine in overseeing the task execution, with
    operations like applying task configuration, resolving its task data and
    computing scheduler tags.

    This gives the server a single API to access logic that is specific to all
    supported types of tasks, whether they are run server-side or worker-side.
    """

    #: 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

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

    #: Work request backing this task in the DB
    work_request: WorkRequest

    #: Collection to use to look up task configuration
    task_configuration = TaskConfigurationCollection(
        stage=inputs.Stage.PENDING, internal=True
    )

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

        Used by DBTask.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, DBTask)

        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, work_request: WorkRequest) -> None:
        """Initialize the task."""
        super().__init__()
        self.work_request = work_request

        #: 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")

        #: Validate task data and dynamic data
        try:
            self.data = self.task_data_type(**work_request.used_task_data)
            self.dynamic_data = (
                None
                if work_request.dynamic_task_data is None
                else self.dynamic_task_data_type(
                    **work_request.dynamic_task_data
                )
            )
        except ValueError as exc:
            raise TaskConfigError(None, original_exception=exc)

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

        # fetch_input() add the downloaded artifacts. Used by
        # `DBTask._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._debug_log_files_directory: (
            None | (tempfile.TemporaryDirectory[str])
        ) = None

    @property
    def workspace(self) -> Workspace:
        """
        Return the workspace for this task.

        This is a shortcut for ``self.work_request.workspace``.
        """
        return self.work_request.workspace

    def post_init(self) -> None:
        """Specific post-init code."""
        # Restore values of task input fields stored in dynamic data
        if self.dynamic_data is not None:
            self.restore_field_values(self.dynamic_data)

    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])

    @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 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
        """
        ws = self.work_request.workspace

        provided = ttags.TaskProvidedTags()
        provided.add(
            ttags.ProvenanceProvided.SYSTEM,
            {
                ttags.SCOPE_PREFIX + ws.scope.name,
                ttags.WORKSPACE_PREFIX + f"{ws.scope.name}:{ws.name}",
            },
        )
        for group in self.work_request.created_by.debusine_groups.values(
            "scope__name", "workspace__name", "name"
        ):
            provided.add(
                ttags.ProvenanceProvided.SYSTEM,
                {
                    ttags.GROUP_PREFIX
                    + group["scope__name"]
                    + ":"
                    + (group["workspace__name"] or "")
                    + ":"
                    + group["name"]
                },
            )

        required = wtags.TaskRequiredTags()

        # The worker must have code that can handle this task version.
        task_version_tag = (
            wtags.TASK_PREFIX + f"{self.TASK_TYPE.lower()}:{self.name}:"
            f"version:{self.TASK_VERSION}"
        )
        # The worker must be of a type that can run this task.
        match self.TASK_TYPE:
            case TaskTypes.WORKER:
                required.add(
                    wtags.ProvenanceRequired.SYSTEM,
                    {wtags.WORKER_TYPE_EXTERNAL, task_version_tag},
                )
            case TaskTypes.SERVER:
                required.add(
                    wtags.ProvenanceRequired.SYSTEM,
                    {wtags.WORKER_TYPE_CELERY, task_version_tag},
                )
            case TaskTypes.SIGNING:
                required.add(
                    wtags.ProvenanceRequired.SYSTEM,
                    {wtags.WORKER_TYPE_SIGNING, task_version_tag},
                )
            case _:
                # Other tasks should never be dispatched to workers, but
                # make sure of that by requiring a tag that will never
                # exist.
                required.add(
                    wtags.ProvenanceRequired.SYSTEM,
                    {wtags.WORKER_TYPE_NOT_ASSIGNABLE},
                )

        return provided, required

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

        This method requires task input fields resolved for stage RUNNING.

        The computed dynamic data contains saved task input fields. When
        overriding this method without calling super, one needs to make sure
        that ``self.save_field_values`` is still called.

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

    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

    @abstractmethod
    def _execute(self) -> WorkRequestResults:
        """
        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.

        :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.
        """

    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["DBTask['Any', 'Any']"]:
        """
        Return class for :param task_name (case-insensitive).

        :param task_type: type of task to look up

        __init_subclass__() registers DBTask subclasses' into
        DBTask._sub_tasks.
        """
        if (registry := DBTask._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

    @staticmethod
    def is_valid_task_name(task_type: TaskTypes, task_name: str) -> bool:
        """Return True if task_name is registered (its class is imported)."""
        if (registry := DBTask._sub_tasks.get(task_type)) is None:
            return False
        return task_name.lower() in registry

    @staticmethod
    def task_names(task_type: TaskTypes) -> list[str]:
        """Return list of sub-task names."""
        return sorted(DBTask._sub_tasks[task_type])

    @staticmethod
    def is_worker_task(task_name: str) -> bool:
        """Check if task_name is a task that can run on external workers."""
        return task_name.lower() in DBTask._sub_tasks[TaskTypes.WORKER]

    @staticmethod
    def worker_task_names() -> list[str]:
        """Return list of sub-task names not of type TaskTypes.SERVER."""
        return sorted(DBTask._sub_tasks[TaskTypes.WORKER].keys())

    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 = []

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

        artifact_ids.extend(self.get_input_artifacts_ids_from_fields())

        return artifact_ids

    def get_input_artifacts_ids_from_fields(self) -> list[int]:
        """Return input artifact IDs from serialized task field values."""
        return []

    def get_label(self) -> str:
        """Return a short human-readable label for the task."""
        if self.dynamic_data is None:
            return self.name

        if self.dynamic_data.parameter_summary is None:
            return self.name

        return f"{self.name}: {self.dynamic_data.parameter_summary}"

    def get_description(self) -> str:
        """Return a description for the task."""
        return self.__doc__ or self.get_label()

    def get_event_reactions(
        self,
        event_name: Literal[
            "on_creation",
            "on_unblock",
            "on_assignment",
            "on_success",
            "on_failure",
        ],
    ) -> list[EventReaction]:
        """
        Return event reactions for this task.

        This allows tasks to provide actions that are processed by the
        server at various points in the lifecycle of the work request.
        """
        event_reactions: list[EventReaction] = []
        if event_name in {"on_success", "on_failure"}:
            event_reactions.append(ActionRecordInTaskHistory())
        return event_reactions

    @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.
        """
        if self._debug_log_files_directory is None:
            self._debug_log_files_directory = tempfile.TemporaryDirectory(
                prefix="debusine-task-debug-log-files-"
            )

        debug_file = Path(self._debug_log_files_directory.name) / filename
        return debug_file.open(mode)

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

        The logs might exist in self._debug_log_files_directory and were
        added via self.open_debug_log_file() or self.create_debug_log_file().

        For each self._source_artifacts_ids: create a relation from
        WorkRequestDebugLogs to source_artifact_id.
        """
        if self._debug_log_files_directory is None:
            return

        work_request_debug_logs = WorkRequestDebugLogs.create(
            files=Path(self._debug_log_files_directory.name).glob("*")
        )

        artifact = Artifact.objects.create_from_local_artifact(
            work_request_debug_logs,
            self.work_request.workspace,
            created_by_work_request=self.work_request,
        )

        for source_artifact_id in self._source_artifacts_ids:
            ArtifactRelation.objects.create(
                artifact=artifact,
                target_id=source_artifact_id,
                type=ArtifactRelation.Relations.RELATES_TO,
            )

        self._debug_log_files_directory.cleanup()
        self._debug_log_files_directory = None

    def resolve_inputs(self, stage: "inputs.Stage") -> None:
        """Resolve input fields."""
        resolver = DBFieldResolver(self)
        resolver.resolve_inputs(self, stage=stage)

    @abstractmethod
    def get_task_configuration_subject_context(
        self,
    ) -> tuple[str | None, str | None]:
        """Get subject and configuration context task configuration lookups."""

    def apply_task_configuration(self) -> None:
        """
        Apply task configuration to this task and work request.

        Task scheduler tags are also computed.

        Previously computed task configuration, scheduler tasks and dynamic
        data are reset.

        This method requires task input fields resolved for stage PENDING.
        """
        import debusine.tasks.tags as ttags
        import debusine.worker.tags as wtags
        from debusine.server.collections.debusine_task_configuration import (
            Changes,
            apply_configuration,
        )

        self.work_request.configured_task_data = None
        self.work_request.dynamic_task_data = None
        self.work_request.scheduler_tags_provided = []
        self.work_request.scheduler_tags_required = []

        # Lookup the task-configuration collection to use
        config_collection: Collection | None = (
            self.task_configuration.collection
            if self.task_configuration is not None
            else None
        )

        # Compute the configured task data

        # Start from the unconfigured task data
        # Use task data from the work request, as if configured_task_data is
        # present, it has likely been used to populate self.data
        self.work_request.configured_task_data = (
            self.work_request.task_data.copy()
        )
        # Apply configuration if we have a config collection
        changes: Changes | None = None
        if config_collection is not None:
            changes = apply_configuration(
                self.work_request.configured_task_data, config_collection, self
            )

        # Compute scheduler tags
        final_provided = make_task_provided_tagset()
        final_required = make_task_required_tagset()
        if changes:
            final_provided.add(
                ttags.ProvenanceProvided.WORKSPACE, changes.provide_tags
            )
            final_required.add(
                wtags.ProvenanceRequired.WORKSPACE, changes.require_tags
            )

        provided, required = self.compute_scheduler_tags()

        final_required.add_candidates(required)
        final_required.finalize()
        self.work_request.scheduler_tags_required = sorted(final_required.tags)

        # Leave computing provided tags last, so that if anything else fails
        # the work request does not accidentally become schedulable
        final_provided.add_candidates(provided)
        final_provided.finalize()
        self.work_request.scheduler_tags_provided = sorted(final_provided.tags)


class DBWorkerTask[TD: BaseTaskData, DTD: BaseDynamicTaskData](
    DBTask[TD, DTD], metaclass=ABCMeta
):
    """Base class for DB tasks of type WORKER."""

    TASK_TYPE = TaskTypes.WORKER

    @override
    def _execute(self) -> WorkRequestResults:  # noqa: C901
        raise AssertionError("Worker tasks cannot be executed server-side")


class DBProxyTask[TD: BaseTaskData, DTD: BaseDynamicTaskData](
    DBTask[TD, DTD], metaclass=ABCMeta
):
    """A :py:class:`DBTask` that wraps a non-db task."""

    #: Wrapped BaseExternalTask subclass
    task_class: type[BaseExternalTask[TD, DTD]]
    #: Wrapped BaseExternalTask
    task: BaseExternalTask[TD, DTD]

    @override
    def __init_subclass__(cls, **kwargs: Any) -> None:
        """Disable automatic registration."""
        # Do not register proxies in the type system automatically: it is done
        # by populate_task_registry

    @classmethod
    def populate_task_registry(cls) -> None:
        """Populate DBWorkerTask types from the BaseExternalTask registry."""
        registry = DBTask._sub_tasks[cls.TASK_TYPE]
        for name, task_cls in BaseExternalTask._sub_tasks[
            cls.TASK_TYPE
        ].items():
            # Skip names already registered, so that the method can be run
            # twice to catch newly registered classes
            if name in registry:
                continue
            db_task_cls = type(
                f"DB{name}",
                (cls,),
                {
                    "TASK_VERSION": task_cls.TASK_VERSION,
                    "name": task_cls.name,
                    "task_data_type": task_cls.task_data_type,
                    "dynamic_task_data_type": task_cls.dynamic_task_data_type,
                    "task_class": task_cls,
                    "__doc__": task_cls.__doc__,
                },
            )
            registry[name] = db_task_cls

    def __init__(self, work_request: WorkRequest) -> None:
        """Initialize the task with the wrapped BaseExternalTask."""
        super().__init__(work_request)
        self.task = self.task_class(
            task_data=self.work_request.used_task_data,
            dynamic_task_data=self.work_request.dynamic_task_data,
            wrapped=True,
        )
        self.child_task_field_container = self.task

    @override
    @classmethod
    def get_provided_worker_tags(cls) -> set[str]:
        tags = super().get_provided_worker_tags()
        tags |= cls.task_class.get_provided_worker_tags()
        return tags

    @override
    def get_input_artifacts_ids_from_fields(self) -> list[int]:
        """Return input artifact IDs from wrapped task field values."""
        return self.task.get_input_artifacts_ids()

    @override
    def get_task_configuration_subject_context(
        self,
    ) -> tuple[str | None, str | None]:
        """Get subject and configuration context task configuration lookups."""
        return self.task.get_subject(), self.task.get_configuration_context()

    @override
    def compute_scheduler_tags(
        self,
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        provided, required = super().compute_scheduler_tags()
        tprovided, trequired = self.task.compute_scheduler_tags()
        provided.add_candidates(tprovided)
        required.add_candidates(trequired)
        return provided, required

    @override
    def compute_dynamic_data(self) -> DTD:
        dynamic_data = self.task.compute_dynamic_data()
        self.save_field_values(dynamic_data)
        return dynamic_data


class DBWorkerProxyTask[TD: BaseTaskData, DTD: BaseDynamicTaskData](
    DBWorkerTask[TD, DTD], DBProxyTask[TD, DTD], metaclass=ABCMeta
):
    """A :py:class:`DBTask` that wraps a worker task."""

    @override
    @classmethod
    def populate_task_registry(cls) -> None:
        # Ensure worker tasks are registered before creating their proxy
        # classes
        import debusine.tasks  # noqa: F401

        super().populate_task_registry()


class DBSigningProxyTask[TD: BaseTaskData, DTD: BaseDynamicTaskData](
    DBProxyTask[TD, DTD], metaclass=ABCMeta
):
    """A :py:class:`DBTask` that wraps a worker task."""

    TASK_TYPE = TaskTypes.SIGNING

    @override
    @classmethod
    def populate_task_registry(cls) -> None:
        # Ensure signing tasks are registered before creating their proxy
        # classes
        import debusine.signing.tasks  # noqa: F401

        super().populate_task_registry()

    @override
    def _execute(self) -> WorkRequestResults:  # noqa: C901
        raise AssertionError("Signing tasks cannot be executed server-side")


DBWorkerProxyTask.populate_task_registry()
DBSigningProxyTask.populate_task_registry()
