Ai/rag/views.py

"""
ویوهای RAG — چت با استریم
"""
import logging

from django.http import StreamingHttpResponse
from drf_spectacular.types import OpenApiTypes
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.request import Request
from rest_framework.response import Response
from rest_framework.views import APIView

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


logger = logging.getLogger(__name__)


RagChatErrorResponseSerializer = build_message_response_serializer(
    "RagChatErrorResponseSerializer"
)
RagIrrigationQueueResponseSerializer = build_envelope_serializer(
    "RagIrrigationQueueResponseSerializer",
    build_task_queue_data_serializer("RagIrrigationQueueDataSerializer"),
)
RagIrrigationStatusResponseSerializer = build_envelope_serializer(
    "RagIrrigationStatusResponseSerializer",
    build_task_status_data_serializer("RagIrrigationStatusDataSerializer"),
)
RagFertilizationQueueResponseSerializer = build_envelope_serializer(
    "RagFertilizationQueueResponseSerializer",
    build_task_queue_data_serializer("RagFertilizationQueueDataSerializer"),
)
RagFertilizationStatusResponseSerializer = build_envelope_serializer(
    "RagFertilizationStatusResponseSerializer",
    build_task_status_data_serializer("RagFertilizationStatusDataSerializer"),
)
RagValidationErrorResponseSerializer = build_envelope_serializer(
    "RagValidationErrorResponseSerializer",
    data_required=False,
    allow_null=True,
)


