app.py

import streamlit as st
import json
import re
import os
import traceback
import logging
from dotenv import load_dotenv

# Configure logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# Import agent logic and message types
try:
    from agent import ClinicalAgent, AgentState, check_red_flags
    from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
except ImportError as e:
    logger.exception("Failed to import from agent.py")
    st.error(f"Failed to import from agent.py: {e}. Make sure agent.py is in the same directory.")
    st.stop()

# --- Environment Variable Loading & Validation ---
load_dotenv()
required_keys = ["UMLS_API_KEY", "GROQ_API_KEY", "TAVILY_API_KEY"]
missing = [key for key in required_keys if not os.getenv(key)]
if missing:
    st.error(f"Missing required API Key(s): {', '.join(missing)}. Please set them in environment variables.")
    st.stop()

# --- App Configuration ---
class ClinicalAppSettings:
    APP_TITLE = "SynapseAI"
    PAGE_LAYOUT = "wide"
    MODEL_NAME_DISPLAY = "Llama3-70b (via Groq)"

# Cache the agent to avoid re-initialization on each rerun
@st.cache_resource
def get_agent():
    try:
        return ClinicalAgent()
    except Exception as e:
        logger.exception("Failed to initialize ClinicalAgent")
        st.error(f"Failed to initialize Clinical Agent: {e}. Check API keys and dependencies.")
        st.stop()

# Sidebar patient intake helper
def load_patient_intake():
    st.header("📄 Patient Intake Form")
    # Demographics
    age = st.number_input("Age", min_value=0, max_value=120, value=55, key="sb_age")
    sex = st.selectbox("Sex", ["Male", "Female", "Other"], key="sb_sex")

    # HPI
    chief_complaint = st.text_input("Chief Complaint", "Chest pain", key="sb_cc")
    hpi_details = st.text_area("HPI Details", "55 y/o male...", height=100, key="sb_hpi")
    symptoms = st.multiselect(
        "Symptoms",
        ["Nausea", "Diaphoresis", "SOB", "Dizziness", "Severe Headache", "Syncope", "Hemoptysis"],
        default=["Nausea", "Diaphoresis"],
        key="sb_sym"
    )

    # History
    pmh = st.text_area("PMH", "HTN, HLD, DM2, History of MI", key="sb_pmh")
    psh = st.text_area("PSH", "Appendectomy", key="sb_psh")

    # Meds & Allergies
    current_meds_str = st.text_area(
        "Current Meds",
        "Lisinopril 10mg daily\nMetformin 1000mg BID\nWarfarin 5mg daily",
        key="sb_meds"
    )
    allergies_str = st.text_area("Allergies", "Penicillin (rash), Aspirin", key="sb_allergies")

    # Social/Family
    social_history = st.text_area("SH", "Smoker", key="sb_sh")
    family_history = st.text_area("FHx", "Father MI", key="sb_fhx")

    # Vitals & Exam
    col1, col2 = st.columns(2)
    with col1:
        temp_c = st.number_input("Temp C", min_value=35.0, max_value=42.0, value=36.8, format="%.1f", key="sb_temp")
        hr_bpm = st.number_input("HR", min_value=30, max_value=250, value=95, key="sb_hr")
        rr_rpm = st.number_input("RR", min_value=5, max_value=50, value=18, key="sb_rr")
    with col2:
        bp_mmhg = st.text_input("BP", "155/90", key="sb_bp")
        spo2_percent = st.number_input("SpO2", min_value=70, max_value=100, value=96, key="sb_spo2")
        pain_scale = st.slider("Pain", min_value=0, max_value=10, value=8, key="sb_pain")

    # Updated minimum height to 68px to satisfy Streamlit requirement
    exam_notes = st.text_area("Exam Notes", "Awake, alert...", height=68, key="sb_exam")

    # Process meds and allergies with comprehensions
    current_meds_list = [m.strip() for m in current_meds_str.splitlines() if m.strip()]
    current_med_names_only = [
        m.group(1).lower()
        for med in current_meds_list
        if (m := re.match(r"^\s*([A-Za-z-]+)", med))
    ]
    allergies_list = [
        (m.group(1).strip().lower() if (m := re.match(r"^\s*([A-Za-z\s/-]+)", a.strip())) else a.strip().lower())
        for a in allergies_str.split(",")
        if a.strip()
    ]

    # Parse blood pressure
    bp_sys, bp_dia = None, None
    if "/" in bp_mmhg:
        try:
            bp_sys, bp_dia = map(int, bp_mmhg.split("/"))
        except ValueError:
            logger.warning(f"Unable to parse BP '{bp_mmhg}'")

    return {
        "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_only},
        "allergies": allergies_list,
        "social_history": {"details": social_history},
        "family_history": {"details": family_history},
        "vitals": {
            "temp_c": temp_c,
            "hr_bpm": hr_bpm,
            "bp_mmhg": bp_mmhg,
            "bp_sys": bp_sys,
            "bp_dia": bp_dia,
            "rr_rpm": rr_rpm,
            "spo2_percent": spo2_percent,
            "pain_scale": pain_scale
        },
        "exam_findings": {"notes": exam_notes},
    }

