# Copyright 2021-2022 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.

"""Schedule WorkRequests to Workers."""

import logging
from typing import Any
from uuid import uuid4

from celery import shared_task

import django.db
from django.conf import settings
from django.db import connection, transaction
from django.db.models import Q
from django.db.models.signals import post_save
from django.dispatch import receiver

from debusine.artifacts.models import CollectionCategory
from debusine.db.models import CollectionItem, WorkRequest, Worker
from debusine.server.celery import run_server_task
from debusine.server.collections.lookup import lookup_multiple, lookup_single
from debusine.server.workflows import Workflow
from debusine.tasks import TaskConfigError
from debusine.tasks.models import LookupMultiple, LookupSingle, TaskTypes
from debusine.tasks.server import TaskDatabaseInterface

logger = logging.getLogger(__name__)


class TaskDatabase(TaskDatabaseInterface):
    """Implementation of database interaction in worker tasks."""

    def __init__(self, work_request: WorkRequest) -> None:
        """Construct a :py:class:`TaskDatabase`."""
        self.work_request = work_request

    def lookup_single_artifact(
        self,
        lookup: LookupSingle,
        default_category: CollectionCategory | None = None,
    ) -> int:
        """Look up a single artifact using :ref:`lookup-single`."""
        return lookup_single(
            lookup=lookup,
            workspace=self.work_request.workspace,
            user=self.work_request.created_by,
            default_category=default_category,
            workflow_root=self.work_request.get_workflow_root(),
            expect_type=CollectionItem.Types.ARTIFACT,
        ).artifact.id

    def lookup_multiple_artifacts(
        self,
        lookup: LookupMultiple,
        default_category: CollectionCategory | None = None,
    ) -> list[int]:
        """Look up multiple artifacts using :ref:`lookup-multiple`."""
        return [
            result.artifact.id
            for result in lookup_multiple(
                lookup=lookup,
                workspace=self.work_request.workspace,
                user=self.work_request.created_by,
                default_category=default_category,
                workflow_root=self.work_request.get_workflow_root(),
                expect_type=CollectionItem.Types.ARTIFACT,
            )
        ]

    def lookup_single_collection(
        self,
        lookup: LookupSingle,
        default_category: CollectionCategory | None = None,
    ) -> int:
        """Look up a single collection using :ref:`lookup-single`."""
        return lookup_single(
            lookup=lookup,
            workspace=self.work_request.workspace,
            user=self.work_request.created_by,
            default_category=default_category,
            workflow_root=self.work_request.get_workflow_root(),
            expect_type=CollectionItem.Types.COLLECTION,
        ).collection.id


def _possible_to_schedule_for_worker(worker):
    if worker.is_busy():
        # Worker already has enough work requests assigned; no scheduling
        # needed
        return False

    try:
        Worker.objects.select_for_update(nowait=True).get(id=worker.id)
    except django.db.DatabaseError:
        logger.debug("[SCHEDULER] Failed to lock Worker %s", worker)
        return False

    return True


@receiver(post_save, sender=Worker, dispatch_uid="ScheduleForWorker")
def _worker_changed(sender, instance: Worker, **kwargs):  # noqa: U100
    if not getattr(settings, "DISABLE_AUTOMATIC_SCHEDULING", False):
        schedule_for_worker(instance)


