Ai/fertilization/views.py

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

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

from .serializers import FertilizationRecommendRequestSerializer


FertilizationQueueResponseSerializer = build_envelope_serializer(
    "FertilizationQueueResponseSerializer",
    build_task_queue_data_serializer("FertilizationQueueDataSerializer"),
)
FertilizationValidationErrorSerializer = build_envelope_serializer(
    "FertilizationValidationErrorSerializer",
    data_required=False,
    allow_null=True,
)
FertilizationStatusResponseSerializer = build_envelope_serializer(
    "FertilizationStatusResponseSerializer",
    build_task_status_data_serializer("FertilizationStatusDataSerializer"),
)


class FertilizationRecommendView(APIView):
    """
    توصیه کودهی با Celery.
    POST با sensor_uuid، plant_name، growth_stage.
    اطلاعات گیاه از plant app دریافت می‌شود.
    نیازی به دریافت نوع آبیاری نیست.
    """

    @extend_schema(
        tags=["Fertilization Recommendation"],
        summary="درخواست توصیه کودهی",
        description=(
            "داده‌های سنسور و گیاه را دریافت کرده و یک تسک Celery "
            "برای تولید توصیه کودهی در صف قرار می‌دهد. "
            "اطلاعات گیاه از جدول Plant بارگذاری می‌شود. "
            "محاسبات مربوط به نیاز آبی در این endpoint انجام نمی‌شود و مستقل از توصیه کودهی است."
        ),
        request=FertilizationRecommendRequestSerializer,
        responses={
            202: build_response(
                FertilizationQueueResponseSerializer,
                "تسک توصیه کودهی در صف قرار گرفت.",
            ),
            400: build_response(
                FertilizationValidationErrorSerializer,
                "پارامتر ورودی نامعتبر است.",
            ),
        },
        examples=[
            OpenApiExample(
                "نمونه درخواست",
                value={
                    "sensor_uuid": "550e8400-e29b-41d4-a716-446655440000",
                    "plant_name": "گوجه‌فرنگی",
                    "growth_stage": "گلدهی",
                },
                request_only=True,
            ),
        ],
    )
    def post(self, request):
        from rag.tasks import fertilization_recommendation_task

        serializer = FertilizationRecommendRequestSerializer(data=request.data)
        if not serializer.is_valid():
            return Response(
                {"code": 400, "msg": "داده نامعتبر.", "data": serializer.errors},
                status=status.HTTP_400_BAD_REQUEST,
            )

        validated = serializer.validated_data
        sensor_uuid = validated["sensor_uuid"]
        plant_name = validated.get("plant_name")
        growth_stage = validated.get("growth_stage")
        query = validated.get("query")

        task = fertilization_recommendation_task.delay(
            sensor_uuid=sensor_uuid,
            plant_name=plant_name,
            growth_stage=growth_stage,
            query=query,
        )
        return Response(
            {
                "code": 202,
                "msg": "تسک توصیه کودهی در صف قرار گرفت.",
                "data": {
                    "task_id": task.id,
                    "status_url": f"/api/fertilization/recommend/{task.id}/status/",
                },
            },
            status=status.HTTP_202_ACCEPTED,
        )


class FertilizationRecommendStatusView(APIView):
    """وضعیت تسک توصیه کودهی."""

    @extend_schema(
        tags=["Fertilization Recommendation"],
        summary="وضعیت تسک توصیه کودهی",
        description="وضعیت تسک Celery توصیه کودهی را برمی‌گرداند.",
        responses={
            200: build_response(
                FertilizationStatusResponseSerializer,
                "وضعیت فعلی تسک توصیه کودهی.",
            ),
        },
    )
    def get(self, request, task_id):
        from celery.result import AsyncResult

        result = AsyncResult(task_id)
        data = {"task_id": task_id, "status": result.state}
        if result.state == "PENDING":
            data["message"] = "تسک در صف یا یافت نشد."
        elif result.state == "PROGRESS":
            data["progress"] = result.info
        elif result.state == "SUCCESS":
            data["result"] = result.result
        elif result.state == "FAILURE":
            data["error"] = str(result.result)
        return Response(
            {"code": 200, "msg": "success", "data": data},
            status=status.HTTP_200_OK,
        )