import gradio as gr
from openai import OpenAI
import base64
from io import BytesIO
from PIL import Image  # Import PIL for image conversion

def analyze_plant_image(image, api_key):
    try:
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )
        
        # Convert NumPy array to PIL Image and ensure RGB mode for JPEG compatibility
        image = Image.fromarray(image).convert("RGB")
        
        # Save the image to a BytesIO buffer as JPEG
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        
        # Encode the image to base64
        encoded_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
        
        # Create the data URL for the image
        image_url = f"data:image/jpeg;base64,{encoded_image}"
        
        # Call the InternVL3 14B model with the image and prompt
        completion = client.chat.completions.create(
            model="opengvlab/internvl3-14b:free",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Identify this plant, provide care instructions, and diagnose any health issues visible in the image."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": image_url
                            }
                        }
                    ]
                }
            ]
        )
        
        # Return the model's response
        return completion.choices[0].message.content
    
    except Exception as e:
        # Return error message if something goes wrong
        return f"An error occurred: {str(e)}"

# Create the Gradio interface
interface = gr.Interface(
    fn=analyze_plant_image,
    inputs=[
        gr.Image(label="Upload Plant Image"),  # Expects a NumPy array
        gr.Textbox(type="password", label="OpenRouter API Key")  # Secure input for API key
    ],
    outputs=gr.Textbox(label="Analysis Result"),  # Displays the result or error
    title="PlantPal: Your AI-Powered Plant Care Assistant",
    description="""
    Upload an image of your plant and enter your OpenRouter API key to get started.
    If you don't have an API key, you can get one from [OpenRouter](https://openrouter.ai/).
    """
)

# Launch the app
interface.launch()