app.py

import streamlit as st
import google.generativeai as genai
import os
from PIL import Image
import io # Needed for handling image bytes
from typing import Optional, Tuple, Any # For type hinting

# --- Configuration and Initialization ---

# Securely load API key
# Prioritize Streamlit secrets, fall back to environment variable for flexibility
GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY", os.environ.get("GEMINI_API_KEY"))

# Configure Gemini Client (only if key is found)
genai_client_configured = False
if GEMINI_API_KEY:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
        genai_client_configured = True
    except Exception as e:
        st.error(f"Failed to configure Google Generative AI: {e}")
        st.stop() # Stop execution if configuration fails
else:
    st.error("⚠️ Gemini API Key not found. Please configure `GEMINI_API_KEY` in Streamlit secrets or environment variables.")
    st.stop() # Stop execution if no API key

# Initialize models (only if client configured)
# Using specific model versions can be good practice for reproducibility
TEXT_MODEL_NAME = 'gemini-1.5-flash' # Or 'gemini-pro' if preferred
VISION_MODEL_NAME = 'gemini-1.5-flash' # Or 'gemini-pro-vision' if preferred

if genai_client_configured:
    try:
        model = genai.GenerativeModel(TEXT_MODEL_NAME)
        vision_model = genai.GenerativeModel(VISION_MODEL_NAME)
    except Exception as e:
        st.error(f"Failed to initialize Gemini models ({TEXT_MODEL_NAME}, {VISION_MODEL_NAME}): {e}")
        st.stop()
else:
    # This state should technically not be reached due to earlier st.stop() calls,
    # but it's good defensive programming.
    st.error("AI Models could not be initialized due to configuration issues.")
    st.stop()


# --- Core AI Interaction Functions ---

# Define more sophisticated prompts emphasizing analysis and *potential* findings
# CRITICAL: Avoid definitive "diagnosis" language. Focus on assisting clinical judgment.
TEXT_ANALYSIS_PROMPT_TEMPLATE = """
**Medical Information Analysis Request:**

**Context:** Analyze the provided medical text (symptoms, history, or reports).
**Task:**
1.  **Identify Key Findings:** Extract significant symptoms, signs, or data points.
2.  **Potential Considerations:** Based *only* on the provided text, list potential underlying conditions or areas of concern that *might* warrant further investigation by a qualified healthcare professional. Use cautious language (e.g., "suggests potential for," "could be consistent with," "warrants investigation into").
3.  **Risk Factors (if applicable):** Mention any potential risk factors identifiable from the text.
4.  **Information Gaps:** Highlight any missing information that would be crucial for a clinical assessment.
5.  **Disclaimer:** Explicitly state that this analysis is AI-generated, not a diagnosis, and cannot replace professional medical evaluation.

**Input Text:**
---
{text_input}
---

**Analysis:**
"""

IMAGE_ANALYSIS_PROMPT_TEMPLATE = """
**Medical Image Analysis Request:**

**Context:** Analyze the provided medical image. User may provide additional context or questions.
**Task:**
1.  **Describe Visible Structures:** Briefly describe the main anatomical structures visible.
2.  **Identify Potential Anomalies:** Point out any visible areas that *appear* abnormal or deviate from typical presentation (e.g., "potential opacity in the lower left lung field," "area of altered signal intensity," "possible asymmetry"). Use cautious, descriptive language.
3.  **Correlate with User Prompt (if provided):** If the user asked a specific question, address it based *only* on the visual information.
4.  **Limitations:** State that image quality, view, and lack of clinical context limit the analysis.
5.  **Disclaimer:** Explicitly state this is an AI-based visual analysis, not a radiological interpretation or diagnosis, and requires review by a qualified radiologist or physician alongside clinical information.

**User's Additional Context/Question (if any):**
---
{user_prompt}
---

**Image Analysis:**
"""

def analyze_medical_text(text_input: str) -> Tuple[Optional[str], Optional[str]]:
    """
    Sends medical text to the Gemini model for analysis using a structured prompt.

    Args:
        text_input: The medical text provided by the user.

    Returns:
        A tuple containing:
        - The analysis text (str) if successful, None otherwise.
        - An error message (str) if an error occurred, None otherwise.
    """
    if not text_input:
        return None, "Input text cannot be empty."
    try:
        prompt = TEXT_ANALYSIS_PROMPT_TEMPLATE.format(text_input=text_input)
        response = model.generate_content(prompt)
        # Add safety check for response structure if needed (e.g., check for blocked content)
        if response.parts:
            return response.text, None
        elif response.prompt_feedback.block_reason:
             return None, f"Analysis blocked due to: {response.prompt_feedback.block_reason.name}. Please revise input."
        else:
             return None, "Received an empty or unexpected response from the AI."

    except Exception as e:
        st.error(f"An error occurred during text analysis: {e}") # Log for debugging
        return None, f"Error communicating with the AI model. Details: {e}"

