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

"""Views for the server application: workers."""

import logging
from typing import Any

import pgtransaction
from celery import shared_task
from django.db import OperationalError
from django.db.models import Exists, OuterRef
from rest_framework import status
from rest_framework.request import Request
from rest_framework.response import Response

from debusine.db.context import context
from debusine.db.models import Token, WorkRequest, Worker
from debusine.db.models.permissions import anonymous_access
from debusine.db.models.work_requests import (
    CannotRetry,
    MAX_AUTOMATIC_RETRIES,
    WorkRequestRetryReason,
)
from debusine.db.transactions import atomic
from debusine.project.celery import (
    apply_async_on_commit,
    celery_tasks_in_progress,
)
from debusine.server.exceptions import DebusineAPIException
from debusine.server.serializers import (
    WorkRequestCompletedSerializer,
    WorkerRegisterSerializer,
    WorkerWorkRequestSerializer,
)
from debusine.server.views.base import BaseAPIView
from debusine.server.views.rest import (
    IsWorkRequestTokenAuthenticated,
    IsWorkerTokenAuthenticated,
)
from debusine.tasks.models import OutputData, OutputDataError

logger = logging.getLogger(__name__)


class RegisterView(BaseAPIView):
    """View used by workers to register to debusine."""

    def post(self, request: Request) -> Response:
        """Worker registers (sends token and fqdn)."""
        self.enforce(anonymous_access)
        worker_register = WorkerRegisterSerializer(data=request.data)

        if not worker_register.is_valid():
            raise DebusineAPIException(
                title="Cannot deserialize worker",
                validation_errors=worker_register.errors,
            )

        valid_activation_token = (
            isinstance(request.auth, Token)
            and request.auth.token_type == Token.Types.WORKER_ACTIVATION
        )

        if (
            request.headers.get("token") is not None
            and not valid_activation_token
        ):
            raise DebusineAPIException(
                title="Invalid worker activation token",
            )

        token_key = worker_register.validated_data['token']
        fqdn = worker_register.validated_data['fqdn']
        worker_type = worker_register.validated_data['worker_type']

        token_hash = Token._generate_hash(token_key)
        token, _ = Token.objects.get_or_create(
            token_type=Token.Types.WORKER, hash=token_hash
        )

        if valid_activation_token:
            assert isinstance(request.auth, Token)
            # Activate an existing worker by registering and automatically
            # enabling a token for it.
            activating_worker = request.auth.activating_worker
            assert activating_worker is not None
            activating_worker.activation_token = None
            activating_worker.token = token
            activating_worker.save()
            token.refresh_from_db()
            token.enable()
            request.auth.delete()
        else:
            Worker.objects.create_with_fqdn(
                fqdn, token=token, worker_type=worker_type
            )

        logger.info('Client registered. Token key: %s', token_key)

        return Response(status=status.HTTP_201_CREATED)


class GetNextWorkRequestView(BaseAPIView):
    """View used by workers to request a task."""

    permission_classes = [IsWorkerTokenAuthenticated]

    def select_work_request(self, worker: Worker) -> WorkRequest | None:
        """Pick our work request from the assigned queue."""
        running_work_requests = WorkRequest.objects.running(worker=worker)

        # If (1) a worker has recently requested that a work request should
        # be marked as completed, but the server hasn't finished dealing
        # with that yet, and (2) a websocket connection was restarted, then
        # we might find ourselves here.  In that case, we don't want to
        # interfere with a work request that's very nearly finished.
        in_progress_task_results = celery_tasks_in_progress(
            "debusine.server.views.workers.update_work_request_as_completed"
        )
        if running_work_requests.filter(
            Exists(
                in_progress_task_results.filter(first_task_arg=OuterRef("pk"))
            )
        ).exists():
            return None

        # Did the worker fail before producing anything?
        # Then it's safe to just re-attempt execution
        for running in running_work_requests.reassignable(worker=worker):
            workflow_data = running.workflow_data
            if workflow_data.retry_count >= MAX_AUTOMATIC_RETRIES:
                logger.info(
                    (
                        "Cannot retry previously-running work request %i: "
                        "automatic retry limit exceeded."
                    ),
                    running.id,
                )
                running.mark_aborted(
                    output_data=OutputData(
                        errors=[
                            OutputDataError(
                                message=(
                                    f"No worker managed to complete work "
                                    f"request {running.id}. Automatic retry "
                                    f"limit exceeded."
                                ),
                                code="work-request-retry-limit-exceeded",
                            )
                        ]
                    )
                )
                continue
            logger.info("Re-attempting work request %i", running.id)
            workflow_data.retry_count += 1
            running.workflow_data = workflow_data
            running.save()
            return running

        pending_work_requests = WorkRequest.objects.pending(worker=worker)

        # Aborting running work requests may change the first pending work
        # request that we select later, but it will still be in one of these
        # workflows.
        try:
            (
                running_work_requests | pending_work_requests[:1]
            ).lock_workflows_for_update(nowait=True)
        except OperationalError:
            logger.debug(
                "Failed to lock work requests assigned to %s", worker.name
            )
            raise DebusineAPIException(
                title="Retry",
                detail="Could not acquire workflow locks; retry later",
                # This is normally a WebDAV response code, but it seems
                # somewhat reasonable here too.
                status_code=status.HTTP_423_LOCKED,
            )

        # This view is normally only called when a worker is idle, but if
        # the worker was restarted then the server may still think it has a
        # running work request.  In that case, retry the work request since
        # it may be in an unclear intermediate state.
        for running in running_work_requests:
            logger.info("Retrying stuck work request %s", running.pk)
            try:
                # Disable permission checks, as retrying a work request is
                # normally not allowed with a worker token
                with context.disable_permission_checks():
                    running.retry(
                        reason=WorkRequestRetryReason.WORKER_FAILED,
                        force_status=True,
                    )
            except CannotRetry as e:
                # retry(force_status=True) starts by aborting the work
                # request; swallowing this exception and committing the
                # transaction anyway has the effect of leaving the work
                # request in the aborted status, which is what we want.
                logger.debug(  # noqa: G200
                    "Cannot retry previously-running work request: %s", e
                )

        return pending_work_requests.first()

    def post(self, request: Request) -> Response:  # noqa: ARG002
        """Return the task to build."""
        # This can't be None at this point, since IsWorkerTokenAuthenticated
        # passed.
        assert context.pc.worker_token
        worker = context.pc.worker_token.worker
        assert worker is not None

        # restricted to workers
        self.enforce(anonymous_access)

        work_request = self.select_work_request(worker)

        status_code: int
        if work_request:
            content = WorkerWorkRequestSerializer(work_request).data
            work_request.mark_running()
            status_code = status.HTTP_200_OK
        else:
            # There is no work request available for the worker
            content = None
            status_code = status.HTTP_204_NO_CONTENT

        return Response(content, status=status_code)


