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"Fatal Error: Failed to configure Google Generative AI. Check API Key. Details: {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 (Consider more powerful model for agentic reasoning if needed)
# Using 1.5 Pro for text (agentic) and 1.5 Flash for vision might be a good balance
TEXT_MODEL_NAME = 'gemini-1.5-pro-latest' # Model for agentic text reasoning
VISION_MODEL_NAME = 'gemini-1.5-flash' # Model for image analysis

if genai_client_configured:
    try:
        text_model = genai.GenerativeModel(TEXT_MODEL_NAME)
        vision_model = genai.GenerativeModel(VISION_MODEL_NAME)
        st.success(f"Successfully initialized models: Text ({TEXT_MODEL_NAME}), Vision ({VISION_MODEL_NAME})", icon="✅")
    except Exception as e:
        st.error(f"Fatal Error: Failed to initialize Gemini models. Text: {TEXT_MODEL_NAME}, Vision: {VISION_MODEL_NAME}. Details: {e}")
        st.stop()
else:
    st.error("AI Models could not be initialized due to configuration issues.")
    st.stop()


# --- Core AI Interaction Functions ---

# AGENTIC prompt for Text Analysis
AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE = """
**Simulated Clinical Reasoning Agent Task:**

**Role:** AI assistant simulating an agentic clinical reasoning process to support a healthcare professional by structuring information, generating possibilities, and suggesting investigation pathways based *strictly* on the provided text. **This is NOT a diagnosis.**

**Input Data:** Unstructured clinical information (e.g., symptoms, history, basic findings).

**Simulated Agentic Steps (Perform sequentially):**

1.  **Information Extraction & Structuring:**
    *   Key Demographics (Age, Sex if provided).
    *   Primary Symptoms/Signs.
    *   Relevant Medical History.
    *   Pertinent Negatives (if mentioned).

2.  **Differential Considerations Generation:**
    *   Based *only* on Step 1, list **potential differential considerations** (possible conditions).
    *   **Use cautious language:** "could be consistent with," "warrants consideration," "less likely but possible." **AVOID definitive statements.**
    *   Briefly justify each consideration based on findings.

3.  **Information Gap Analysis:**
    *   Identify critical missing information (e.g., lab results, imaging, exam specifics, duration/onset).

4.  **Suggested Next Steps for Investigation (for Clinician):**
    *   Propose logical next steps a **healthcare professional might consider**.
    *   Categorize (e.g., Further History, Exam Points, Labs, Imaging).
    *   Frame as *suggestions* (e.g., "Consider ordering...", "Assessment of X may be informative").

5.  **Mandatory Disclaimer:** Conclude with: "This AI-generated analysis is for informational support only. It is **NOT** a diagnosis and cannot replace the judgment of a qualified healthcare professional."

**Input Clinical Information:**
---
{text_input}
---

**Agentic Analysis:**
"""

# Standard prompt for Image 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 main anatomical structures.
2.  **Identify Potential Anomalies:** Point out areas that *appear* abnormal or deviate from typical presentation (e.g., "potential opacity," "altered signal intensity," "possible asymmetry"). Use cautious, descriptive language.
3.  **Correlate with User Prompt (if provided):** Address specific user questions based *only* on visual information.
4.  **Limitations:** State that image quality, view, and lack of clinical context limit analysis.
5.  **Disclaimer:** Explicitly state this is AI visual analysis, not radiological interpretation or diagnosis, requiring review by a qualified professional with clinical context.

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

**Image Analysis:**
"""

def run_agentic_text_analysis(text_input: str) -> Tuple[Optional[str], Optional[str]]:
    """
    Sends clinical text to the Gemini text model for simulated agentic analysis.

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

    Returns:
        Tuple: (analysis_text, error_message)
    """
    if not text_input or not text_input.strip():
        return None, "Input text cannot be empty."
    try:
        prompt = AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE.format(text_input=text_input)
        # Using the designated text model
        response = text_model.generate_content(prompt)

        if response.parts:
            return response.text, None
        elif response.prompt_feedback.block_reason:
             return None, f"Analysis blocked by safety filters: {response.prompt_feedback.block_reason.name}. Review input."
        else:
             candidate = response.candidates[0] if response.candidates else None
             if candidate and candidate.finish_reason != "STOP":
                 return None, f"Analysis stopped prematurely. Reason: {candidate.finish_reason.name}."
             else:
                 return None, "Received an empty or unexpected response from the AI model."

    except Exception as e:
        st.error(f"Error during agentic text analysis: {e}", icon="🚨")
        return None, f"Error communicating with the AI model for text analysis. 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: Uploaded image file object from Streamlit.
        user_prompt: Optional text context/questions from the user.

    Returns:
        Tuple: (analysis_text, error_message)
    """
    if not image_file:
        return None, "Image file cannot be empty."
    try:
        try:
            image = Image.open(image_file)
            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")
        model_input = [prompt_text, image]
        # Using the designated vision model
        response = vision_model.generate_content(model_input)

        if response.parts:
            return response.text, None
        elif response.prompt_feedback.block_reason:
            return None, f"Image analysis blocked by safety filters: {response.prompt_feedback.block_reason.name}. May relate to sensitive content policies."
        else:
             candidate = response.candidates[0] if response.candidates else None
             if candidate and candidate.finish_reason != "STOP":
                 return None, f"Image analysis stopped prematurely. Reason: {candidate.finish_reason.name}."
             else:
                 return None, "Received an empty or unexpected response from the AI model for image analysis."

    except Exception as e:
        st.error(f"Error during image analysis: {e}", icon="🖼️")
        return None, f"Error communicating with the AI model for image analysis. Details: {e}"


# --- Streamlit User Interface ---

