app.py

import streamlit as st
from langchain_groq import ChatGroq
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import tool
from typing import Optional, List, Dict, Any
import json
import re # For parsing vitals like BP

# --- Configuration & Constants ---
class ClinicalAppSettings:
    APP_TITLE = "SynapseAI: Advanced Clinical Decision Support"
    PAGE_LAYOUT = "wide"
    MODEL_NAME = "llama3-70b-8192" # Use a powerful model like Groq's Llama3 70b
    TEMPERATURE = 0.1
    MAX_SEARCH_RESULTS = 3

class ClinicalPrompts:
    SYSTEM_PROMPT = """
    You are SynapseAI, an expert AI clinical assistant designed to support healthcare professionals.
    Your primary function is to analyze patient data, provide differential diagnoses, suggest evidence-based management plans, and identify potential risks according to the latest medical guidelines and safety protocols.

    **Core Directives:**
    1.  **Comprehensive Analysis:** Thoroughly analyze ALL provided patient data (demographics, HPI, PMH, PSH, Allergies, Meds, SH, FH, ROS, Vitals, Exam).
    2.  **Structured Output:** ALWAYS format your response using the following JSON structure:
        ```json
        {
          "assessment": "Concise summary of the patient's presentation and key findings.",
          "differential_diagnosis": [
            {"diagnosis": "Primary Diagnosis", "likelihood": "High/Medium/Low", "rationale": "Supporting evidence..."},
            {"diagnosis": "Alternative Diagnosis 1", "likelihood": "Medium/Low", "rationale": "Supporting/Refuting evidence..."},
            {"diagnosis": "Alternative Diagnosis 2", "likelihood": "Low", "rationale": "Why it's less likely but considered..."}
          ],
          "risk_assessment": {
            "identified_red_flags": ["List any triggered red flags based on input"],
            "immediate_concerns": ["Specific urgent issues requiring attention (e.g., sepsis risk, ACS rule-out)"],
            "potential_complications": ["Possible future issues based on presentation"]
          },
          "recommended_plan": {
            "investigations": ["List specific lab tests or imaging required. Use 'order_lab_test' tool."],
            "therapeutics": ["Suggest specific treatments or prescriptions. Use 'prescribe_medication' tool."],
            "consultations": ["Recommend specialist consultations if needed."],
            "patient_education": ["Key points for patient communication."]
          },
          "rationale_summary": "Brief justification for the overall assessment and plan, referencing guidelines or evidence where possible. Use 'tavily_search_results' tool if needed to find supporting evidence/guidelines.",
          "interaction_check_summary": "Summary of findings from the 'check_drug_interactions' tool IF a new medication was considered or prescribed."
        }
        ```
    3.  **Safety First - Red Flags:** Immediately identify and report any conditions matching the defined RED_FLAGS. Use the `flag_risk` tool if critical.
    4.  **Safety First - Drug Interactions:** BEFORE suggesting *any* new prescription, you MUST use the `check_drug_interactions` tool to verify against the patient's current medications and allergies. Mention the result in `interaction_check_summary`.
    5.  **Tool Utilization:** Employ the provided tools (`order_lab_test`, `prescribe_medication`, `check_drug_interactions`, `flag_risk`, `tavily_search_results`) precisely when indicated by your plan. Adhere strictly to tool schemas. Do NOT hallucinate tool usage results; wait for actual tool output if required in a multi-turn scenario (though this implementation focuses on single-turn analysis with tool calls).
    6.  **Evidence-Based:** Briefly cite reasoning, drawing on general medical knowledge. Use Tavily Search for specific guideline checks or novel information when necessary.
    7.  **Clarity and Conciseness:** Be clear, avoiding ambiguity. Use standard medical terminology.
    """

# --- Mock Data / Helpers ---
# (In a real system, this would be a proper API/database)
MOCK_INTERACTION_DB = {
    ("Lisinopril", "Spironolactone"): "High risk of hyperkalemia. Monitor potassium closely.",
    ("Warfarin", "Amiodarone"): "Increased bleeding risk. Monitor INR frequently and adjust Warfarin dose.",
    ("Simvastatin", "Clarithromycin"): "Increased risk of myopathy/rhabdomyolysis. Avoid combination or use lower statin dose.",
    ("Aspirin", "Ibuprofen"): "Concurrent use may decrease Aspirin's cardioprotective effect. Potential for increased GI bleeding."
}

