import os
from PIL import Image as PILImage
from agno.agent import Agent
from agno.models.google import Gemini
from agno.tools.duckduckgo import DuckDuckGoTools
from agno.media import Image as AgnoImage
import gradio as gr

# Set your API Key (Replace with your actual key)
GOOGLE_API_KEY = "AIzaSyCUPYFW5ESk1YIwSjwpGd3jZJJ7oPAX4_s"
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY

# Ensure API Key is provided
if not GOOGLE_API_KEY:
    raise ValueError("⚠️ Please set your Google API Key in GOOGLE_API_KEY")

# Initialize the Medical Agent
medical_agent = Agent(
    model=Gemini(id="gemini-2.0-flash-exp"),
    tools=[DuckDuckGoTools()],
    markdown=True
)

# Medical Analysis Query
query = """
You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging. Analyze the medical image and structure your response as follows:

### 1. Image Type & Region
- Identify imaging modality (X-ray/MRI/CT/Ultrasound/etc.).
- Specify anatomical region and positioning.
- Evaluate image quality and technical adequacy.

### 2. Key Findings
- Highlight primary observations systematically.
- Identify potential abnormalities with detailed descriptions.
- Include measurements and densities where relevant.

### 3. Diagnostic Assessment
- Provide primary diagnosis with confidence level.
- List differential diagnoses ranked by likelihood.
- Support each diagnosis with observed evidence.
- Highlight critical/urgent findings.

### 4. Patient-Friendly Explanation
- Simplify findings in clear, non-technical language.
- Avoid medical jargon or provide easy definitions.
- Include relatable visual analogies.

### 5. Research Context
- Use DuckDuckGo search to find recent medical literature.
- Search for standard treatment protocols.
- Provide 2-3 key references supporting the analysis.

Ensure a structured and medically accurate response using clear markdown formatting.
"""

# Function to analyze medical image
def analyze_medical_image(image):
    """Processes and analyzes a medical image using AI."""
    if image is None:
        return "⚠️ Please upload an image to analyze."

    # Save the input image to a temporary file
    temp_path = "temp_image.png"
    image.save(temp_path)

    # Create AgnoImage object
    agno_image = AgnoImage(filepath=temp_path)

    # Run AI analysis
    try:
        response = medical_agent.run(query, images=[agno_image])
        return response.content
    except Exception as e:
        return f"⚠️ Analysis error: {e}"
    finally:
        # Clean up temporary file
        if os.path.exists(temp_path):
            os.remove(temp_path)

# Create Gradio interface
with gr.Blocks(title="Medical Image Analysis") as demo:
    gr.Markdown(
        """
        # 🩺 Medical Image Analysis Tool 🔬

        Welcome to the **Medical Image Analysis** tool! 📸

        Upload a medical image (X-ray, MRI, CT, Ultrasound, etc.), and our AI-powered system will analyze it,
        providing detailed findings, diagnosis, and research insights.

        Let's get started!
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Upload Medical Image")
            analyze_button = gr.Button("Analyze Image", variant="primary")

        with gr.Column(scale=2):
            output_report = gr.Markdown(label="Analysis Report")

    analyze_button.click(
        fn=analyze_medical_image,
        inputs=input_image,
        outputs=output_report
    )

    # gr.Examples(
    #     examples=[
    #         "examples/xray_chest.jpg",
    #         "examples/mri_brain.jpg",
    #         "examples/ct_abdomen.jpg"
    #     ],
    #     inputs=input_image
    # )

# Launch the application
if __name__ == "__main__":
    demo.launch(debug=True, share=True)