# Main application
def main():
    st.set_page_config(page_title=ClinicalAppSettings.APP_TITLE, layout=ClinicalAppSettings.PAGE_LAYOUT)
    st.title(f"🩺 {ClinicalAppSettings.APP_TITLE}")
    st.caption(f"Interactive Assistant | LangGraph/Groq/Tavily/UMLS/OpenFDA | Model: {ClinicalAppSettings.MODEL_NAME_DISPLAY}")

    # Initialize session state
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "patient_data" not in st.session_state:
        st.session_state.patient_data = None
    if "summary" not in st.session_state:
        st.session_state.summary = None
    if "agent" not in st.session_state:
        st.session_state.agent = get_agent()

    # Sidebar intake
    with st.sidebar:
        patient_data = load_patient_intake()
        if st.button("Start/Update Consultation", key="sb_start"):
            st.session_state.patient_data = patient_data
            red_flags = check_red_flags(patient_data)
            st.sidebar.markdown("---")
            if red_flags:
                st.sidebar.warning("**Initial Red Flags:**")
                for flag in red_flags:
                    st.sidebar.warning(f"- {flag.replace('Red Flag: ', '')}")
            else:
                st.sidebar.success("No immediate red flags.")
            st.session_state.messages = [HumanMessage(content="Initiate consultation. Review patient data and begin analysis.")]
            st.session_state.summary = None
            st.success("Patient data loaded/updated.")
            st.rerun()

    # Chat area
    st.header("💬 Clinical Consultation")
    for msg in st.session_state.messages:
        if isinstance(msg, HumanMessage):
            with st.chat_message("user"):
                st.markdown(msg.content)
        elif isinstance(msg, AIMessage):
            with st.chat_message("assistant"):
                ai_content = msg.content
                structured_output = None
                try:
                    match = re.search(r"```json\s*(\{.*?\})\s*```", ai_content, re.DOTALL | re.IGNORECASE)
                    if match:
                        payload = match.group(1)
                        structured_output = json.loads(payload)
                        prefix = ai_content[:match.start()].strip()
                        suffix = ai_content[match.end():].strip()
                        if prefix:
                            st.markdown(prefix)
                        if suffix:
                            st.markdown(suffix)
                    else:
                        st.markdown(ai_content)
                except (AttributeError, json.JSONDecodeError) as e:
                    logger.warning(f"JSON parse error: {e}")
                    st.markdown(ai_content)

                if structured_output and isinstance(structured_output, dict):
                    st.divider()
                    # Display structured JSON sections
                    cols = st.columns(2)
                    with cols[0]:
                        st.markdown("**Assessment:**")
                        st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
                        st.markdown("**Differential Diagnosis:**")
                        ddx = structured_output.get('differential_diagnosis', [])
                        if ddx:
                            for item in ddx:
                                likelihood = item.get('likelihood', 'Low')
                                icon = '🥇' if likelihood == 'High' else ('🥈' if likelihood == 'Medium' else '🥉')
                                with st.expander(f"{icon} {item.get('diagnosis', 'Unknown')} ({likelihood})"):
                                    st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
                        else:
                            st.info("No DDx provided.")

                        st.markdown("**Risk Assessment:**")
                        risk = structured_output.get('risk_assessment', {})
                        for key, style in [('identified_red_flags', st.warning), ('immediate_concerns', st.warning), ('potential_complications', st.info)]:
                            items = risk.get(key, [])
                            if items:
                                style(f"**{key.replace('_', ' ').capitalize()}:** {', '.join(items)}")
                        if not any(risk.get(k) for k in ['identified_red_flags', 'immediate_concerns', 'potential_complications']):
                            st.success("No specific risks highlighted.")

                    with cols[1]:
                        st.markdown("**Recommended Plan:**")
                        plan = structured_output.get('recommended_plan', {})
                        for section in ["investigations","therapeutics","consultations","patient_education"]:
                            st.markdown(f"_{section.replace('_',' ').capitalize()}:_")
                            items = plan.get(section)
                            if isinstance(items, list):
                                for it in items:
                                    st.markdown(f"- {it}")
                            elif items:
                                st.markdown(f"- {items}")
                            else:
                                st.markdown("_None_")

                    st.markdown("**Rationale & Guideline Check:**")
                    st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}")
                    if interaction := structured_output.get('interaction_check_summary'):
                        st.markdown("**Interaction Check Summary:**")
                        st.markdown(f"> {interaction}")
                    st.divider()

        elif isinstance(msg, ToolMessage):
            tool_name = getattr(msg, 'name', 'tool_execution')
            with st.chat_message(tool_name, avatar="🛠️"):
                try:
                    data = json.loads(msg.content)
                    status = data.get('status', 'info')
                    message = data.get('message', msg.content)
                    if tool_name == "flag_risk" and status == "flagged":
                        st.error(f"🚨 **RISK FLAGGED:** {message}")
                    elif status in ("success", "clear"):
                        st.success(message)
                    elif status == "warning":
                        st.warning(message)
                    else:
                        st.error(message)
                    if details := data.get('details'):
                        st.caption(f"Details: {details}")
                except json.JSONDecodeError:
                    st.info(msg.content)

    # --- Chat Input ---
    if prompt := st.chat_input("Your message or follow-up query..."):
        if not st.session_state.patient_data:
            st.warning("Please load patient data first.")
            st.stop()
        user_msg = HumanMessage(content=prompt)
        st.session_state.messages.append(user_msg)
        with st.chat_message("user"):
            st.markdown(prompt)
        current_state = {
            "messages": st.session_state.messages,
            "patient_data": st.session_state.patient_data,
            "summary": st.session_state.summary,
            "interaction_warnings": None
        }
        with st.spinner("SynapseAI is processing..."):
            try:
                final_state = st.session_state.agent.invoke_turn(current_state)
                st.session_state.messages = final_state.get('messages', [])
                st.session_state.summary = final_state.get('summary')
            except Exception as e:
                logger.exception("Error during agent.invoke_turn")
                st.error(f"Error: {e}")
                st.session_state.messages.append(AIMessage(content=f"Error processing request: {e}"))
        st.rerun()

    # Disclaimer
    st.markdown("---")
    st.warning("**Disclaimer:** SynapseAI is for demonstration only and not for clinical use.")

if __name__ == "__main__":
    main()