import gradio as gr
from transformers import pipeline
from PIL import Image
import torch
import os
import spaces

# --- Initialize the Model Pipeline ---
print("Loading MedGemma model...")
try:
    pipe = pipeline(
        "image-text-to-text",
        model="google/medgemma-4b-it",
        torch_dtype=torch.bfloat16,
        device_map="auto",
        token=os.environ.get("HF_TOKEN")
    )
    model_loaded = True
    print("Model loaded successfully!")
except Exception as e:
    model_loaded = False
    print(f"Error loading model: {e}")

# --- Core Analysis Function ---
@spaces.GPU()
def analyze_symptoms(symptom_image, symptoms_text):
    """
    Analyzes user's symptoms using the MedGemma model, based on text and an optional image.
    """
    if not model_loaded:
        return "Error: The AI model could not be loaded. Please check the Space logs."
        
    if symptom_image is None and not symptoms_text.strip():
        return "Please describe your symptoms or upload an image for analysis."

    try:
        if symptoms_text.strip():
            prompt_for_model = (
                "Analyze the potential medical condition based on the following information. "
                f"The user's description of their symptoms is: '{symptoms_text}'. "
                "If an image is provided, describe what you see in the image and factor it into your analysis. "
                "Provide a list of possible conditions and a clear, actionable plan."
            )
        else:
            prompt_for_model = (
                "An image has been provided without a text description. "
                "Describe what you see in the image in detail. Based on the visual evidence, "
                "provide a list of possible conditions and a clear, actionable plan."
            )

        user_content = [{"type": "text", "text": prompt_for_model}]
        if symptom_image:
            user_content.append({"type": "image", "image": symptom_image})
        
        messages = [
            {
                "role": "system",
                "content": [{
                    "type": "text",
                    "text": "You are an expert, empathetic AI medical assistant. Your role is to analyze a user's symptoms and provide a helpful, safe, and informative response."
                }]
            },
            {
                "role": "user",
                "content": user_content
            }
        ]

        # For debugging, let's see what the output looks like in the logs
        print("Generating pipeline output...")
        output = pipe(messages, max_new_tokens=512, do_sample=True, temperature=0.7)
        print("Pipeline Output:", output) # This will show in the Hugging Face Space logs

        # --- THE FIX ---
        # The error proves the output is not a complex list of dicts.
        # We will treat it as a simple string, which is more robust.
        result = output[0]["generated_text"]

        disclaimer = "\n\n***Disclaimer: I am an AI assistant and not a medical professional. This is not a diagnosis. Please consult a doctor for any health concerns.***"
        
        return result + disclaimer

    except Exception as e:
        # Added more detailed error logging
        print(f"An exception occurred during analysis: {type(e).__name__}: {e}")
        return f"Error during analysis: {str(e)}"

# --- Create the Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft(), title="AI Symptom Analyzer") as demo:
    # Header and disclaimer remain the same
    gr.HTML("""
        <div style="text-align: center; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 2rem; border-radius: 10px; margin-bottom: 2rem;">
            <h1>🩺 AI Symptom Analyzer</h1>
            <p>Advanced symptom analysis powered by Google's MedGemma AI</p>
        </div>
    """)
    gr.HTML("""
        <div style="background-color: #fff3cd; border: 1px solid #ffeaa7; border-radius: 8px; padding: 1rem; margin: 1rem 0; color: #856404;">
            <strong>⚠️ Medical Disclaimer:</strong> This AI tool is for informational purposes only and is not a substitute for professional medical diagnosis or treatment.
        </div>
    """)
    
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            gr.Markdown("### 1. Describe Your Symptoms")
            symptoms_input = gr.Textbox(
                label="Symptoms",
                placeholder="e.g., 'I have a rash on my arm that is red and itchy, and it has been there for two days.'",
                lines=5,
            )

            gr.Markdown("### 2. Upload an Image (Optional)")
            image_input = gr.Image(
                label="Symptom Image",
                type="pil",
                height=300
            )
            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
                analyze_btn = gr.Button("🔍 Analyze Symptoms", variant="primary", size="lg")
                
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Analysis Report")
            output_text = gr.Textbox(
                label="AI Analysis",
                lines=25,
                show_copy_button=True,
                placeholder="Analysis results will appear here..."
            )

    # --- Event Handlers ---
    def clear_all():
        return None, "", ""

    analyze_btn.click(
        fn=analyze_symptoms,
        inputs=[image_input, symptoms_input],
        outputs=output_text
    )

    clear_btn.click(
        fn=clear_all,
        outputs=[image_input, symptoms_input, output_text]
    )

# --- Launch the App ---
if __name__ == "__main__":
    print("Starting Gradio interface...")
    demo.launch(debug=True)