def analyze_medical_image(image_file: Any, user_prompt: str = "") -> Tuple[Optional[str], Optional[str]]:
    """
    Sends a medical image (and optional prompt) to the Gemini Vision model for analysis.

    Args:
        image_file: The uploaded image file object from Streamlit.
        user_prompt: Optional text context or specific questions from the user.

    Returns:
        A tuple containing:
        - The analysis text (str) if successful, None otherwise.
        - An error message (str) if an error occurred, None otherwise.
    """
    if not image_file:
        return None, "Image file cannot be empty."
    try:
        # Ensure image is opened correctly, handle potential errors
        try:
            image = Image.open(image_file)
            # Optional: Convert image to a supported format if needed, e.g., RGB
            if image.mode != 'RGB':
                 image = image.convert('RGB')
        except Exception as img_e:
             return None, f"Error opening or processing image file: {img_e}"

        prompt_text = IMAGE_ANALYSIS_PROMPT_TEMPLATE.format(user_prompt=user_prompt if user_prompt else "N/A")

        # Prepare content for the vision model
        model_input = [prompt_text, image]

        response = vision_model.generate_content(model_input)

        # Add safety check for response structure
        if response.parts:
            return response.text, None
        elif response.prompt_feedback.block_reason:
            return None, f"Analysis blocked due to: {response.prompt_feedback.block_reason.name}. This might be due to sensitive content policies."
        else:
            return None, "Received an empty or unexpected response from the AI."

    except Exception as e:
        st.error(f"An error occurred during image analysis: {e}") # Log for debugging
        return None, f"Error communicating with the AI model. Details: {e}"


# --- Streamlit User Interface ---

def main():
    st.set_page_config(page_title="AI Medical Information Assistant", layout="wide")

    st.title("🤖 AI Medical Information Assistant")
    st.caption(f"Powered by Google Gemini ({TEXT_MODEL_NAME} / {VISION_MODEL_NAME})")

    # --- CRITICAL DISCLAIMER ---
    st.warning(
        """
        **IMPORTANT DISCLAIMER:**
        *   This tool uses AI to analyze information but **DOES NOT PROVIDE MEDICAL ADVICE OR DIAGNOSIS.**
        *   The analysis is based solely on the input provided and may be incomplete, inaccurate, or lack clinical context.
        *   **ALWAYS consult a qualified healthcare professional** for any health concerns, diagnosis, or treatment decisions.
        *   Do not rely on this tool for medical decisions. It is for informational and educational purposes only, potentially assisting clinicians by highlighting areas for review.
        *   **Do not upload identifiable patient information** unless you have explicit consent and comply with all privacy regulations (e.g., HIPAA).
        """
    )

    st.sidebar.header("Input Options")
    input_method = st.sidebar.radio(
        "Select Input Type:",
        ("Text Description", "Medical Image"),
        help="Choose whether to analyze text-based medical information or a medical image."
    )

    st.sidebar.markdown("---") # Visual separator

    col1, col2 = st.columns(2)

    with col1:
        st.header("Input")
        if input_method == "Text Description":
            st.subheader("Enter Medical Text")
            text_input = st.text_area(
                "Paste symptoms, patient history excerpt, or non-identifiable report sections:",
                height=300,
                placeholder="Example: 55-year-old male presents with intermittent chest pain, worse on exertion. History of hypertension. Non-smoker...",
                key="text_input_area" # Add key for potential state management
            )
            analyze_button = st.button("Analyze Text", key="analyze_text_button", type="primary")

            if analyze_button and text_input:
                with st.spinner("🧠 Analyzing Text... Please wait."):
                    analysis_result, error_message = analyze_medical_text(text_input)

                if error_message:
                     with col2:
                        st.error(f"Analysis Failed: {error_message}")
                elif analysis_result:
                     with col2:
                        st.header("Analysis Results")
                        st.markdown(analysis_result) # Use markdown for better formatting
                else:
                    # Handle unexpected case where both are None
                     with col2:
                         st.error("An unknown error occurred during analysis.")

            elif analyze_button and not text_input:
                 st.warning("Please enter some text to analyze.")


        elif input_method == "Medical Image":
            st.subheader("Upload Medical Image")
            image_file = st.file_uploader(
                "Upload an image (e.g., X-ray, CT slice, pathology slide). Supported: PNG, JPG, JPEG.",
                type=["png", "jpg", "jpeg"],
                key="image_uploader"
            )
            user_image_prompt = st.text_input(
                "Optional: Add context or specific questions for the image analysis:",
                placeholder="Example: 'Check for abnormalities in the left lung apex' or 'Patient context: Follow-up scan after treatment'",
                key="image_prompt_input"
            )
            analyze_button = st.button("Analyze Image", key="analyze_image_button", type="primary")

            if analyze_button and image_file:
                 # Display uploaded image immediately for confirmation
                 st.image(image_file, caption="Uploaded Image Preview", use_column_width=True)
                 with st.spinner("🖼️ Analyzing Image... Please wait."):
                     analysis_result, error_message = analyze_medical_image(image_file, user_image_prompt)

                 if error_message:
                     with col2:
                        st.error(f"Analysis Failed: {error_message}")
                 elif analysis_result:
                     with col2:
                        st.header("Image Analysis Results")
                        st.markdown(analysis_result)
                 else:
                      with col2:
                        st.error("An unknown error occurred during analysis.")

            elif analyze_button and not image_file:
                 st.warning("Please upload an image file to analyze.")


    # Display placeholder in result column if nothing submitted yet
    # Check if 'analysis_result' or 'error_message' exists in session state if needed for persistence,
    # but for this simpler flow, checking button presses is often sufficient.
    # A simple way is to check if the button was pressed in the current run.
    button_pressed = st.session_state.get('analyze_text_button', False) or st.session_state.get('analyze_image_button', False)
    if not button_pressed:
         with col2:
             st.info("Analysis results will appear here once input is submitted.")


    st.sidebar.markdown("---")
    st.sidebar.info("Remember: This AI tool is an assistant, not a substitute for professional medical expertise.")

if __name__ == "__main__":
    main()