ALLERGY_INTERACTIONS = {
    "Penicillin": ["Amoxicillin", "Ampicillin", "Piperacillin"],
    "Sulfa": ["Sulfamethoxazole", "Sulfasalazine"],
    "Aspirin": ["Ibuprofen", "Naproxen"] # Cross-reactivity example for NSAIDs
}

def parse_bp(bp_string: str) -> Optional[tuple[int, int]]:
    """Parses BP string like '120/80' into (systolic, diastolic) integers."""
    match = re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", bp_string)
    if match:
        return int(match.group(1)), int(match.group(2))
    return None

def check_red_flags(patient_data: dict) -> List[str]:
    """Checks patient data against predefined red flags."""
    flags = []
    symptoms = patient_data.get("hpi", {}).get("symptoms", [])
    vitals = patient_data.get("vitals", {})
    history = patient_data.get("pmh", {}).get("conditions", "")

    # Symptom Flags
    if "chest pain" in [s.lower() for s in symptoms]: flags.append("Red Flag: Chest Pain reported.")
    if "shortness of breath" in [s.lower() for s in symptoms]: flags.append("Red Flag: Shortness of Breath reported.")
    if "severe headache" in [s.lower() for s in symptoms]: flags.append("Red Flag: Severe Headache reported.")
    if "sudden vision loss" in [s.lower() for s in symptoms]: flags.append("Red Flag: Sudden Vision Loss reported.")
    if "weakness on one side" in [s.lower() for s in symptoms]: flags.append("Red Flag: Unilateral Weakness reported (potential stroke).")

    # Vital Sign Flags (add more checks as needed)
    if "temp_c" in vitals and vitals["temp_c"] >= 38.5: flags.append(f"Red Flag: Fever (Temperature: {vitals['temp_c']}°C).")
    if "hr_bpm" in vitals and vitals["hr_bpm"] >= 120: flags.append(f"Red Flag: Tachycardia (Heart Rate: {vitals['hr_bpm']} bpm).")
    if "rr_rpm" in vitals and vitals["rr_rpm"] >= 24: flags.append(f"Red Flag: Tachypnea (Respiratory Rate: {vitals['rr_rpm']} rpm).")
    if "spo2_percent" in vitals and vitals["spo2_percent"] <= 92: flags.append(f"Red Flag: Hypoxia (SpO2: {vitals['spo2_percent']}%).")
    if "bp_mmhg" in vitals:
        bp = parse_bp(vitals["bp_mmhg"])
        if bp:
            if bp[0] >= 180 or bp[1] >= 110: flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {vitals['bp_mmhg']} mmHg).")
            if bp[0] <= 90 and bp[1] <= 60: flags.append(f"Red Flag: Hypotension (BP: {vitals['bp_mmhg']} mmHg).")

    # History Flags (Simple examples)
    if "history of mi" in history.lower() and "chest pain" in [s.lower() for s in symptoms]: flags.append("Red Flag: History of MI with current Chest Pain.")

    return flags

# --- Enhanced Tool Definitions ---

# Use Pydantic models for robust argument validation
class LabOrderInput(BaseModel):
    test_name: str = Field(..., description="Specific name of the lab test or panel (e.g., 'CBC', 'BMP', 'Troponin I', 'Urinalysis').")
    reason: str = Field(..., description="Clinical justification for ordering the test (e.g., 'Rule out infection', 'Assess renal function', 'Evaluate for ACS').")
    priority: str = Field("Routine", description="Priority of the test (e.g., 'STAT', 'Routine').")

@tool("order_lab_test", args_schema=LabOrderInput)
def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
    """Orders a specific lab test with clinical justification and priority."""
    return json.dumps({
        "status": "success",
        "message": f"Lab Ordered: {test_name} ({priority})",
        "details": f"Reason: {reason}"
    })

