src/utils/farm_classifier.py

"""
Farm Analyzer with Nebius AI Studio
====================================

Analyzes farm/crop images using Qwen2.5-VL-72B-Instruct via Nebius API.
Provides overall health assessment, disease detection, and recommendations.

Model: Qwen/Qwen2.5-VL-72B-Instruct
API: OpenAI-compatible

Usage:
    from src.nebius_analyzer import analyze_farm_image
    
    result = analyze_farm_image(pil_image)
    print(result["health_status"])  # "Diseased"
"""

import os
import json
import base64
from io import BytesIO
from PIL import Image


# ============================================================
# CONFIGURATION
# ============================================================

NEBIUS_API_URL = "https://api.studio.nebius.com/v1/"
MODEL_NAME = "Qwen/Qwen2.5-VL-72B-Instruct"


# ============================================================
# MODULE STATE
# ============================================================

_client = None


# ============================================================
# PRIVATE FUNCTIONS
# ============================================================

def _get_api_key() -> str:
    """Get Nebius API key from environment variable."""
    return os.environ.get("NEBIUS_API_KEY", "")


def _get_client():
    """
    Get OpenAI client configured for Nebius.
    
    Returns:
        OpenAI client or None if API key not set
    """
    global _client
    
    if _client is not None:
        return _client
    
    api_key = _get_api_key()
    if not api_key:
        print("⚠️ NEBIUS_API_KEY not set")
        return None
    
    try:
        from openai import OpenAI
        
        print("🤖 Initializing Nebius client...")
        _client = OpenAI(
            base_url=NEBIUS_API_URL,
            api_key=api_key,
        )
        print(f"✅ Nebius client ready (model: {MODEL_NAME})")
        
        return _client
    
    except ImportError:
        print("❌ openai package not installed. Run: pip install openai")
        return None
    except Exception as e:
        print(f"❌ Error initializing Nebius client: {e}")
        return None


def _encode_image_to_base64(image: Image.Image) -> str:
    """
    Encode PIL Image to base64 string.
    
    Args:
        image: PIL Image
    
    Returns:
        Base64 encoded string
    """
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")


def _build_prompt() -> str:
    """Build the analysis prompt for the vision model."""
    
    return """You are an expert agronomist analyzing a farm/crop image.

        Analyze this image and provide a detailed assessment of the crop health.

        Respond ONLY with valid JSON (no markdown, no code blocks, no extra text):
        {
            "crop_identified": "name of the crop or 'Unknown'",
            "health_status": "Healthy" or "Diseased" or "Stressed" or "Unknown",
            "confidence": 0-100,
            "issues_detected": [
                {
                    "name": "issue name",
                    "severity": "Low" or "Medium" or "High" or "Critical",
                    "affected_area": "percentage or description"
                }
            ],
            "overall_severity": "None" or "Low" or "Medium" or "High" or "Critical",
            "recommendations": [
                "recommendation 1",
                "recommendation 2",
                "recommendation 3"
            ],
            "observations": "any additional observations about the crop/field"
        }

        Important:
        - If you cannot identify the crop, set crop_identified to "Unknown"
        - Always provide at least one recommendation
        - Be specific about issues you observe
        - Output ONLY the JSON, nothing else"""


def _parse_response(text: str) -> dict:
    """Parse model response to dict."""
    
    # Clean markdown if present
    cleaned = text.strip()
    
    if "```json" in cleaned:
        cleaned = cleaned.split("```json")[1].split("```")[0]
    elif "```" in cleaned:
        cleaned = cleaned.split("```")[1].split("```")[0]
    
    cleaned = cleaned.strip()
    
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError as e:
        # Try to extract JSON with regex as fallback
        import re
        match = re.search(r'\{[\s\S]*\}', cleaned)
        if match:
            try:
                return json.loads(match.group())
            except:
                pass
        
        return {
            "parse_error": str(e),
            "raw_response": text
        }


# ============================================================
# PUBLIC FUNCTION
# ============================================================

def analyze_farm_image(image: Image.Image) -> dict:
    """
    Analyze a farm/crop image using Qwen2.5-VL-72B via Nebius.
    
    Args:
        image: PIL Image of farm/crop
    
    Returns:
        dict with result:
        {
            "success": True,
            "crop_identified": "Tomato",
            "health_status": "Diseased",
            "confidence": 85,
            "issues_detected": [...],
            "overall_severity": "Medium",
            "recommendations": [...],
            "observations": "..."
        }
        
        On error:
        {
            "success": False,
            "error": "Error description"
        }
    """
    
    # Validate input
    if image is None:
        return {
            "success": False,
            "error": "No image provided"
        }
    
    if not isinstance(image, Image.Image):
        return {
            "success": False,
            "error": f"Invalid image type: {type(image)}. Expected PIL.Image"
        }
    
    # Check API key
    if not _get_api_key():
        return {
            "success": False,
            "error": "NEBIUS_API_KEY not configured. Set it as environment variable."
        }
    
    # Get client
    client = _get_client()
    if client is None:
        return {
            "success": False,
            "error": "Failed to initialize Nebius client"
        }
    
    try:
        # Convert to RGB if needed
        image = image.convert("RGB")
        
        # Encode image to base64
        image_base64 = _encode_image_to_base64(image)
        
        # Build prompt
        prompt = _build_prompt()
        
        # Call Nebius API with vision model
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}"
                            }
                        },
                        {
                            "type": "text",
                            "text": prompt
                        }
                    ]
                }
            ],
            max_tokens=1024,
            temperature=0.3  # Lower temperature for more consistent JSON output
        )
        
        # Extract response text
        response_text = response.choices[0].message.content
        
        # Parse response
        result = _parse_response(response_text)
        
        # Check for parse error
        if "parse_error" in result:
            return {
                "success": False,
                "error": f"Failed to parse response: {result['parse_error']}",
                "raw_response": result.get("raw_response", "")
            }
        
        # Add success flag
        result["success"] = True
        return result
    
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }