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

"""Schedule WorkRequests to Workers."""

import logging
from operator import attrgetter
from uuid import uuid4

import django.db
import pgtransaction
from celery import shared_task
from django.conf import settings
from django.db import OperationalError, transaction
from django.db.models import (
    DateTimeField,
    Exists,
    F,
    OuterRef,
    Q,
    QuerySet,
    Window,
)
from django.db.models.fields.json import KT
from django.db.models.functions import Cast, RowNumber
from django.utils import timezone

import debusine.tasks.tags as ttags
import debusine.worker.tags as wtags
from debusine.artifacts.models import TaskTypes
from debusine.db.context import context
from debusine.db.models import WorkRequest, Worker
from debusine.db.models.work_requests import (
    SkipWorkRequest,
    compute_workflow_last_activity,
    compute_workflow_runtime_status,
    workflow_ancestors,
)
from debusine.db.transactions import atomic
from debusine.project.celery import (
    apply_async_on_commit,
    celery_tasks_in_progress,
)
from debusine.server.celery import run_server_task
from debusine.server.workflows.celery import run_workflow_task
from debusine.tasks import TaskConfigError
from debusine.tasks.models import Confirmation, OutputData, OutputDataError

logger = logging.getLogger(__name__)


class AssignToWorkerError(Exception):
    """Assigning a work request to a worker failed."""

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


def _possible_to_schedule_for_worker(worker: Worker) -> bool:
    if worker.is_at_capacity():
        # 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