class PrescriptionInput(BaseModel):
    medication_name: str = Field(..., description="Name of the medication.")
    dosage: str = Field(..., description="Dosage amount and unit (e.g., '500 mg', '10 mg').")
    route: str = Field(..., description="Route of administration (e.g., 'PO', 'IV', 'IM', 'Topical').")
    frequency: str = Field(..., description="How often the medication should be taken (e.g., 'BID', 'QDaily', 'Q4-6H PRN').")
    duration: str = Field("As directed", description="Duration of treatment (e.g., '7 days', '1 month', 'Until follow-up').")
    reason: str = Field(..., description="Clinical indication for the prescription.")

@tool("prescribe_medication", args_schema=PrescriptionInput)
def prescribe_medication(medication_name: str, dosage: str, route: str, frequency: str, duration: str, reason: str) -> str:
    """Prescribes a medication with detailed instructions and clinical indication."""
    # In a real scenario, this would trigger an e-prescription workflow
    return json.dumps({
        "status": "success",
        "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}",
        "details": f"Duration: {duration}. Reason: {reason}"
    })

class InteractionCheckInput(BaseModel):
    potential_prescription: str = Field(..., description="The name of the NEW medication being considered.")
    current_medications: List[str] = Field(..., description="List of the patient's CURRENT medication names.")
    allergies: List[str] = Field(..., description="List of the patient's known allergies.")

@tool("check_drug_interactions", args_schema=InteractionCheckInput)
def check_drug_interactions(potential_prescription: str, current_medications: List[str], allergies: List[str]) -> str:
    """Checks for potential drug-drug and drug-allergy interactions BEFORE prescribing."""
    warnings = []
    potential_med_lower = potential_prescription.lower()

    # Check Allergies
    for allergy in allergies:
        allergy_lower = allergy.lower()
        # Simple direct check
        if allergy_lower == potential_med_lower:
            warnings.append(f"CRITICAL ALLERGY: Patient allergic to {allergy}. Cannot prescribe {potential_prescription}.")
            continue
        # Check cross-reactivity (using simplified mock data)
        if allergy_lower in ALLERGY_INTERACTIONS:
            for cross_reactant in ALLERGY_INTERACTIONS[allergy_lower]:
                if cross_reactant.lower() == potential_med_lower:
                    warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to {allergy}. High risk with {potential_prescription}.")

    # Check Drug-Drug Interactions (using simplified mock data)
    current_meds_lower = [med.lower() for med in current_medications]
    for current_med in current_meds_lower:
        # Check pairs in both orders
        pair1 = (current_med, potential_med_lower)
        pair2 = (potential_med_lower, current_med)
        if pair1 in MOCK_INTERACTION_DB:
            warnings.append(f"Interaction Found: {potential_prescription} with {current_med.capitalize()} - {MOCK_INTERACTION_DB[pair1]}")
        elif pair2 in MOCK_INTERACTION_DB:
             warnings.append(f"Interaction Found: {potential_prescription} with {current_med.capitalize()} - {MOCK_INTERACTION_DB[pair2]}")

    if not warnings:
        return json.dumps({"status": "clear", "message": f"No major interactions identified for {potential_prescription} with current meds/allergies.", "warnings": []})
    else:
        return json.dumps({"status": "warning", "message": f"Potential interactions identified for {potential_prescription}.", "warnings": warnings})

class FlagRiskInput(BaseModel):
    risk_description: str = Field(..., description="Specific critical risk identified (e.g., 'Suspected Sepsis', 'Acute Coronary Syndrome', 'Stroke Alert').")
    urgency: str = Field("High", description="Urgency level (e.g., 'Critical', 'High', 'Moderate').")

@tool("flag_risk", args_schema=FlagRiskInput)
def flag_risk(risk_description: str, urgency: str) -> str:
    """Flags a critical risk identified during analysis for immediate attention."""
    st.error(f"🚨 **{urgency.upper()} RISK FLAGGED:** {risk_description}", icon="🚨")
    return json.dumps({
        "status": "flagged",
        "message": f"Risk '{risk_description}' flagged with {urgency} urgency."
    })


# Initialize Search Tool
search_tool = TavilySearchResults(max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS)