def _assign_work_request_to_worker(
    work_request: WorkRequest, worker: Worker
) -> WorkRequest | None:
    """
    Attempt to assign a particular work request to the worker.

    :param work_request: the work request to try to assign.
    :param worker: the worker that needs a new work request to be
      assigned.
    :return: The assigned work request, or None if it could not be assigned.
    """
    try:
        task = work_request.get_task()
    except TaskConfigError as exc:
        logger.warning(  # noqa: G200
            "WorkRequest %s failed to configure, aborting it. "
            "Task data: %s Error: %s",
            work_request.id,
            work_request.task_data,
            exc,
        )
        work_request.mark_completed(WorkRequest.Results.ERROR)
        return None

    if not task.can_run_on(worker.metadata()):
        return None

    try:
        work_request = WorkRequest.objects.select_for_update(nowait=True).get(
            id=work_request.id
        )
    except django.db.DatabaseError:
        logger.debug(
            "[SCHEDULER] Failed to lock WorkRequest %s",
            work_request,
        )
        return None

    if work_request.worker:  # pragma: no cover
        # work_request did not have a worker assigned on the
        # initial qs but has one now - nothing to do
        return None

    try:
        dynamic_task_data = task.compute_dynamic_data(
            TaskDatabase(work_request)
        )
    except Exception as exc:
        logger.warning(  # noqa: G200
            "WorkRequest %s failed to pre-process task data, aborting it. "
            "Task data: %s Error: %s",
            work_request.id,
            work_request.task_data,
            exc,
        )
        work_request.mark_completed(WorkRequest.Results.ERROR)
        return None

    try:
        work_request._disable_signals = True  # type: ignore[attr-defined]

        work_request.dynamic_task_data = (
            None if dynamic_task_data is None else dynamic_task_data.dict()
        )
        work_request.save()

        work_request.assign_worker(worker)
    finally:
        del work_request._disable_signals  # type: ignore[attr-defined]

    if task.TASK_TYPE is TaskTypes.SERVER:
        # Use a slightly more informative task ID than the default.
        task_id = f"{work_request.task_name}_{work_request.id}_{uuid4()}"
        transaction.on_commit(
            lambda: run_server_task.apply_async(
                args=(work_request.id,), task_id=task_id
            )
        )

    return work_request


@transaction.atomic
def schedule_for_worker(worker: Worker) -> WorkRequest | None:
    """
    Schedule a new work request for the worker.

    :param worker: the worker that needs a new work request to be
      assigned.
    :return: The assigned work request. None if no suitable work request could
      be found.
    :rtype: WorkRequest | None.
    """
    if not _possible_to_schedule_for_worker(worker):
        return None

    worker_metadata = worker.metadata()
    tasks_allowlist = worker_metadata.get("tasks_allowlist", None)
    tasks_denylist = worker_metadata.get("tasks_denylist", [])

    # TODO: Optimization: pre-filter on task type as appropriate.
    qs = WorkRequest.objects.pending(exclude_assigned=True).filter(
        task_type__in=(
            TaskTypes.WORKER,
            TaskTypes.SERVER,
            TaskTypes.SIGNING,
        )
    )
    if tasks_allowlist is not None:
        qs = qs.filter(task_name__in=tasks_allowlist)
    elif tasks_denylist:
        qs = qs.exclude(task_name__in=tasks_denylist)

    for work_request in qs:
        assigned_work_request = _assign_work_request_to_worker(
            work_request, worker
        )
        if assigned_work_request is not None:
            return assigned_work_request

    return None


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

        if dynamic_task_data is not None:
            workflow.dynamic_task_data = dynamic_task_data.dict()
            workflow.save()
            orchestrator.dynamic_data = orchestrator.dynamic_task_data_type(
                **workflow.dynamic_task_data
            )

    return True


def populate_workflow(workflow: WorkRequest) -> bool:
    """Populate a workflow with initial work requests."""
    try:
        orchestrator = Workflow.from_name(workflow.task_name)(workflow)
    except TaskConfigError as exc:
        logger.warning(  # noqa: G200
            "Workflow orchestrator failed to configure for WorkRequest %s; "
            "aborting it. Error: %s",
            workflow.id,
            exc,
        )
        return False

    if not ensure_workflow_dynamic_data(workflow, orchestrator):
        return False

    try:
        orchestrator.populate()
    except Exception as exc:
        logger.warning(  # noqa: G200
            "Workflow orchestrator failed for WorkRequest %s. Error: %s",
            workflow.id,
            exc,
        )
        return False

    workflow.unblock_workflow_children()

    return True


