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 با {"service_id": "...", "query": "متن سوال", "user_id": "شناسه کاربر"}
    service_id اجباری است. user_id فقط برای سرویس‌هایی که user embeddings دارند اجباری می‌شود.
    """

    @extend_schema(
        tags=["RAG Chat"],
        summary="چت RAG با استریم",
        description="پیام کاربر را دریافت و پاسخ را به صورت استریم برمی‌گرداند.",
        request=inline_serializer(
            name="ChatRequest",
            fields={
                "service_id": drf_serializers.CharField(help_text="شناسه سرویس"),
                "query": drf_serializers.CharField(required=False, help_text="متن سوال کاربر"),
                "message": drf_serializers.CharField(required=False, help_text="نام قبلی فیلد query"),
                "user_id": drf_serializers.CharField(required=False, help_text="شناسه کاربر"),
                "sensor_uuid": drf_serializers.CharField(required=False, help_text="نام قبلی فیلد user_id"),
            },
        ),
        responses={
            200: OpenApiResponse(
                response=OpenApiTypes.STR,
                description="پاسخ استریم متنی (text/plain)",
            ),
            400: build_response(
                RagChatErrorResponseSerializer,
                "پارامترهای ورودی نامعتبر هستند.",
            ),
        },
        examples=[
            OpenApiExample(
                "نمونه درخواست",
                value={
                    "service_id": "support_bot",
                    "user_id": "12345",
                    "query": "How do I reset my password?",
                },
                request_only=True,
            ),
        ],
    )
    def post(self, request: Request):
        from .config import load_rag_config, get_service_config

        data = request.data if request.method == "POST" else request.query_params
        service_id = data.get("service_id")
        message = data.get("query", data.get("message"))
        user_id = data.get("user_id", data.get("sensor_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 service_id or not isinstance(service_id, str):
            return Response(
                {"code": 400, "msg": "پارامتر service_id الزامی است."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        service_id = str(service_id).strip()
        if not service_id:
            return Response(
                {"code": 400, "msg": "service_id نباید خالی باشد."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        cfg = load_rag_config()
        try:
            service = get_service_config(service_id, cfg)
        except KeyError:
            return Response(
                {"code": 400, "msg": f"service_id نامعتبر است: {service_id}"},
                status=status.HTTP_400_BAD_REQUEST,
            )
        if user_id is not None:
            user_id = str(user_id).strip()
            if not user_id:
                user_id = None
        if service.use_user_embeddings and not user_id:
            return Response(
                {"code": 400, "msg": "برای این service_id، پارامتر user_id الزامی است."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        def generate():
            try:
                for chunk in chat_rag_stream(
                    message,
                    sensor_uuid=user_id,
                    service_id=service_id,
                    config=cfg,
                ):
                    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,
        )