# --- Core Agent Logic ---
class ClinicalAgent:
    def __init__(self):
        self.model = ChatGroq(
            temperature=ClinicalAppSettings.TEMPERATURE,
            model=ClinicalAppSettings.MODEL_NAME
        )
        # Combine all tools
        self.tools = [
            order_lab_test,
            prescribe_medication,
            check_drug_interactions,
            flag_risk,
            search_tool
        ]
        # Bind tools to the model
        self.model_with_tools = self.model.bind_tools(self.tools)
        # History for context (simple implementation)
        self.history = []

    def _format_patient_data_for_prompt(self, data: dict) -> str:
        """Formats the patient dictionary into a readable string for the LLM."""
        prompt_str = "Patient Data:\n"
        for key, value in data.items():
            if isinstance(value, dict):
                prompt_str += f"  {key.replace('_', ' ').title()}:\n"
                for sub_key, sub_value in value.items():
                     if sub_value: # Only include if there's data
                        prompt_str += f"    - {sub_key.replace('_', ' ').title()}: {sub_value}\n"
            elif isinstance(value, list) and value:
                 prompt_str += f"  {key.replace('_', ' ').title()}: {', '.join(map(str, value))}\n"
            elif value: # Only include non-empty fields
                 prompt_str += f"  {key.replace('_', ' ').title()}: {value}\n"
        return prompt_str.strip()


    def analyze(self, patient_data: dict) -> tuple[Optional[dict], List[dict]]:
        """Runs the analysis, handling tool calls and parsing the structured output."""
        try:
            # Add System Prompt and formatted Patient Data
            # Simple history management: add previous messages if any
            messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)]
            # Include history if needed - consider token limits
            # messages.extend(self.history)
            formatted_data = self._format_patient_data_for_prompt(patient_data)
            messages.append(HumanMessage(content=formatted_data))

            # Invoke the model
            ai_response = self.model_with_tools.invoke(messages)

            # Store conversation turn
            # self.history.append(HumanMessage(content=formatted_data))
            # self.history.append(ai_response) # AIMessage includes tool calls

            response_content = None
            tool_calls = []

            if isinstance(ai_response, AIMessage):
                # Check if the response contains the structured JSON output
                try:
                    # Sometimes the JSON is embedded in the content, sometimes it's the primary content
                    # Look for ```json ... ``` block first
                    json_match = re.search(r"```json\n(\{.*?\})\n```", ai_response.content, re.DOTALL)
                    if json_match:
                        response_content = json.loads(json_match.group(1))
                    else:
                        # Try parsing the whole content as JSON
                        response_content = json.loads(ai_response.content)
                except json.JSONDecodeError:
                    st.warning("AI did not return valid JSON in the expected format. Displaying raw content.")
                    st.code(ai_response.content, language=None) # Display raw if not JSON
                    response_content = {"assessment": ai_response.content, "error": "Output format incorrect"}

                # Extract tool calls separately
                if ai_response.tool_calls:
                    tool_calls = ai_response.tool_calls

            return response_content, tool_calls

        except Exception as e:
            st.error(f"Error during AI analysis: {str(e)}")
            return None, []

    def process_tool_call(self, tool_call: Dict[str, Any]) -> Any:
        """Executes a single tool call."""
        tool_name = tool_call.get("name")
        tool_args = tool_call.get("args", {})
        selected_tool = {t.name: t for t in self.tools}.get(tool_name)

        if not selected_tool:
            return json.dumps({"status": "error", "message": f"Unknown tool: {tool_name}"})

        try:
            # Ensure args are correctly passed (Pydantic models handle validation)
            return selected_tool.invoke(tool_args)
        except Exception as e:
            st.error(f"Error executing tool '{tool_name}': {str(e)}")
            return json.dumps({"status": "error", "message": f"Failed to execute {tool_name}: {str(e)}"})