def run_workflow_callback(work_request: WorkRequest) -> bool:
    """Run a workflow callback."""
    workflow: WorkRequest | None = work_request
    while workflow is not None and workflow.task_type != TaskTypes.WORKFLOW:
        workflow = workflow.parent
    if workflow is None:
        logger.warning(
            "Workflow callback %s is not contained in a workflow; aborting it.",
            work_request.id,
        )
        return False

    try:
        orchestrator = Workflow.from_name(workflow.task_name)(workflow)
    except TaskConfigError as exc:
        logger.warning(  # noqa: G200
            "Workflow orchestrator failed to configure for WorkRequest %s; "
            "aborting it. Error: %s",
            work_request.id,
            exc,
        )
        return False

    if not ensure_workflow_dynamic_data(workflow, orchestrator):
        return False

    try:
        orchestrator.callback(work_request)
    except Exception as exc:
        logger.warning(  # noqa: G200
            "Workflow orchestrator failed for WorkRequest %s. Error: %s",
            work_request.id,
            exc,
        )
        return False

    return True


@transaction.atomic
def schedule_internal() -> list[WorkRequest]:
    """
    Schedule internal work requests.

    These don't require a worker, but need some kind of action from the
    scheduler.
    """
    qs = WorkRequest.objects.pending().filter(
        Q(
            task_type=TaskTypes.INTERNAL,
            task_name__in={"synchronization_point", "workflow"},
        )
        | Q(task_type=TaskTypes.WORKFLOW),
    )

    result: list[WorkRequest] = []
    for work_request in qs:
        match (work_request.task_type, work_request.task_name):
            case (TaskTypes.INTERNAL, "synchronization_point"):
                # Synchronization points do nothing; when pending, they are
                # immediately marked as completed, thus unblocking work
                # requests that depend on them.
                work_request.mark_completed(WorkRequest.Results.SUCCESS)
                result.append(work_request)

            case (TaskTypes.INTERNAL, "workflow"):
                # Workflow callbacks are handled by the corresponding
                # workflow orchestrator.  For now we just run this inline;
                # if this becomes a performance problem then we can delegate
                # it to a Celery task instead.
                if run_workflow_callback(work_request):
                    work_request.mark_completed(WorkRequest.Results.SUCCESS)
                    result.append(work_request)
                else:
                    work_request.mark_completed(WorkRequest.Results.ERROR)

            case (TaskTypes.WORKFLOW, _):
                # Workflows are populated by the corresponding workflow
                # orchestrator.  For now we just run this inline; if this
                # becomes a performance problem then we can delegate it to a
                # Celery task instead.
                work_request.mark_running()
                if populate_workflow(work_request):
                    # The workflow is left running until all its children
                    # have finished.
                    result.append(work_request)
                else:
                    work_request.mark_completed(WorkRequest.Results.ERROR)

            case _ as unreachable:
                raise AssertionError(
                    f"Unexpected internal task name: {unreachable}"
                )

    return result


@receiver(post_save, sender=WorkRequest, dispatch_uid="ScheduleForWorkRequest")
def _work_request_changed(
    sender, instance: WorkRequest, **kwargs  # noqa: U100
):
    if not getattr(
        settings, "DISABLE_AUTOMATIC_SCHEDULING", False
    ) and not hasattr(instance, "_disable_signals"):
        # Run the scheduler on commit, but only once per savepoint.  This
        # uses undocumented Django internals, but if it goes wrong without
        # an obvious unit test failure, then the worst case should be some
        # extra calls to the scheduler, which is relatively harmless.
        if (set(connection.savepoint_ids), schedule_task.delay) not in [
            hook[:2] for hook in connection.run_on_commit
        ]:
            transaction.on_commit(schedule_task.delay)


def schedule() -> list[WorkRequest]:
    """
    Try to assign work requests to free workers.

    The function loops over workers that have no work request assigned
    and tries to assign them new work requests.

    :return: the list of work requests that got assigned.
    """
    available_workers = Worker.objects.waiting_for_work_request()

    result = list(schedule_internal())
    for available_worker in available_workers:
        work_request = schedule_for_worker(available_worker)
        if work_request:
            result.append(work_request)

    return result


@shared_task()
def schedule_task():
    """Run the scheduler as a celery task."""
    schedule()
