# 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 workflow orchestrators."""

import logging

from celery import shared_task
from django.db import transaction
from django.utils import timezone
from prometheus_client import Histogram

from debusine.artifacts.models import TaskTypes
from debusine.db.context import context
from debusine.db.metrics import emit_metric
from debusine.db.models import WorkRequest
from debusine.server.workflows.base import orchestrate_workflow
from debusine.tasks.models import OutputData, OutputDataError

logger = logging.getLogger(__name__)

workflow_orchestrator_duration = Histogram(
    "debusine_workflow_orchestrator_duration",
    "Histogram of how long workflow orchestrators took to run.",
    labelnames=("task_name",),
    # Orchestrators can take some time.  Try to provide a reasonable range
    # of buckets.
    buckets=(
        1.0,
        5.0,
        15.0,
        60.0,  # 1 minute
        300.0,  # 5 minutes
        900.0,  # 15 minutes
        3600.0,  # 1 hour
        10800.0,  # 3 hours
        21600.0,  # 6 hours
        43200.0,  # 12 hours
        86400.0,  # 1 day
        float("inf"),
    ),
)


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task  # type: ignore[untyped-decorator]
@transaction.atomic
def run_workflow_task(work_request_id: int) -> WorkRequest.Results:
    """Run a workflow callback or workflow in Celery."""
    try:
        work_request = (
            WorkRequest.objects.filter(id=work_request_id)
            .lock_workflows_for_update()
            .get()
        )
    except WorkRequest.DoesNotExist:
        logger.error("Work request %d does not exist", work_request_id)
        raise

    # orchestrate_workflow will mark the workflow running, but we need to
    # make sure that that was already done before dispatching the Celery
    # task in order that it's visible outside this transaction.
    if work_request.status != WorkRequest.Statuses.RUNNING:
        message = (
            f"Work request is in status {work_request.status}, not running"
        )
        logger.error(
            "Error running work request %s/%s (%s): %s",
            work_request.task_type,
            work_request.task_name,
            work_request.id,
            message,
        )
        work_request.mark_completed(
            WorkRequest.Results.ERROR,
            output_data=OutputData(
                errors=[OutputDataError(message=message, code="wrong-status")]
            ),
        )
        return WorkRequest.Results.ERROR

    if (
        work_request.task_type == TaskTypes.WORKFLOW
        and not work_request.is_workflow_root
    ):
        # This task must only be called to populate root workflows.
        # Sub-workflows are populated by their parent workflows instead.
        message = "Must be populated by its parent workflow instead"
        logger.error(
            "Error running work request %s/%s (%s): %s",
            work_request.task_type,
            work_request.task_name,
            work_request.id,
            message,
        )
        work_request.mark_completed(
            WorkRequest.Results.ERROR,
            output_data=OutputData(
                errors=[
                    OutputDataError(message=message, code="non-root-workflow")
                ]
            ),
        )
        return WorkRequest.Results.ERROR

    context.reset()
    work_request.set_current()
    start_time = timezone.now()

    result = (
        WorkRequest.Results.SUCCESS
        if orchestrate_workflow(work_request)
        else WorkRequest.Results.ERROR
    )

    emit_metric(
        metric_type="histogram",
        name="debusine_workflow_orchestrator_duration",
        labels={"task_name": work_request.task_name},
        value=(timezone.now() - start_time).total_seconds(),
    )

    return result
