model.py

import os
import sys
import platform
import argparse
import json
from dotenv import load_dotenv
from groq import Groq

# Import the new prompts
from prompts import SYSTEM_PROMPT

# Load environment variables
load_dotenv()

# Common configuration
OLLAMA_MODEL = "llama-3.3-70b-versatile"  # Upgraded model for better Tamil accuracy

# Initialize Groq client
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
client = None
if GROQ_API_KEY:
    client = Groq(api_key=GROQ_API_KEY)

def print_header() -> None:
    print("AgroGPT (Mobile Backend) starting...", flush=True)
    print(f"Python: {platform.python_version()} ({sys.executable})", flush=True)

def check_ollama_connection() -> bool:
    """Check if Groq API is reachable and key is valid."""
    if not GROQ_API_KEY:
        print("Error: GROQ_API_KEY not found in .env file.", flush=True)
        return False
    
    try:
        if client:
            client.models.list()
            print("Connected to Groq (JSON Mode Ready).", flush=True)
            return True
        return False
    except Exception as e:
        print(f"Error: Could not connect to backend: {str(e)}", flush=True)
        return False

def generate_with_ollama(user_prompt: str, system_prompt: str = SYSTEM_PROMPT, model: str = OLLAMA_MODEL) -> dict:
    """
    Generate a JSON response using Groq.
    Returns a dictionary parsed from the JSON response.
    """
    if not client:
        return {"error": "Groq client not initialized. Check your API key."}
        
    try:
        completion = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            response_format={"type": "json_object"},
            stream=False
        )
        response_text = completion.choices[0].message.content
        return json.loads(response_text)
    except json.JSONDecodeError:
        return {"error": "Failed to parse JSON response from LLM", "raw_response": response_text}
    except Exception as e:
        return {"error": f"Error generating response: {str(e)}"}

# Import the disease detection module
try:
    from disease_detection import get_disease_detector
    HAS_DISEASE_DETECTION = True
except ImportError:
    HAS_DISEASE_DETECTION = False
    print("Warning: disease_detection module not found. Vision features disabled.")

def analyze_image_for_disease(image_path: str) -> dict:
    """
    Analyzes a plant image using the local vision model, then generates
    expert advice using Groq (returning JSON).
    """
    if not HAS_DISEASE_DETECTION:
        return {"error": "Disease detection module not available."}

    try:
        detector = get_disease_detector()
        result = detector.predict_disease(image_path)
        
        if "error" in result:
             return {"error": result.get('error')}

        disease_name = result.get('prediction', 'Unknown')
        confidence = result.get('confidence', 0.0)
        is_simulated = result.get('simulation', False)
        
        # Construct prompt for the LLM to get structured advice
        # We reuse the same system prompt structure but adapt the user input
        input_data = {
            "task": "disease_analysis",
            "disease_name": disease_name,
            "confidence_score": confidence,
            "is_simulated": is_simulated,
            "user_query": "Provide detailed treatment and prevention advice for this disease."
        }
        
        prompt = f"Analyze this disease detection result and provide structured advice:\n{json.dumps(input_data, indent=2)}"
        
        print(f"Requesting advice for {disease_name}...", flush=True)
        advice_json = generate_with_ollama(prompt)
        
        # Merge vision results with LLM advice
        return {
            "disease_detection": {
                "name": disease_name,
                "confidence": confidence,
                "is_simulated": is_simulated
            },
            "advice": advice_json
        }

    except Exception as e:
        return {"error": f"Error during analysis: {str(e)}"}

# --- Main function to handle the interactive loop ---
def main() -> None:
    print_header()
    
    if not check_ollama_connection():
        print("Fatal: Could not connect to backend. Exiting.", flush=True)
        sys.exit(1)

    parser = argparse.ArgumentParser(description="AgroGPT Mobile Backend CLI")
    parser.add_argument("--prompt", type=str, default=None, help="Single question to answer")
    args = parser.parse_args()

    if args.prompt:
        print(f"Prompt: {args.prompt}", flush=True)
        print("-" * 40)
        response = generate_with_ollama(args.prompt)
        print(json.dumps(response, indent=2), flush=True)
        return
    
    print("Interactive mode (JSON). Type your question and press Enter.", flush=True)
    while True:
        try:
            user_input = input("AgroGPT> ").strip()
        except EOFError:
            break
        
        if not user_input or user_input.lower() in {"exit", "quit"}:
            break
        
        # Simple wrapper for CLI testing
        response = generate_with_ollama(user_input)
        print(json.dumps(response, indent=2), flush=True)
        print("-" * 40)

if __name__ == "__main__":
    main()