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

"""Celery integration for debusine tasks."""

import logging
from contextlib import ExitStack

from celery import shared_task
from django.db import transaction

from debusine.artifacts.models import TaskTypes
from debusine.db.context import context
from debusine.db.models import WorkRequest
from debusine.server.tasks import BaseServerTask
from debusine.tasks import TaskConfigError
from debusine.tasks.models import OutputData, OutputDataError

logger = logging.getLogger(__name__)


class WorkRequestNotPending(Exception):
    """We only run pending work requests."""


class ServerTaskRunError(Exception):
    """Running a server task failed."""

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


def _run_server_task_or_error(work_request: WorkRequest) -> WorkRequest.Results:
    """Run a server task, raising :py:class:`ServerTaskRunError` on errors."""
    if work_request.task_type != TaskTypes.SERVER:
        raise ServerTaskRunError(
            "Cannot run on a Celery worker", "wrong-task-type"
        )

    task_name = work_request.task_name
    try:
        task = work_request.get_task()
    except ValueError as exc:
        raise ServerTaskRunError(str(exc), "setup-failed") from exc
    except TaskConfigError as exc:
        raise ServerTaskRunError(
            f"Failed to configure: {exc}", "configure-failed"
        ) from exc
    assert isinstance(task, BaseServerTask)

    try:
        with ExitStack() as stack:
            if not task.TASK_MANAGES_TRANSACTIONS:
                stack.enter_context(transaction.atomic())
            result = task.execute_logging_exceptions()
    except Exception as exc:
        raise ServerTaskRunError(
            f"Execution failed: {exc}", "execute-failed"
        ) from exc
    else:
        if task.aborted:
            logger.info("Task: %s has been aborted", task_name)
            # No need to update DB state
            return WorkRequest.Results.FAILURE
        else:
            # Lock the whole workflow and mark the work request as
            # completed.
            with transaction.atomic():
                WorkRequest.objects.filter(
                    id=work_request.id
                ).lock_workflows_for_update()
                work_request.mark_completed(result)
                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]
# Unlike some other Celery tasks, this must _not_ be decorated with
# @transaction.atomic, because it needs to do finer-grained transaction
# control.
def run_server_task(work_request_id: int) -> WorkRequest.Results:
    """
    Run a :py:class:`BaseServerTask` via Celery.

    :param work_request_id: The ID of the work request to run.
    :raises WorkRequest.DoesNotExist: if ``work_request_id`` does not exist.
    :raises WorkRequestNotPending: if the given work request's status is not
      "pending".
    :return: SUCCESS if the server task succeeded; ERROR if it failed or raised
      an exception.
    """
    try:
        work_request = WorkRequest.objects.get(pk=work_request_id)
    except WorkRequest.DoesNotExist:
        logger.error("Work request %d does not exist", work_request_id)
        raise

    if work_request.status != WorkRequest.Statuses.PENDING:
        logger.error(
            "Work request %d is in status %s, not pending",
            work_request_id,
            work_request.status,
        )
        raise WorkRequestNotPending

    # We would normally lock the whole workflow as close to the start of the
    # operation as possible, holding that lock until the end of the
    # transaction to minimize the risk of deadlocks when updating multiple
    # work requests in succession.  However, server tasks may take a long
    # time to run, and some of them manage their own transactions.
    #
    # WorkRequest.mark_running has no cascading effects on other work
    # requests, and at this point we're in Django's default autocommit mode,
    # so the worst case at this point is that it fails to obtain a lock on
    # this work request; it cannot cause a deadlock.
    work_request.mark_running()

    context.reset()
    work_request.set_current()

    try:
        return _run_server_task_or_error(work_request)
    except ServerTaskRunError as exc:
        # Lock the whole workflow and record the error.
        with transaction.atomic():
            logger.error(
                "Error running work request %s/%s (%s): %s",
                work_request.task_type,
                work_request.task_name,
                work_request.id,
                exc.message,
            )
            WorkRequest.objects.filter(
                id=work_request.id
            ).lock_workflows_for_update()
            work_request.mark_completed(
                WorkRequest.Results.ERROR,
                output_data=OutputData(
                    errors=[OutputDataError(message=exc.message, code=exc.code)]
                ),
            )
            return WorkRequest.Results.ERROR
