# 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 List, Optional

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

from debusine.db.models import WorkRequest, Worker
from debusine.tasks import Task, TaskConfigError

logger = logging.getLogger(__name__)


def _possible_to_schedule_for_worker(worker):
    if WorkRequest.objects.running_or_pending_exists(worker):
        # Worker has an already assigned work_request in pending or running:
        # 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)


@transaction.atomic()
def schedule_for_worker(worker: Worker) -> Optional[WorkRequest]:
    """
    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: Optional[WorkRequest].
    """
    if not _possible_to_schedule_for_worker(worker):
        return

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

    qs = WorkRequest.objects.pending(exclude_assigned=True)
    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:
        task = Task.class_from_name(work_request.task_name)()
        try:
            task.configure(work_request.task_data)
        except TaskConfigError:
            logger.warning(
                "WorkRequest %s failed to configure, aborting it."
                "Task data: %s",
                work_request.id,
                work_request.task_data,
            )
            work_request.mark_aborted()
            continue

        if task.can_run_on(worker_metadata):
            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,
                )
                continue

            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
                continue

            work_request.assign_worker(worker)
            return work_request


@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):
        schedule()


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 = []
    for available_worker in available_workers:
        work_request = schedule_for_worker(available_worker)
        if work_request:
            result.append(work_request)

    return result