# mypy complains that celery.shared_task is untyped, which is true, but we
# can't fix that here.
@shared_task  # type: ignore[untyped-decorator]
@context.local()
@context.disable_permission_checks()
def update_work_request_as_completed(
    work_request_id: int,
    result: WorkRequest.Results,
    output_data: dict[str, Any] | None = None,
) -> None:
    """Mark a work request as completed."""
    found_work_request = False

    @atomic(isolation_level=pgtransaction.REPEATABLE_READ)
    def mark_completed() -> None:
        nonlocal found_work_request

        WorkRequest.objects.filter(
            id=work_request_id
        ).lock_workflows_for_update()

        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
        found_work_request = True

        work_request.mark_completed(
            result,
            output_data=(
                None
                if output_data is None
                else OutputData.model_validate(output_data)
            ),
        )

    @atomic(isolation_level=pgtransaction.REPEATABLE_READ)
    def retry_if_running() -> None:
        if tuple(
            WorkRequest.objects.filter(pk=work_request_id).values_list(
                "status", flat=True
            )
        ) == (WorkRequest.Statuses.RUNNING,):
            work_request = (
                WorkRequest.objects.filter(id=work_request_id)
                .lock_workflows_for_update()
                .get()
            )

            # Since the worker won't be informed of this Celery task
            # failing, we must take care to avoid it becoming stuck.  Retry
            # this so that the scheduler will give the worker something else
            # to do.
            work_request.retry(
                reason=WorkRequestRetryReason.WORKER_FAILED,
                force_status=True,
            )

    try:
        mark_completed()
    except Exception:
        if not found_work_request:
            # Re-raise any exceptions from before we got as far as trying to
            # mark the work request as completed.
            raise

        logger.exception(
            "Failed to mark work request %d as completed", work_request_id
        )
        retry_if_running()


class UpdateWorkRequestAsCompletedView(BaseAPIView):
    """
    View used by the workers to mark a task as completed.

    This defers most of the work to a Celery task, to avoid lock contention
    causing timeouts and retries, and potentially even workers having to
    give up and lose work.

    There is a slight risk in deferring this work.  If the worker is
    restarted between receiving the result from this view and the Celery
    task completing, then the server may send it another
    ``work_request_available`` message, causing the worker to request the
    next work request, causing the server to think that the worker failed
    and retry the work request.  However, given a successful response from
    this view, that scenario should only happen if the worker is restarted
    manually.
    """

    permission_classes = [IsWorkRequestTokenAuthenticated]

    def put(self, request: Request, work_request_id: int) -> Response:
        """Mark a work request as completed."""
        self.enforce(anonymous_access)  # restricted to workers

        try:
            work_request = WorkRequest.objects.can_display(context.pc).get(
                pk=work_request_id
            )
        except WorkRequest.DoesNotExist:
            raise DebusineAPIException(
                title="Work request not found",
                status_code=status.HTTP_404_NOT_FOUND,
            )

        self.check_object_permissions(request, work_request)

        work_request_completed_serializer = WorkRequestCompletedSerializer(
            data=request.data
        )

        if work_request_completed_serializer.is_valid():
            data = work_request_completed_serializer.validated_data
            apply_async_on_commit(
                update_work_request_as_completed,
                args=(work_request_id,),
                kwargs={
                    "result": data["result"],
                    "output_data": data.get("output_data"),
                },
            )
            content = None
            status_code = status.HTTP_200_OK
        else:
            raise DebusineAPIException(
                title="Cannot change work request as completed",
                validation_errors=work_request_completed_serializer.errors,
            )

        return Response(content, status=status_code)


class UpdateWorkerDynamicMetadataView(BaseAPIView):
    """View used by the workers to post dynamic metadata."""

    permission_classes = [IsWorkerTokenAuthenticated]

    def put(self, request: Request) -> Response:
        """Update Worker dynamic metadata."""
        self.enforce(anonymous_access)
        token_key = request.headers['token']
        worker = Worker.objects.get_worker_by_token_key_or_none(token_key)
        # Checked by IsWorkerTokenAuthenticated.
        assert worker is not None

        worker.set_dynamic_metadata(request.data)

        return Response(status=status.HTTP_204_NO_CONTENT)