def main():
    st.set_page_config(
        page_title="AI Clinical Support Demonstrator",
        layout="wide",
        initial_sidebar_state="expanded"
        )

    st.title("🤖 AI Clinical Support Demonstrator")
    st.caption(f"Agentic Text Analysis ({TEXT_MODEL_NAME}) & Image Analysis ({VISION_MODEL_NAME})")
    st.markdown("---")

    # --- CRITICAL DISCLAIMER ---
    st.warning(
        """
        **🔴 IMPORTANT SAFETY & USE DISCLAIMER 🔴**
        *   This tool **DEMONSTRATES** AI capabilities. It **DOES NOT** provide medical advice or diagnosis.
        *   **Agentic Text Analysis:** Simulates reasoning on text input. Output is illustrative, not diagnostic.
        *   **Image Analysis:** Provides observations on images. Output is **NOT** a radiological interpretation.
        *   AI analysis lacks full clinical context, may be inaccurate, and **CANNOT** replace professional judgment.
        *   **ALWAYS consult qualified healthcare professionals** for diagnosis and treatment.
        *   **PRIVACY:** Do **NOT** upload identifiable patient information (PHI) without explicit consent and adherence to all privacy laws (e.g., HIPAA).
        """,
        icon="⚠️"
    )
    st.markdown("---")


    st.sidebar.header("Analysis Options")
    input_method = st.sidebar.radio(
        "Select Analysis Type:",
        ("Agentic Text Analysis", "Medical Image Analysis"),
        key="input_method_radio",
        help="Choose 'Agentic Text Analysis' for reasoning simulation on clinical text, or 'Medical Image Analysis' for visual observations on images."
    )
    st.sidebar.markdown("---") # Visual separator

    col1, col2 = st.columns(2)

    with col1:
        st.header("Input Data")
        analysis_result = None # Initialize results variables
        error_message = None
        output_header = "Analysis Results" # Default header for the output column

        # --- Agentic Text Analysis Input ---
        if input_method == "Agentic Text Analysis":
            st.subheader("Clinical Text for Agentic Analysis")
            text_input = st.text_area(
                "Paste de-identified clinical information (symptoms, history, findings):",
                height=350, # Slightly larger text area
                placeholder="Example: 68yo male, sudden SOB & pleuritic chest pain post-flight. HR 110, SpO2 92% RA. No known cardiac hx...",
                key="text_input_area"
            )
            analyze_button = st.button("▶️ Run Agentic Text Analysis", key="analyze_text_button", type="primary")

            if analyze_button:
                if text_input:
                    with st.spinner("🧠 Simulating agentic reasoning..."):
                        analysis_result, error_message = run_agentic_text_analysis(text_input)
                    output_header = "Simulated Agentic Analysis Output"
                else:
                    st.warning("Please enter clinical text to analyze.", icon="☝️")

        # --- Medical Image Analysis Input ---
        elif input_method == "Medical Image Analysis":
            st.subheader("Medical Image for Analysis")
            image_file = st.file_uploader(
                "Upload a de-identified medical image (e.g., X-ray, CT slice). Supported: PNG, JPG, JPEG.",
                type=["png", "jpg", "jpeg"],
                key="image_uploader"
            )
            user_image_prompt = st.text_input(
                "Optional: Add context or specific question for image analysis:",
                placeholder="Example: 'Describe findings in the lung fields' or 'Any visible fractures?'",
                key="image_prompt_input"
            )
            analyze_button = st.button("🖼️ Analyze Medical Image", key="analyze_image_button", type="primary")

            if analyze_button:
                if image_file:
                    st.image(image_file, caption="Uploaded Image Preview", use_column_width=True)
                    with st.spinner("👁️ Analyzing image..."):
                        analysis_result, error_message = analyze_medical_image(image_file, user_image_prompt)
                    output_header = "Medical Image Analysis Output"
                else:
                    st.warning("Please upload an image file to analyze.", icon="☝️")

    # --- Output Column ---
    with col2:
        st.header(output_header)
        # Check if a button was pressed in this run (using session state keys is more robust for complex apps,
        # but checking the result variables works here as they are reset on each run unless persisted).
        button_pressed = st.session_state.get('analyze_text_button', False) or st.session_state.get('analyze_image_button', False)

        if button_pressed: # Only show results if a button was pressed in this run cycle
            if error_message:
                 st.error(f"Analysis Failed: {error_message}", icon="❌")
            elif analysis_result:
                 st.markdown(analysis_result) # Display the successful result
            # Removed the "unknown error" case here, as the functions should return either result or error message
        else:
            st.info("Analysis results will appear here after providing input and clicking the corresponding analysis button.")


    # --- Sidebar Explanations ---
    st.sidebar.markdown("---")
    st.sidebar.header("About The Prompts")
    with st.sidebar.expander("View Agentic Text Prompt Structure"):
        st.markdown(f"```plaintext\n{AGENTIC_TEXT_ANALYSIS_PROMPT_TEMPLATE.split('---')[0]} ... [Input Text] ...\n```")
        st.caption("Guides the AI through structured reasoning steps for text.")
    with st.sidebar.expander("View Image Analysis Prompt Structure"):
        st.markdown(f"```plaintext\n{IMAGE_ANALYSIS_PROMPT_TEMPLATE.split('---')[0]} ... [User Prompt] ...\n```")
        st.caption("Guides the AI to describe visual features and potential anomalies in images.")

    st.sidebar.markdown("---")
    st.sidebar.error(
        "**Ethical Use Reminder:** AI in medicine requires extreme caution. This tool is for demonstration and education, not clinical practice. Verify all information and rely on professional expertise."
        )

if __name__ == "__main__":
    main()