Ai/tasks/views.py

from celery.result import AsyncResult

from drf_spectacular.utils import (
    OpenApiExample,
    OpenApiResponse,
    extend_schema,
    inline_serializer,
)
from rest_framework import status
from rest_framework import serializers as drf_serializers
from rest_framework.response import Response
from rest_framework.views import APIView

from config.openapi import (
    build_envelope_serializer,
    build_response,
    build_task_status_data_serializer,
)

from .celery_tasks import sample_task


TaskTriggerResponseSerializer = build_envelope_serializer(
    "TaskTriggerResponseSerializer",
    inline_serializer(
        name="TaskTriggerPayloadSerializer",
        fields={
            "task_id": drf_serializers.CharField(),
        },
    ),
)
TaskStatusResponseSerializer = build_envelope_serializer(
    "TaskStatusEnvelopeSerializer",
    build_task_status_data_serializer("TaskStatusPayloadSerializer"),
)


class TaskTriggerView(APIView):
    """
    ثبت و اجرای تسک.
    POST با بدنه اختیاری: {"duration": 3} - مدت زمان تسک به ثانیه.
    """

    @extend_schema(
        tags=["Tasks"],
        summary="ثبت و اجرای تسک",
        description="یک تسک نمونه را در صف Celery قرار می‌دهد.",
        request=inline_serializer(
            name="TaskTriggerRequest",
            fields={
                "duration": drf_serializers.IntegerField(
                    required=False, default=1, help_text="مدت زمان تسک به ثانیه (۱–۶۰)"
                ),
            },
        ),
        responses={
            200: build_response(
                TaskTriggerResponseSerializer,
                "تسک نمونه با موفقیت در صف قرار گرفت.",
            ),
        },
        examples=[
            OpenApiExample(
                "نمونه درخواست",
                value={"duration": 3},
                request_only=True,
            ),
            OpenApiExample(
                "نمونه پاسخ",
                value={"code": 200, "msg": "success", "data": {"task_id": "abc123-def456"}},
                response_only=True,
            ),
        ],
    )
    def post(self, request):
        duration = request.data.get("duration", 1)
        try:
            duration = int(duration)
            duration = max(1, min(duration, 60))
        except (TypeError, ValueError):
            duration = 1
        result = sample_task.delay(duration)
        return Response(
            {"code": 200, "msg": "success", "data": {"task_id": result.id}},
            status=status.HTTP_200_OK,
        )


class TaskStatusView(APIView):
    """
    وضعیت تسک بر اساس task_id.
    GET /api/tasks/<task_id>/status/
    """

    @extend_schema(
        tags=["Tasks"],
        summary="وضعیت تسک",
        description="وضعیت یک تسک Celery را بر اساس task_id برمی‌گرداند.",
        responses={
            200: build_response(
                TaskStatusResponseSerializer,
                "وضعیت فعلی تسک Celery.",
            ),
        },
        examples=[
            OpenApiExample(
                "تسک موفق",
                value={"code": 200, "msg": "success", "data": {"task_id": "abc123", "status": "SUCCESS", "result": "done"}},
                response_only=True,
            ),
        ],
    )
    def get(self, request, task_id):
        result = AsyncResult(task_id)
        state = result.state
        data = {"task_id": task_id, "status": state}
        if state == "PENDING":
            data["message"] = "تسک در صف یا یافت نشد."
        elif state == "PROGRESS":
            data["progress"] = result.info
        elif state == "SUCCESS":
            data["result"] = result.result
        elif state == "FAILURE":
            data["error"] = str(result.result)
        return Response(
            {"code": 200, "msg": "success", "data": data},
            status=status.HTTP_200_OK,
        )