Ai/rag/services/water_need_prediction.py

from __future__ import annotations

import json
import logging
from typing import Any

from farm_data.services import get_farm_details
from rag.api_provider import get_chat_client
from rag.chat import (
    _complete_audit_log,
    _create_audit_log,
    _fail_audit_log,
    _load_service_tone,
    build_rag_context,
)
from rag.config import RAGConfig, get_service_config, load_rag_config

logger = logging.getLogger(__name__)

KB_NAME = "water_need_prediction"
SERVICE_ID = "water_need_prediction"

WATER_NEED_PROMPT = (
    "شما یک دستیار تخصصی تحليل نياز آبي کوتاه مدت مزرعه هستي. "
    "ورودي شامل محاسبات ساختاريافته نياز آبي، اطلاعات مزرعه و متن هاي بازيابي شده از پايگاه دانش است. "
    "فقط JSON معتبر با اين کليدها برگردان: "
    "summary, irrigation_outlook, recommended_action, risk_note, confidence. "
    "confidence عددي بين 0 و 1 باشد. "
    "اعداد اصلي را از داده ورودي بگير و عدد متناقض جديد نساز."
)


def _clean_json(raw: str) -> dict[str, Any]:
    cleaned = (raw or "").strip()
    if cleaned.startswith("```"):
        cleaned = cleaned.strip("`")
        if cleaned.startswith("json"):
            cleaned = cleaned[4:]
        cleaned = cleaned.strip()
    if not cleaned:
        return {}
    try:
        return json.loads(cleaned)
    except (json.JSONDecodeError, ValueError):
        logger.warning("Invalid JSON returned by water_need_prediction LLM: %s", cleaned[:500])
        return {}


def _load_farm_or_error(farm_uuid: str) -> dict[str, Any]:
    farm_details = get_farm_details(farm_uuid)
    if farm_details is None:
        raise ValueError("farm_uuid نامعتبر است یا اطلاعات مزرعه پیدا نشد.")
    return farm_details


def _build_service_client(cfg: RAGConfig):
    service = get_service_config(SERVICE_ID, cfg)
    service_cfg = RAGConfig(
        embedding=cfg.embedding,
        qdrant=cfg.qdrant,
        chunking=cfg.chunking,
        llm=service.llm,
        knowledge_bases=cfg.knowledge_bases,
        services=cfg.services,
        chromadb=cfg.chromadb,
    )
    client = get_chat_client(service_cfg)
    return service, client, service.llm.model


def _fallback_from_payload(prediction_payload: dict[str, Any]) -> dict[str, Any]:
    total = float(prediction_payload.get("totalNext7Days") or 0.0)
    daily = prediction_payload.get("dailyBreakdown") or []
    peak_day = max(daily, key=lambda item: float(item.get("gross_irrigation_mm", 0.0) or 0.0), default=None)
    if total <= 0:
        return {
            "summary": "براي چند روز آينده نياز آبي معناداري برآورد نشد.",
            "irrigation_outlook": "بارش موثر يا شرايط فعلي باعث شده نياز خالص آبياري پايين باشد.",
            "recommended_action": "پايش رطوبت خاک ادامه يابد و قبل از هر آبياري جديد forecast دوباره بررسي شود.",
            "risk_note": "اگر forecast تغيير کند يا بارش موثر رخ ندهد، برآورد بايد به روز شود.",
            "confidence": 0.58,
        }

    peak_text = ""
    if peak_day:
        peak_text = (
            f" بيشترين فشار آبي در {peak_day.get('forecast_date')} "
            f"با حدود {peak_day.get('gross_irrigation_mm')} ميلي متر برآورد شده است."
        )
    return {
        "summary": f"جمع نياز آبي 7 روز آينده حدود {round(total, 2)} ميلي متر برآورد شده است.",
        "irrigation_outlook": "الگوي آبياري بايد در چند روز آينده بر اساس نياز روزانه و بارش موثر تنظيم شود." + peak_text,
        "recommended_action": "برنامه آبياري کوتاه مدت بر اساس روزهاي اوج نياز تنظيم و صبح زود يا نزديک غروب اجرا شود.",
        "risk_note": "در صورت تغيير دما، باد يا بارش، مقادير gross irrigation ممکن است تغيير کنند.",
        "confidence": 0.72,
    }


def _build_messages(
    *,
    service: Any,
    cfg: RAGConfig,
    query: str,
    rag_context: str,
    structured_context: dict[str, Any],
) -> tuple[str, list[dict[str, str]]]:
    tone = _load_service_tone(service, cfg)
    system_parts = [tone] if tone else []
    if service.system_prompt:
        system_parts.append(service.system_prompt)
    system_parts.append(WATER_NEED_PROMPT)
    system_parts.append(
        "[کانتکست ساختاريافته نياز آبي]\n"
        + json.dumps(structured_context, ensure_ascii=False, indent=2, default=str)
    )
    if rag_context:
        system_parts.append(rag_context)
    system_prompt = "\n\n".join(part for part in system_parts if part)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": query},
    ]
    return system_prompt, messages


def get_water_need_prediction_insight(
    *,
    farm_uuid: str,
    prediction_payload: dict[str, Any],
    query: str | None = None,
) -> dict[str, Any]:
    cfg = load_rag_config()
    service, client, model = _build_service_client(cfg)
    farm_details = _load_farm_or_error(farm_uuid)
    fallback = _fallback_from_payload(prediction_payload)
    user_query = query or "نياز آبي کوتاه مدت اين مزرعه را تفسير کن و اقدام عملياتي پيشنهاد بده."
    structured_context = {
        "farm_uuid": farm_uuid,
        "prediction_payload": prediction_payload,
        "fallback_summary": fallback,
    }
    rag_context = build_rag_context(
        query=user_query,
        sensor_uuid=farm_uuid,
        config=cfg,
        kb_name=KB_NAME,
        service_id=SERVICE_ID,
        farm_details=farm_details,
    )
    system_prompt, messages = _build_messages(
        service=service,
        cfg=cfg,
        query=user_query,
        rag_context=rag_context,
        structured_context=structured_context,
    )
    audit_log = _create_audit_log(
        farm_uuid=farm_uuid,
        service_id=SERVICE_ID,
        model=model,
        query=user_query,
        system_prompt=system_prompt,
        messages=messages,
    )
    try:
        response = client.chat.completions.create(model=model, messages=messages)
        raw = response.choices[0].message.content.strip()
        parsed = _clean_json(raw)
        _complete_audit_log(audit_log, raw)
    except Exception as exc:
        logger.error("Water need prediction insight failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc), json.dumps(fallback, ensure_ascii=False))
        return {
            **fallback,
            "farm_uuid": farm_uuid,
            "knowledge_base": KB_NAME,
            "tone_file": service.tone_file,
            "raw_response": None,
        }

    if not parsed:
        parsed = fallback
    parsed.setdefault("summary", fallback["summary"])
    parsed.setdefault("irrigation_outlook", fallback["irrigation_outlook"])
    parsed.setdefault("recommended_action", fallback["recommended_action"])
    parsed.setdefault("risk_note", fallback["risk_note"])
    parsed.setdefault("confidence", fallback["confidence"])
    parsed["farm_uuid"] = farm_uuid
    parsed["knowledge_base"] = KB_NAME
    parsed["tone_file"] = service.tone_file
    parsed["raw_response"] = raw
    return parsed