import gradio as gr
from ultralytics import YOLO
import PIL.Image
from groq import Groq
import os

# 1. Initialize YOLO and Groq
# Make sure "best.pt" is uploaded to your Hugging Face Space
model = YOLO("best.pt")
client = Groq(api_key=os.environ.get("Bot"))

def get_urdu_advice(disease_name):
    """Fetches cure and advice in Urdu from Groq LLM"""
    prompt = f"The plant has been diagnosed with {disease_name}. Please provide a brief description of this disease and its cure in Urdu language for a farmer."
    
    completion = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[{"role": "user", "content": prompt}],
    )
    return completion.choices[0].message.content

def predict_and_advise(input_img):
    # Detection
    results = model.predict(source=input_img, conf=0.25)
    res_plotted = results[0].plot()
    
    # Get the name of the top detected disease
    if len(results[0].boxes) > 0:
        # Extract class name of the first detection
        class_id = int(results[0].boxes.cls[0])
        disease_name = results[0].names[class_id]
        urdu_advice = get_urdu_advice(disease_name)
    else:
        disease_name = "No disease detected"
        urdu_advice = "پودا صحت مند لگ رہا ہے یا کوئی بیماری نہیں ملی۔"

    # Convert BGR to RGB for output
    output_img = PIL.Image.fromarray(res_plotted[:, :, ::-1])
    return output_img, urdu_advice

# 3. Gradio Blocks UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌾 AgriBot PK: Rice Disease Detector")
    gr.Markdown("پودے کی تصویر اپ لوڈ کریں تاکہ بیماری کی شناخت ہو سکے اور علاج معلوم کیا جا سکے۔")
    
    with gr.Row():
        input_file = gr.Image(type="pil", label="Upload Leaf Image")
        output_image = gr.Image(type="pil", label="Detected Disease")
        
    urdu_text = gr.Textbox(label="Disease Info & Cure (Urdu)", lines=10)
    
    submit_btn = gr.Button("Detect & Get Advice")
    submit_btn.click(fn=predict_and_advise, inputs=input_file, outputs=[output_image, urdu_text])

demo.launch()