@transaction.atomic
@context.local()
@context.disable_permission_checks()
def _assign_work_request_to_worker(
    work_request: WorkRequest, worker: Worker, locked_workflow_ids: set[int]
) -> 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; must already be locked for update.
    :return: The assigned work request, or None if it could not be assigned.
    :raises AssignToWorkerError: if the work request is broken.
    """
    # Enable tag-based scheduling by default for now, but allow to disable it
    # with a feature flag, in case there are issues with it.
    FEATURE_TAG_BASED_SCHEDULING = getattr(
        settings, "FEATURE_TAG_BASED_SCHEDULING", True
    )

    try:
        task = work_request.get_task()
    except TaskConfigError as exc:
        raise AssignToWorkerError(
            f"Failed to configure: {exc}", "configure-failed"
        ) from exc

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

    try:
        # We don't need to try to lock the whole workflow here, since
        # WorkRequest.assign_worker has no cascading effects on other work
        # requests.
        work_request = WorkRequest.objects.select_for_update(nowait=True).get(
            id=work_request.id
        )
    except django.db.DatabaseError:
        if work_request.workflow_root_id is not None:
            # The whole workflow is probably locked, so don't spam logs
            # trying each pending work request in it.
            locked_workflow_ids.add(work_request.workflow_root_id)
        logger.debug(
            "[SCHEDULER] Failed to lock WorkRequest %s"
            " (skipping this workflow)",
            work_request,
        )
        transaction.set_rollback(True)
        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:
        work_request.assign_worker(worker)
    except SkipWorkRequest:
        logger.debug(
            "Skipping work request %s due to on_assignment action", work_request
        )
        return None
    except Exception as exc:
        raise AssignToWorkerError(
            f"Failed to compute dynamic data: {exc}", "dynamic-data-failed"
        ) from exc

    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()}"
        apply_async_on_commit(
            run_server_task, args=(work_request.id,), task_id=task_id
        )

    return work_request


def _log_work_request_assign_error(
    work_request: WorkRequest, exc: AssignToWorkerError
) -> None:
    """
    Log an error arisen when assigning a work request to a worker.

    This is in a separate function to make it easier to mock.
    """
    logger.error(
        "Error assigning work request %s/%s (%s): %s",
        work_request.task_type,
        work_request.task_name,
        work_request.id,
        exc.message,
    )


def _assign_work_request_to_worker_or_error(
    work_request: WorkRequest, worker: Worker, locked_workflow_ids: set[int]
) -> WorkRequest | None:
    """
    Attempt to assign a particular work request to the worker, handling errors.

    This needs special care around transaction management.  Assigning the
    worker shouldn't take a whole-workflow lock, because it doesn't need to
    and doing so would increase contention; but if the work request is
    broken then we must take a whole-workflow lock because marking the work
    request as completed might cascade to other work requests in the same
    workflow, and that must therefore be done in a separate savepoint from
    the assignment to avoid a potential deadlock.

    :param work_request: the work request to try to assign.
    :param worker: the worker that needs a new work request to be
      assigned; must already be locked for update.
    :return: The assigned work request, or None if it could not be assigned.
    """
    try:
        return _assign_work_request_to_worker(
            work_request, worker, locked_workflow_ids
        )
    except AssignToWorkerError as exc:
        with transaction.atomic():
            # _assign_work_request_to_worker runs in its own savepoint, so
            # any update locks it holds will be released in case an
            # exception is raised, allowing this to be called in the
            # exception handler without risking a deadlock on WorkRequest
            # locks.  We must still use nowait=True, both because we're in
            # the scheduler and can't block, and because we're within a
            # Worker update lock and don't want to risk something else
            # trying to acquire WorkRequest and Worker locks in the opposite
            # order.
            try:
                WorkRequest.objects.filter(
                    id=work_request.id
                ).lock_workflows_for_update(nowait=True)
            except django.db.DatabaseError:
                # Not ideal, but we'll try the same work request again
                # later.
                logger.debug(
                    "[SCHEDULER] Failed to lock WorkRequest %s to record error",
                    work_request,
                )
                transaction.set_rollback(True)
            else:
                _log_work_request_assign_error(work_request, exc)
                work_request.mark_completed(
                    WorkRequest.Results.ERROR,
                    output_data=OutputData(
                        errors=[
                            OutputDataError(message=exc.message, code=exc.code)
                        ]
                    ),
                )
        return None


def get_provided_worker_tags(worker: Worker) -> set[str]:
    """Get the set of tags provided by a worker."""
    from debusine.server.tags import make_worker_provided_tagset

    tagset = make_worker_provided_tagset()
    tagset.add(
        wtags.ProvenanceProvided.WORKER,
        worker.dynamic_metadata.get("provided_tags"),
    )
    tagset.add(
        wtags.ProvenanceProvided.ADMIN,
        worker.static_metadata.get("provided_tags"),
    )
    tagset.add(
        wtags.ProvenanceProvided.SYSTEM,
        {wtags.WORKER_TYPE_PREFIX + worker.worker_type},
    )
    tagset.finalize()
    return tagset.tags


def get_required_task_tags(worker: Worker) -> set[str]:
    """Get the set of task tags required by a worker."""
    from debusine.server.tags import make_worker_required_tagset

    tagset = make_worker_required_tagset()
    tagset.add(
        ttags.ProvenanceRequired.ADMIN,
        worker.static_metadata.get("required_tags"),
    )
    tagset.finalize()
    return tagset.tags


def get_candidate_work_requests(worker: Worker) -> QuerySet[WorkRequest]:
    """Get pending work requests that can be scheduled on this worker."""
    # Disable tag-based scheduling by default for now, but allow to enable it
    # with a feature flag so we can test it, before offering it for mainstream
    # adoption
    FEATURE_TAG_BASED_SCHEDULING = getattr(
        settings, "FEATURE_TAG_BASED_SCHEDULING", True
    )

    provided_worker_tags = get_provided_worker_tags(worker)
    required_task_tags = get_required_task_tags(worker)
    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)

    if FEATURE_TAG_BASED_SCHEDULING:
        qs = qs.filter(
            scheduler_tags_required__contained_by=list(provided_worker_tags),
            scheduler_tags_provided__contains=list(required_task_tags),
        )

    return qs


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

    :param worker: the worker that needs a new work request to be
      assigned.
    :param locked_workflow_ids: a set of workflows that have so far
      discovered to have been locked.
    :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

    for work_request in get_candidate_work_requests(worker):
        if work_request.workflow_root_id in locked_workflow_ids:
            continue
        assigned_work_request = _assign_work_request_to_worker_or_error(
            work_request, worker, locked_workflow_ids
        )
        if assigned_work_request is not None:
            return assigned_work_request

    return None


# We may encounter serialization failures, but if so it's OK to just let the
# next attempt catch up.
@atomic(isolation_level=pgtransaction.REPEATABLE_READ)
def update_workflows() -> None:
    """
    Update expensive properties of workflows.

    The ``workflow_last_activity_at`` and ``workflow_runtime_status`` fields
    of workflows are used in the web UI and need to be kept up to date.
    """
    try:
        need_update = WorkRequest.objects.filter(
            workflows_need_update=True
        ).lock_workflows_for_update()
    except OperationalError as exc:
        # Something like select_for_update(skip_locked=True) might be
        # better, allowing us to make some progress, but it's unclear how to
        # acquire a lock on only those work requests that are in workflows
        # that can be locked as a whole without blocking.
        logger.warning(
            "update_workflows() deferred due to a concurrent update: %s",
            exc,
        )
        return
    for workflow in workflow_ancestors(need_update):
        assert workflow.task_type == TaskTypes.WORKFLOW
        workflow.workflow_last_activity_at = compute_workflow_last_activity(
            workflow
        )
        workflow.workflow_runtime_status = compute_workflow_runtime_status(
            workflow
        )
        workflow.save()
    need_update.update(workflows_need_update=False)


def run_workflow_task_via_celery(work_request: WorkRequest) -> None:
    """Run a workflow callback or workflow via Celery."""
    # Use a slightly more informative task ID than the default.
    task_id = (
        f"{work_request.task_type.lower()}_{work_request.task_name}_"
        f"{work_request.id}_{uuid4()}"
    )
    apply_async_on_commit(
        run_workflow_task, args=(work_request.id,), task_id=task_id
    )


@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.
    """
    in_progress_task_results = celery_tasks_in_progress(
        "debusine.server.workflows.celery.run_workflow_task"
    )

    pending_wrs = WorkRequest.objects.pending()
    pending_synchronization_points = pending_wrs.synchronization_points()
    pending_workflow_callbacks = (
        pending_wrs.workflow_callbacks()
        .annotate(
            row_number=Window(
                RowNumber(), partition_by=F("parent"), order_by="id"
            )
        )
        # Only one workflow callback for a given workflow may run at once.
        .exclude(
            Exists(
                in_progress_task_results.filter(
                    first_task_arg__in=WorkRequest.objects.filter(
                        parent=OuterRef(OuterRef("parent"))
                    ),
                )
            )
        )
        .exclude(row_number__gt=1)
    )
    pending_workflows = (
        pending_wrs.filter(task_type=TaskTypes.WORKFLOW)
        .annotate(
            row_number=Window(
                RowNumber(), partition_by=F("workflow_root"), order_by="id"
            )
        )
        # Only one sub-workflow within a given root workflow may run at once.
        .exclude(
            Exists(
                in_progress_task_results.filter(
                    first_task_arg=OuterRef("workflow_root")
                )
            )
        )
        .exclude(row_number__gt=1)
    )
    pending_waits = pending_wrs.filter(task_type=TaskTypes.WAIT)
    pending_combined = WorkRequest.objects.filter(
        Q(id__in=pending_synchronization_points)
        | Q(id__in=pending_workflow_callbacks)
        | Q(id__in=pending_workflows)
        | Q(id__in=pending_waits)
    )

    try:
        pending_combined.lock_workflows_for_update(nowait=True)
    except django.db.DatabaseError:
        # Something like select_for_update(skip_locked=True) might be
        # better, allowing us to make some progress, but it's unclear how to
        # acquire a lock on only those work requests that are in workflows
        # that can be locked as a whole without blocking.  This should be
        # fine provided that nothing holds a lock for too long on workflows
        # that had pending internal work requests at the start of its
        # transaction.
        logger.debug(
            "[SCHEDULER] Failed to lock workflows with pending internal"
            " work requests"
        )
        return []

    workflow_roots: set[WorkRequest] = set()
    result: list[WorkRequest] = []
    for work_request in pending_combined:
        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.  This may take a while, so we
                # delegate it to a Celery task.
                work_request.mark_running()
                run_workflow_task_via_celery(work_request)
                result.append(work_request)

            case (TaskTypes.WORKFLOW, _):
                # Workflow population always happens via the orchestrator
                # for the root workflow.  Sub-workflow orchestrators are not
                # called directly from the scheduler, although if a
                # sub-workflow is pending then its root workflow may deal
                # with populating it.
                # Workflows always have a root.
                assert work_request.workflow_root is not None
                workflow_roots.add(work_request.workflow_root)

            case (TaskTypes.WAIT, _):
                # Wait tasks are marked running immediately, but otherwise
                # we do nothing; some later event will mark them as
                # completed.
                work_request.mark_running()
                result.append(work_request)

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

    for workflow_root in sorted(workflow_roots, key=attrgetter("id")):
        # Workflows are populated by the corresponding workflow
        # orchestrator.  This may take a while, so we delegate it to a
        # Celery task.
        workflow_root.mark_running()
        run_workflow_task_via_celery(workflow_root)
        result.append(workflow_root)

    return result


