Ai/rag/services/pest_disease.py

"""
سرویس RAG برای تشخیص تصویری و پیش بینی ریسک آفات و بیماری گیاه.
"""
from __future__ import annotations

import json
import logging
from typing import Any

from farm_data.services import build_ai_farm_snapshot
from rag.api_provider import get_chat_client
from rag.chat import (
    _build_content_parts,
    _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
from rag.failure_contract import RAGServiceError
from rag.user_data import build_plant_text

logger = logging.getLogger(__name__)

KB_NAME = "pest_disease"
SERVICE_ID = "pest_disease"

DETECTION_PROMPT = (
    "شما یک دستیار تخصصی تشخیص آفات و بیماری گیاهی هستی. "
    "با استفاده از تصویر، اطلاعات مزرعه، و متن های بازیابی شده از پایگاه دانش تحلیل کن. "
    "پاسخ فقط JSON معتبر باشد و این کلیدها را داشته باشد: "
    "has_issue, category, confidence, severity, summary, detected_signs, possible_causes, immediate_actions, reasoning. "
    "category فقط یکی از no_issue, pest, disease, nutrient_stress, abiotic_stress, unknown باشد. "
    "severity فقط یکی از low, medium, high باشد."
)

RISK_PROMPT = (
    "شما یک دستیار تخصصی پیش بینی ریسک آفات و بیماری گیاهی هستی. "
    "با استفاده از داده های مزرعه، آب و هوا، مرحله رشد، و متن های بازیابی شده از پایگاه دانش تحلیل کن. "
    "پاسخ فقط JSON معتبر باشد و این کلیدها را داشته باشد: "
    "summary, forecast_window, overall_risk, disease_risk, pest_risk, key_drivers, recommended_actions. "
    "overall_risk فقط یکی از low, medium, high باشد. "
    "disease_risk و pest_risk باید آبجکت هایی با کلیدهای score, level, likely_conditions, reasoning باشند و level فقط یکی از low, medium, high باشد."
)


def _safe_float(value: Any, default: float = 0.0) -> float:
    try:
        if value in (None, ""):
            return default
        return float(value)
    except (TypeError, ValueError):
        return default


def _normalize_images(images: list[dict[str, str]] | None) -> list[dict[str, str]]:
    output: list[dict[str, str]] = []
    for item in images or []:
        if not isinstance(item, dict):
            continue
        url = item.get("url")
        if not isinstance(url, str) or not url.strip():
            continue
        output.append({"url": url.strip(), "detail": item.get("detail", "auto")})
    return output


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:
        raise RAGServiceError(
            error_code="empty_response",
            message="Pest disease LLM response was empty.",
            source="llm",
            retriable=True,
            details={"service_id": SERVICE_ID},
            http_status=502,
        )
    try:
        parsed = json.loads(cleaned)
    except (json.JSONDecodeError, ValueError) as exc:
        logger.warning("Invalid JSON returned by pest_disease LLM: %s", cleaned[:500])
        raise RAGServiceError(
            error_code="invalid_json",
            message="Pest disease LLM response was not valid JSON.",
            source="llm",
            retriable=True,
            details={"service_id": SERVICE_ID},
            http_status=502,
        ) from exc
    if not isinstance(parsed, dict):
        raise RAGServiceError(
            error_code="invalid_schema",
            message="Pest disease LLM response root must be a JSON object.",
            source="llm",
            details={"service_id": SERVICE_ID},
            http_status=502,
        )
    return parsed


def _load_farm_or_error(farm_uuid: str) -> dict[str, Any]:
    farm_details = build_ai_farm_snapshot(farm_uuid)
    if farm_details is None:
        raise RAGServiceError(
            error_code="farm_not_found",
            message="farm_uuid نامعتبر است یا اطلاعات مزرعه پیدا نشد.",
            source="farm_data",
            details={"farm_uuid": farm_uuid},
            http_status=404,
        )
    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 _weather_risk_summary(farm_details: dict[str, Any]) -> dict[str, Any]:
    weather = ((farm_details.get("weather") or {}).get("forecast") or {})
    soil = (farm_details.get("farm_metrics") or {}).get("resolved_metrics") or {}
    humidity = _safe_float(weather.get("humidity_mean"), 55.0)
    temp = _safe_float(weather.get("temperature_mean"), 24.0)
    rain = _safe_float(weather.get("precipitation"), 0.0)
    moisture = _safe_float(soil.get("soil_moisture"), _safe_float(soil.get("wv0033"), 35.0))
    ec = _safe_float(soil.get("electrical_conductivity"), 0.0)
    ph = _safe_float(soil.get("soil_ph") or soil.get("phh2o"), 7.0)

    fungal_score = min(max(round((humidity * 0.45) + (moisture * 0.35) + (rain * 2.5) - 25, 2), 0.0), 100.0)
    pest_score = min(max(round((temp * 2.2) + max(0.0, 45.0 - moisture) + (ec * 3.0) - 20, 2), 0.0), 100.0)
    abiotic_stress = min(max(round((abs(ph - 6.8) * 18.0) + (ec * 8.0), 2), 0.0), 100.0)
    return {
        "humidity_mean": humidity,
        "temperature_mean": temp,
        "precipitation": rain,
        "soil_moisture": moisture,
        "ec": ec,
        "ph": ph,
        "fungal_score": fungal_score,
        "pest_score": pest_score,
        "abiotic_stress_score": abiotic_stress,
    }


def _risk_level(score: float) -> str:
    if score >= 70:
        return "high"
    if score >= 40:
        return "medium"
    return "low"


def _build_risk_context(farm_details: dict[str, Any], plant_name: str | None, growth_stage: str | None) -> dict[str, Any]:
    risk = _weather_risk_summary(farm_details)
    disease_level = _risk_level(risk["fungal_score"])
    pest_level = _risk_level(risk["pest_score"])
    overall_score = max(risk["fungal_score"], risk["pest_score"], risk["abiotic_stress_score"])
    overall_level = _risk_level(overall_score)
    drivers = []
    if risk["humidity_mean"] >= 70:
        drivers.append("رطوبت بالا")
    if risk["soil_moisture"] >= 60:
        drivers.append("رطوبت خاک بالا")
    if risk["temperature_mean"] >= 30:
        drivers.append("دمای بالا")
    if risk["precipitation"] > 2:
        drivers.append("بارش موثر")
    if risk["ec"] > 2.5:
        drivers.append("EC بالا")
    if abs(risk["ph"] - 6.8) > 0.8:
        drivers.append("خروج pH از محدوده مطلوب")
    if not drivers:
        drivers.append("شرایط فعلی مزرعه نسبتا پایدار است")

    return {
        "summary": "برآورد ریسک آفات و بیماری بر اساس داده های فعلی مزرعه ساخته شد.",
        "forecast_window": "24 تا 72 ساعت آینده",
        "overall_risk": overall_level,
        "disease_risk": {
            "score": risk["fungal_score"],
            "level": disease_level,
            "likely_conditions": [
                "فشار قارچی و بیماری برگی" if disease_level != "low" else "ریسک بیماری فعلا پایین است",
            ],
            "reasoning": [
                f"رطوبت میانگین حدود {risk['humidity_mean']} درصد است.",
                f"رطوبت خاک حدود {risk['soil_moisture']} درصد برآورد شده است.",
            ],
        },
        "pest_risk": {
            "score": risk["pest_score"],
            "level": pest_level,
            "likely_conditions": [
                "فشار آفات مکنده یا تنش زا" if pest_level != "low" else "ریسک آفت فعلا پایین است",
            ],
            "reasoning": [
                f"دمای میانگین حدود {risk['temperature_mean']} درجه است.",
                f"EC فعلی حدود {risk['ec']} و pH حدود {risk['ph']} است.",
            ],
        },
        "key_drivers": drivers,
        "recommended_actions": [
            "بازدید مزرعه و بررسی برگ ها و پشت برگ انجام شود.",
            "در صورت مشاهده علائم مشکوک، نمونه برداری تصویری نزدیک تر انجام شود.",
            "رطوبت ماندگار و یکنواختی آبیاری پایش شود.",
        ],
        "farm_context": {
            "plant_name": plant_name,
            "growth_stage": growth_stage,
            "risk_summary": risk,
        },
    }


def _validate_detection_result(parsed: dict[str, Any]) -> dict[str, Any]:
    required_keys = {
        "has_issue",
        "category",
        "confidence",
        "severity",
        "summary",
        "detected_signs",
        "possible_causes",
        "immediate_actions",
        "reasoning",
    }
    missing = [key for key in required_keys if key not in parsed]
    if missing:
        raise RAGServiceError(
            error_code="invalid_schema",
            message="Pest disease detection response is missing required fields: " + ", ".join(missing),
            source="llm",
            details={"missing_fields": missing, "service_id": SERVICE_ID},
            http_status=502,
        )
    return parsed


def _validate_risk_result(parsed: dict[str, Any]) -> dict[str, Any]:
    required_keys = {
        "summary",
        "forecast_window",
        "overall_risk",
        "disease_risk",
        "pest_risk",
        "key_drivers",
        "recommended_actions",
    }
    missing = [key for key in required_keys if key not in parsed]
    if missing:
        raise RAGServiceError(
            error_code="invalid_schema",
            message="Pest disease risk response is missing required fields: " + ", ".join(missing),
            source="llm",
            details={"missing_fields": missing, "service_id": SERVICE_ID},
            http_status=502,
        )
    return parsed


def _build_detection_messages(
    *,
    service: Any,
    cfg: RAGConfig,
    query: str,
    rag_context: str,
    plant_text: str,
    images: list[dict[str, str]],
) -> tuple[str, list[dict[str, Any]]]:
    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(DETECTION_PROMPT)
    if plant_text:
        system_parts.append("[اطلاعات گیاه]\n" + plant_text)
    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": _build_content_parts(query, images)},
    ]
    return system_prompt, messages