# --- Streamlit UI ---
def main():
    st.set_page_config(page_title=ClinicalAppSettings.APP_TITLE, layout=ClinicalAppSettings.PAGE_LAYOUT)
    st.title(f"🩺 {ClinicalAppSettings.APP_TITLE}")
    st.caption(f"Powered by Langchain & Groq ({ClinicalAppSettings.MODEL_NAME})")

    # Initialize Agent in session state
    if 'agent' not in st.session_state:
        st.session_state.agent = ClinicalAgent()
    if 'analysis_complete' not in st.session_state:
        st.session_state.analysis_complete = False
    if 'analysis_result' not in st.session_state:
        st.session_state.analysis_result = None
    if 'tool_call_results' not in st.session_state:
        st.session_state.tool_call_results = []
    if 'red_flags' not in st.session_state:
         st.session_state.red_flags = []

    # --- Patient Data Input Sidebar ---
    with st.sidebar:
        st.header("📄 Patient Intake Form")

        # Demographics
        st.subheader("Demographics")
        age = st.number_input("Age", min_value=0, max_value=120, value=55)
        sex = st.selectbox("Biological Sex", ["Male", "Female", "Other/Prefer not to say"])

        # History of Present Illness (HPI)
        st.subheader("History of Present Illness (HPI)")
        chief_complaint = st.text_input("Chief Complaint", "Chest pain")
        hpi_details = st.text_area("Detailed HPI", "55 y/o male presents with substernal chest pain started 2 hours ago, described as pressure, radiating to left arm. Associated with nausea and diaphoresis. Pain is 8/10 severity. No relief with rest.")
        symptoms = st.multiselect("Associated Symptoms", ["Nausea", "Diaphoresis", "Shortness of Breath", "Dizziness", "Palpitations", "Fever", "Cough"], default=["Nausea", "Diaphoresis"])

        # Past Medical/Surgical History (PMH/PSH)
        st.subheader("Past History")
        pmh = st.text_area("Past Medical History (PMH)", "Hypertension (HTN), Hyperlipidemia (HLD), Type 2 Diabetes Mellitus (DM2)")
        psh = st.text_area("Past Surgical History (PSH)", "Appendectomy (2005)")

        # Medications & Allergies
        st.subheader("Medications & Allergies")
        current_meds = st.text_area("Current Medications (name, dose, freq)", "Lisinopril 10mg daily\nMetformin 1000mg BID\nAtorvastatin 40mg daily\nAspirin 81mg daily")
        allergies = st.text_area("Allergies (comma separated)", "Penicillin (rash)")

        # Social & Family History (SH/FH)
        st.subheader("Social/Family History")
        social_history = st.text_area("Social History (SH)", "Smoker (1 ppd x 30 years), occasional alcohol.")
        family_history = st.text_area("Family History (FHx)", "Father had MI at age 60. Mother has HTN.")

        # Review of Systems (ROS) - Simplified
        # st.subheader("Review of Systems (ROS)") # Keep UI cleaner for now
        # ros_constitutional = st.checkbox("ROS: Constitutional (Fever, Chills, Weight loss)")
        # ros_cardiac = st.checkbox("ROS: Cardiac (Chest pain, Palpitations)", value=True) # Pre-check based on HPI

        # Vitals & Basic Exam
        st.subheader("Vitals & Exam Findings")
        col1, col2 = st.columns(2)
        with col1:
            temp_c = st.number_input("Temperature (°C)", 35.0, 42.0, 36.8, format="%.1f")
            hr_bpm = st.number_input("Heart Rate (bpm)", 30, 250, 95)
            rr_rpm = st.number_input("Respiratory Rate (rpm)", 5, 50, 18)
        with col2:
            bp_mmhg = st.text_input("Blood Pressure (SYS/DIA)", "155/90")
            spo2_percent = st.number_input("SpO2 (%)", 70, 100, 96)
            pain_scale = st.slider("Pain (0-10)", 0, 10, 8)
        exam_notes = st.text_area("Brief Physical Exam Notes", "Awake, alert, oriented x3. Mild distress. Lungs clear. Cardiac exam: Regular rhythm, no murmurs/gallops. Abdomen soft. No edema.")

        # Clean medication list and allergies for processing
        current_meds_list = [med.strip() for med in current_meds.split('\n') if med.strip()]
        current_med_names = [med.split(' ')[0].strip() for med in current_meds_list] # Simplified name extraction
        allergies_list = [a.strip() for a in allergies.split(',') if a.strip()]

        # Compile Patient Data Dictionary
        patient_data = {
            "demographics": {"age": age, "sex": sex},
            "hpi": {"chief_complaint": chief_complaint, "details": hpi_details, "symptoms": symptoms},
            "pmh": {"conditions": pmh},
            "psh": {"procedures": psh},
            "medications": {"current": current_meds_list, "names_only": current_med_names},
            "allergies": allergies_list,
            "social_history": {"details": social_history},
            "family_history": {"details": family_history},
            # "ros": {"constitutional": ros_constitutional, "cardiac": ros_cardiac}, # Add if using ROS inputs
            "vitals": {
                "temp_c": temp_c, "hr_bpm": hr_bpm, "bp_mmhg": bp_mmhg,
                "rr_rpm": rr_rpm, "spo2_percent": spo2_percent, "pain_scale": pain_scale
            },
            "exam_findings": {"notes": exam_notes}
        }

    # --- Main Analysis Area ---
    st.header("🤖 AI Clinical Analysis")

    # Action Button
    if st.button("Analyze Patient Data", type="primary", use_container_width=True):
        st.session_state.analysis_complete = False
        st.session_state.analysis_result = None
        st.session_state.tool_call_results = []
        st.session_state.red_flags = []

        # 1. Initial Red Flag Check (Client-side before LLM)
        st.session_state.red_flags = check_red_flags(patient_data)
        if st.session_state.red_flags:
            st.warning("**Initial Red Flags Detected:**")
            for flag in st.session_state.red_flags:
                st.warning(f"- {flag}")
            st.warning("Proceeding with AI analysis, but these require immediate attention.")

        # 2. Call AI Agent
        with st.spinner("SynapseAI is processing the case... Please wait."):
            analysis_output, tool_calls = st.session_state.agent.analyze(patient_data)

        if analysis_output:
            st.session_state.analysis_result = analysis_output
            st.session_state.analysis_complete = True

            # 3. Process any Tool Calls requested by the AI
            if tool_calls:
                st.info(f"AI recommended {len(tool_calls)} action(s). Executing...")
                tool_results = []
                with st.spinner("Executing recommended actions..."):
                    for call in tool_calls:
                        st.write(f"⚙️ Requesting: `{call['name']}` with args `{call['args']}`")
                        # Pass patient context if needed (e.g., for interaction check)
                        if call['name'] == 'check_drug_interactions':
                           call['args']['current_medications'] = patient_data['medications']['names_only']
                           call['args']['allergies'] = patient_data['allergies']
                        elif call['name'] == 'prescribe_medication':
                             # Pre-flight check: Ensure interaction check was requested *before* this prescribe call
                             interaction_check_requested = any(tc['name'] == 'check_drug_interactions' and tc['args'].get('potential_prescription') == call['args'].get('medication_name') for tc in tool_calls)
                             if not interaction_check_requested:
                                 st.error(f"**Safety Violation:** AI attempted to prescribe '{call['args'].get('medication_name')}' without requesting `check_drug_interactions` first. Prescription blocked.")
                                 tool_results.append({"tool_call_id": call['id'], "name": call['name'], "output": json.dumps({"status":"error", "message": "Interaction check not performed prior to prescription attempt."})})
                                 continue # Skip this tool call

                        result = st.session_state.agent.process_tool_call(call)
                        tool_results.append({"tool_call_id": call['id'], "name": call['name'], "output": result}) # Store result with ID

                        # Display tool result immediately
                        try:
                            result_data = json.loads(result)
                            if result_data.get("status") == "success" or result_data.get("status") == "clear" or result_data.get("status") == "flagged":
                                st.success(f"✅ Action `{call['name']}`: {result_data.get('message')}", icon="✅")
                                if result_data.get("details"): st.caption(f"Details: {result_data.get('details')}")
                            elif result_data.get("status") == "warning":
                                st.warning(f"⚠️ Action `{call['name']}`: {result_data.get('message')}", icon="⚠️")
                                if result_data.get("warnings"):
                                    for warn in result_data["warnings"]: st.caption(f"- {warn}")
                            else:
                                st.error(f"❌ Action `{call['name']}`: {result_data.get('message')}", icon="❌")
                        except json.JSONDecodeError:
                             st.error(f"Tool `{call['name']}` returned non-JSON: {result}") # Fallback for non-JSON results

                st.session_state.tool_call_results = tool_results
                # Optionally: Send results back to LLM for final summary (requires multi-turn agent)
        else:
            st.error("Analysis failed. Please check the input data or try again.")

    # --- Display Analysis Results ---
    if st.session_state.analysis_complete and st.session_state.analysis_result:
        st.divider()
        st.header("📊 Analysis & Recommendations")

        res = st.session_state.analysis_result

        # Layout columns for better readability
        col_assessment, col_plan = st.columns(2)

        with col_assessment:
            st.subheader("📋 Assessment")
            st.write(res.get("assessment", "N/A"))

            st.subheader("🤔 Differential Diagnosis")
            ddx = res.get("differential_diagnosis", [])
            if ddx:
                for item in ddx:
                    likelihood = item.get('likelihood', 'Unknown').capitalize()
                    icon = "🥇" if likelihood=="High" else ("🥈" if likelihood=="Medium" else "🥉")
                    with st.expander(f"{icon} {item.get('diagnosis', 'Unknown Diagnosis')} ({likelihood} Likelihood)", expanded=(likelihood=="High")):
                        st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
            else:
                st.info("No differential diagnosis provided.")

            st.subheader("🚨 Risk Assessment")
            risk = res.get("risk_assessment", {})
            flags = risk.get("identified_red_flags", []) + [f.replace("Red Flag: ", "") for f in st.session_state.red_flags] # Combine AI and initial flags
            if flags:
                st.warning(f"**Identified Red Flags:** {', '.join(flags)}")
            else:
                st.success("No immediate red flags identified by AI in this analysis.")

            if risk.get("immediate_concerns"):
                st.warning(f"**Immediate Concerns:** {', '.join(risk.get('immediate_concerns'))}")
            if risk.get("potential_complications"):
                 st.info(f"**Potential Complications:** {', '.join(risk.get('potential_complications'))}")


        with col_plan:
            st.subheader("📝 Recommended Plan")
            plan = res.get("recommended_plan", {})

            st.markdown("**Investigations:**")
            if plan.get("investigations"):
                 st.markdown("\n".join([f"- {inv}" for inv in plan.get("investigations")]))
            else: st.markdown("_None suggested._")

            st.markdown("**Therapeutics:**")
            if plan.get("therapeutics"):
                 st.markdown("\n".join([f"- {thx}" for thx in plan.get("therapeutics")]))
            else: st.markdown("_None suggested._")

            st.markdown("**Consultations:**")
            if plan.get("consultations"):
                 st.markdown("\n".join([f"- {con}" for con in plan.get("consultations")]))
            else: st.markdown("_None suggested._")

            st.markdown("**Patient Education:**")
            if plan.get("patient_education"):
                 st.markdown("\n".join([f"- {edu}" for edu in plan.get("patient_education")]))
            else: st.markdown("_None specified._")

        # Display Rationale and Interaction Summary below the columns
        st.subheader("🧠 AI Rationale & Checks")
        with st.expander("Show AI Reasoning Summary", expanded=False):
            st.write(res.get("rationale_summary", "No rationale summary provided."))

        interaction_summary = res.get("interaction_check_summary", "")
        if interaction_summary: # Only show if interaction check was relevant/performed
             with st.expander("Drug Interaction Check Summary", expanded=True):
                 st.write(interaction_summary)
                 # Also show detailed results from the tool call itself if available
                 for tool_res in st.session_state.tool_call_results:
                     if tool_res['name'] == 'check_drug_interactions':
                         try:
                            data = json.loads(tool_res['output'])
                            if data.get('warnings'):
                                st.warning("Interaction Details:")
                                for warn in data['warnings']:
                                    st.caption(f"- {warn}")
                            else:
                                st.success("Interaction Details: " + data.get('message', 'Check complete.'))
                         except: pass # Ignore parsing errors here

        # Display raw JSON if needed for debugging
        with st.expander("Show Raw AI Output (JSON)"):
            st.json(res)

        st.divider()
        st.success("Analysis Complete.")

    # Disclaimer
    st.markdown("---")
    st.warning(
        """**Disclaimer:** SynapseAI is an AI assistant for clinical decision support and does not replace professional medical judgment.
        All outputs should be critically reviewed by a qualified healthcare provider before making any clinical decisions.
        Verify all information, especially dosages and interactions, independently."""
    )


if __name__ == "__main__":
    main()