# Copyright 2024 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.models import WorkRequest
from debusine.server.tasks import BaseServerTask
from debusine.tasks import BaseTask, TaskConfigError
from debusine.tasks.models import TaskTypes


logger = logging.getLogger(__name__)


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


@shared_task
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
    task_name = work_request.task_name
    task_data = work_request.task_data
    work_request.mark_running()

    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 = BaseTask.class_from_name(TaskTypes.SERVER, task_name)(task_data)
    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