def _build_risk_messages(
    *,
    service: Any,
    cfg: RAGConfig,
    query: str,
    rag_context: str,
    structured_context: dict[str, Any],
    plant_text: str,
) -> 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(RISK_PROMPT)
    if plant_text:
        system_parts.append("[اطلاعات گیاه]\n" + plant_text)
    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_pest_disease_detection(
    *,
    farm_uuid: str,
    plant_name: str | None = None,
    query: str | None = None,
    images: list[dict[str, str]] | None = None,
) -> dict[str, Any]:
    normalized_images = _normalize_images(images)
    if not normalized_images:
        raise RAGServiceError(
            error_code="missing_images",
            message="حداقل یک تصویر برای تشخیص لازم است.",
            source="request",
            http_status=400,
        )

    cfg = load_rag_config()
    service, client, model = _build_service_client(cfg)
    farm_details = _load_farm_or_error(farm_uuid)
    resolved_plant_name = plant_name or (farm_details.get("plants") or [{}])[0].get("name")
    user_query = query or "این تصویر را بررسی کن و بگو آیا گیاه دچار آفت یا بیماری شده است یا نه."
    plant_text = build_plant_text(resolved_plant_name, "") if resolved_plant_name else ""
    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_detection_messages(
        service=service,
        cfg=cfg,
        query=user_query,
        rag_context=rag_context,
        plant_text=plant_text or "",
        images=normalized_images,
    )
    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 RAGServiceError as exc:
        logger.error("Pest disease detection failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise
    except Exception as exc:
        logger.error("Pest disease detection failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise RAGServiceError(
            error_code="upstream_failure",
            message=f"Pest disease detection failed for farm {farm_uuid}.",
            source="llm",
            retriable=True,
            details={"farm_uuid": farm_uuid, "service_id": SERVICE_ID},
            http_status=503,
        ) from exc

    parsed = _validate_detection_result(parsed)
    parsed["status"] = "success"
    parsed["source"] = "llm"
    parsed["farm_uuid"] = farm_uuid
    parsed["raw_response"] = raw
    return parsed


def get_pest_disease_risk(
    *,
    farm_uuid: str,
    plant_name: str | None = None,
    growth_stage: str | None = None,
    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)
    resolved_plant_name = plant_name or (farm_details.get("plants") or [{}])[0].get("name")
    risk_context = _build_risk_context(farm_details, resolved_plant_name, growth_stage)
    user_query = query or "ریسک آفات و بیماری این مزرعه را برای چند روز آینده پیش بینی کن."
    plant_text = build_plant_text(resolved_plant_name, growth_stage or "") if resolved_plant_name else ""
    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_risk_messages(
        service=service,
        cfg=cfg,
        query=user_query,
        rag_context=rag_context,
        structured_context=risk_context,
        plant_text=plant_text or "",
    )
    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 RAGServiceError as exc:
        logger.error("Pest disease risk prediction failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise
    except Exception as exc:
        logger.error("Pest disease risk prediction failed for %s: %s", farm_uuid, exc)
        _fail_audit_log(audit_log, str(exc))
        raise RAGServiceError(
            error_code="upstream_failure",
            message=f"Pest disease risk prediction failed for farm {farm_uuid}.",
            source="llm",
            retriable=True,
            details={"farm_uuid": farm_uuid, "service_id": SERVICE_ID},
            http_status=503,
        ) from exc

    parsed = _validate_risk_result(parsed)
    parsed["status"] = "success"
    parsed["source"] = "llm"
    parsed["farm_uuid"] = farm_uuid
    parsed["raw_response"] = raw
    return parsed