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

logger = logging.getLogger(__name__)


class CeleryWorkerRequiresServerTask(Exception):
    """Only server tasks can run on Celery workers."""


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


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task  # type: ignore[misc]
def run_server_task(work_request_id: int) -> bool:
    """Run a :class:`BaseTask` via Celery."""
    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

    task_name = work_request.task_name
    work_request.mark_running()

    context.reset()
    work_request.set_current()

    if work_request.task_type != TaskTypes.SERVER:
        logger.error("Task: %s cannot run on a Celery worker", task_name)
        work_request.mark_completed(WorkRequest.Results.ERROR)
        raise CeleryWorkerRequiresServerTask

    try:
        task = work_request.get_task()
    except ValueError:
        logger.error("Task: %s does not exist", task_name)
        work_request.mark_completed(WorkRequest.Results.ERROR)
        raise
    except TaskConfigError:
        logger.error("Task: %s failed to configure", task_name)
        work_request.mark_completed(WorkRequest.Results.ERROR)
        raise
    assert isinstance(task, BaseServerTask)

    task.set_work_request(work_request)

    try:
        with ExitStack() as stack:
            if not task.TASK_MANAGES_TRANSACTIONS:
                stack.enter_context(transaction.atomic())
            result = task.execute_logging_exceptions()
    except Exception:
        logger.error("Task: %s failed to execute", task_name)
        work_request.mark_completed(WorkRequest.Results.ERROR)
        raise
    else:
        if task.aborted:
            logger.info("Task: %s has been aborted", task_name)
            # No need to update DB state
            return False
        else:
            work_request.mark_completed(
                WorkRequest.Results.SUCCESS
                if result
                else WorkRequest.Results.FAILURE
            )
            return result