@transaction.atomic
def complete_wait_tasks() -> None:
    """Mark certain wait tasks as completed."""
    qs = (
        WorkRequest.objects.running()
        .filter(task_type=TaskTypes.WAIT)
        .filter(
            # Delay tasks complete automatically once their delay has
            # expired.
            Q(task_name="delay")
            # Confirm tasks with auto_confirm_if_no_failure=True complete
            # automatically if all their dependencies are successful.
            | (
                Q(
                    task_name="confirm",
                    task_data__auto_confirm_if_no_failure=True,
                )
                & ~Exists(
                    WorkRequest.objects.filter(
                        Q(reverse_dependencies=OuterRef("id"))
                        & ~Q(
                            status=WorkRequest.Statuses.COMPLETED,
                            result=WorkRequest.Results.SUCCESS,
                        )
                    )
                )
            )
        )
        .annotate(
            delay_until=Cast(
                KT("task_data__delay_until"),
                DateTimeField(blank=True, null=True),
            )
        )
        .exclude(task_name="delay", delay_until__gt=timezone.now())
    )

    try:
        qs.lock_workflows_for_update(nowait=True)
    except django.db.DatabaseError:
        # Something like select_for_update(skip_locked=True) might be
        # better, allowing us to make some progress, but it's unclear how to
        # acquire a lock on only those work requests that are in workflows
        # that can be locked as a whole without blocking.  This should be
        # fine provided that nothing holds a lock for too long on workflows
        # that had pending internal work requests at the start of its
        # transaction.
        logger.debug(
            "[SCHEDULER] Failed to lock workflows with pending internal"
            " work requests"
        )
        return

    for work_request in qs:
        match work_request.task_name:
            case "delay":
                work_request.mark_completed(WorkRequest.Results.SUCCESS)
            case "confirm":
                work_request.mark_completed(
                    WorkRequest.Results.SUCCESS,
                    output_data=OutputData(
                        confirmation=Confirmation(confirmed=True)
                    ),
                )
            case _ as unreachable:
                raise AssertionError(f"Unexpected task name: {unreachable}")


def schedule() -> list[WorkRequest]:
    """
    Perform work request scheduling.

    This runs frequently, and does a few different steps:

    * Keep various properties of workflows up to date.
    * Check pending work requests that can be acted on directly without a
      worker, as well as pending workflow instances and workflow callbacks
      that need to be orchestrated by a Celery task.
    * Loop over idle workers and try to assign pending work requests to
      them.

    :return: the list of work requests that got assigned.
    """
    update_workflows()

    available_workers = Worker.objects.waiting_for_work_request()

    result = list(schedule_internal())
    locked_workflow_ids: set[int] = set()
    for available_worker in available_workers:
        work_request = schedule_for_worker(
            available_worker, locked_workflow_ids
        )
        if work_request:
            result.append(work_request)

    complete_wait_tasks()

    return result


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task()  # type: ignore[untyped-decorator]
def schedule_task() -> None:
    """Run the scheduler as a celery task."""
    schedule()


def schedule_soon() -> None:
    """Ask the scheduler to run soon."""
    if not getattr(settings, "DISABLE_AUTOMATIC_SCHEDULING", False):
        apply_async_on_commit(schedule_task, once_per_savepoint=True)