class ChatView(APIView):
    """
    چت RAG با استریم.
    POST با {"query": "متن سوال", "farm_uuid": "شناسه مزرعه"}.
    همیشه از سرویس ثابت `chat` استفاده می‌کند و اطلاعات مزرعه را مستقیم به مدل می‌فرستد.
    """

    @extend_schema(
        tags=["RAG Chat"],
        summary="چت RAG با استریم",
        description="پیام کاربر را دریافت و پاسخ را به صورت استریم برمی‌گرداند.",
        request=inline_serializer(
            name="ChatRequest",
            fields={
                "query": drf_serializers.CharField(required=False, help_text="متن سوال کاربر"),
                "message": drf_serializers.CharField(required=False, help_text="نام قبلی فیلد query"),
                "farm_uuid": drf_serializers.CharField(help_text="شناسه مزرعه"),
            },
        ),
        responses={
            200: OpenApiResponse(
                response=OpenApiTypes.STR,
                description="پاسخ استریم متنی (text/plain)",
            ),
            400: build_response(
                RagChatErrorResponseSerializer,
                "پارامترهای ورودی نامعتبر هستند.",
            ),
            404: build_response(
                RagChatErrorResponseSerializer,
                "مزرعه پیدا نشد.",
            ),
        },
        examples=[
            OpenApiExample(
                "نمونه درخواست",
                value={
                    "farm_uuid": "550e8400-e29b-41d4-a716-446655440000",
                    "query": "وضعیت مزرعه من چطور است؟",
                },
                request_only=True,
            ),
        ],
    )
    def post(self, request: Request):
        from farm_data.services import get_farm_details
        from .config import load_rag_config

        data = request.data if request.method == "POST" else request.query_params
        message = data.get("query", data.get("message"))
        farm_uuid = data.get("farm_uuid")
        if not message or not isinstance(message, str):
            return Response(
                {"code": 400, "msg": "پارامتر query الزامی است."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        message = str(message).strip()
        if not message:
            return Response(
                {"code": 400, "msg": "پیام نباید خالی باشد."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        if not farm_uuid or not isinstance(farm_uuid, str):
            return Response(
                {"code": 400, "msg": "پارامتر farm_uuid الزامی است."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        farm_uuid = str(farm_uuid).strip()
        if not farm_uuid:
            return Response(
                {"code": 400, "msg": "farm_uuid نباید خالی باشد."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        cfg = load_rag_config()
        farm_details = get_farm_details(farm_uuid)
        if farm_details is None:
            return Response(
                {"code": 404, "msg": "farm پیدا نشد."},
                status=status.HTTP_404_NOT_FOUND,
            )

        def generate():
            try:
                for chunk in chat_rag_stream(
                    message,
                    farm_uuid=farm_uuid,
                    config=cfg,
                    farm_details=farm_details,
                ):
                    yield chunk
            except Exception as e:
                yield f"\n[خطا: {e}]"

        return StreamingHttpResponse(
            generate(),
            content_type="text/plain; charset=utf-8",
        )


class IrrigationRecommendationView(APIView):
    """
    توصیه آبیاری با Celery.
    POST با sensor_uuid، plant_name، growth_stage، irrigation_method_name.
    تسک در صف قرار می‌گیرد و task_id برگشت داده می‌شود.
    """

    @extend_schema(
        tags=["RAG Recommendations"],
        summary="درخواست توصیه آبیاری",
        description=(
            "داده‌های سنسور، گیاه و روش آبیاری را دریافت کرده و یک تسک Celery "
            "برای تولید توصیه آبیاری در صف قرار می‌دهد."
        ),
        request=inline_serializer(
            name="IrrigationRecommendationRequest",
            fields={
                "sensor_uuid": drf_serializers.CharField(help_text="شناسه یکتای سنسور (اجباری)"),
                "plant_name": drf_serializers.CharField(required=False, help_text="نام گیاه"),
                "growth_stage": drf_serializers.CharField(required=False, help_text="مرحله رشد گیاه"),
                "irrigation_method_name": drf_serializers.CharField(required=False, help_text="نام روش آبیاری"),
                "query": drf_serializers.CharField(required=False, help_text="سوال اختیاری"),
            },
        ),
        responses={
            202: build_response(
                RagIrrigationQueueResponseSerializer,
                "تسک توصیه آبیاری در صف قرار گرفت.",
            ),
            400: build_response(
                RagValidationErrorResponseSerializer,
                "پارامتر ورودی نامعتبر است.",
            ),
        },
        examples=[
            OpenApiExample(
                "نمونه درخواست",
                value={
                    "sensor_uuid": "550e8400-e29b-41d4-a716-446655440000",
                    "plant_name": "گوجه‌فرنگی",
                    "growth_stage": "میوه‌دهی",
                    "irrigation_method_name": "آبیاری قطره‌ای",
                },
                request_only=True,
            ),
        ],
    )
    def post(self, request: Request):
        from rag.tasks import irrigation_recommendation_task

        sensor_uuid = request.data.get("sensor_uuid")
        if not sensor_uuid:
            return Response(
                {"code": 400, "msg": "پارامتر sensor_uuid الزامی است.", "data": None},
                status=status.HTTP_400_BAD_REQUEST,
            )

        task = irrigation_recommendation_task.delay(
            sensor_uuid=str(sensor_uuid),
            plant_name=request.data.get("plant_name"),
            growth_stage=request.data.get("growth_stage"),
            irrigation_method_name=request.data.get("irrigation_method_name"),
            query=request.data.get("query"),
        )
        return Response(
            {
                "code": 202,
                "msg": "تسک توصیه آبیاری در صف قرار گرفت.",
                "data": {
                    "task_id": task.id,
                    "status_url": f"/api/rag/recommend/irrigation/{task.id}/status/",
                },
            },
            status=status.HTTP_202_ACCEPTED,
        )


class IrrigationRecommendationStatusView(APIView):
    """وضعیت تسک توصیه آبیاری."""

    @extend_schema(
        tags=["RAG Recommendations"],
        summary="وضعیت تسک توصیه آبیاری",
        description="وضعیت تسک Celery توصیه آبیاری را برمی‌گرداند.",
        responses={
            200: build_response(
                RagIrrigationStatusResponseSerializer,
                "وضعیت فعلی تسک توصیه آبیاری.",
            ),
        },
    )
    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,
        )


class FertilizationRecommendationView(APIView):
    """
    توصیه کودهی با Celery.
    POST با sensor_uuid، plant_name، growth_stage.
    تسک در صف قرار می‌گیرد و task_id برگشت داده می‌شود.
    """

    @extend_schema(
        tags=["RAG Recommendations"],
        summary="درخواست توصیه کودهی",
        description=(
            "داده‌های سنسور و گیاه را دریافت کرده و یک تسک Celery "
            "برای تولید توصیه کودهی در صف قرار می‌دهد."
        ),
        request=inline_serializer(
            name="FertilizationRecommendationRequest",
            fields={
                "sensor_uuid": drf_serializers.CharField(help_text="شناسه یکتای سنسور (اجباری)"),
                "plant_name": drf_serializers.CharField(required=False, help_text="نام گیاه"),
                "growth_stage": drf_serializers.CharField(required=False, help_text="مرحله رشد گیاه"),
                "query": drf_serializers.CharField(required=False, help_text="سوال اختیاری"),
            },
        ),
        responses={
            202: build_response(
                RagFertilizationQueueResponseSerializer,
                "تسک توصیه کودهی در صف قرار گرفت.",
            ),
            400: build_response(
                RagValidationErrorResponseSerializer,
                "پارامتر ورودی نامعتبر است.",
            ),
        },
        examples=[
            OpenApiExample(
                "نمونه درخواست",
                value={
                    "sensor_uuid": "550e8400-e29b-41d4-a716-446655440000",
                    "plant_name": "گوجه‌فرنگی",
                    "growth_stage": "رویشی",
                },
                request_only=True,
            ),
        ],
    )
    def post(self, request: Request):
        from rag.tasks import fertilization_recommendation_task

        sensor_uuid = request.data.get("sensor_uuid")
        if not sensor_uuid:
            return Response(
                {"code": 400, "msg": "پارامتر sensor_uuid الزامی است.", "data": None},
                status=status.HTTP_400_BAD_REQUEST,
            )

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


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

    @extend_schema(
        tags=["RAG Recommendations"],
        summary="وضعیت تسک توصیه کودهی",
        description="وضعیت تسک Celery توصیه کودهی را برمی‌گرداند.",
        responses={
            200: build_response(
                RagFertilizationStatusResponseSerializer,
                "وضعیت فعلی تسک توصیه کودهی.",
            ),
        },
    )
    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,
        )