components/chatbot.py

"""Reusable chatbot component for all personas.

Renders a styled chat interface with persona-specific context.
Uses st.session_state for message persistence within a session.
"""

import streamlit as st
from services.chat import get_chat_response, get_bot_name, get_greeting


def render_chatbot(
    persona: str,
    resident_name: str = "",
    container_key: str = "",
):
    """Render a full chatbot UI for the given persona.

    Args:
        persona: One of 'caregiver', 'family', 'executive', 'lab_admin'.
        resident_name: Optional resident name for family persona context.
        container_key: Optional unique key suffix to prevent Streamlit widget conflicts.
    """
    bot_name = get_bot_name(persona)
    session_key = f"chat_{persona}_{container_key}" if container_key else f"chat_{persona}"

    # ── Header ──
    icon_map = {
        "caregiver": "🩺",
        "family": "💙",
        "executive": "📊",
        "lab_admin": "🔬",
    }
    icon = icon_map.get(persona, "💬")

    st.markdown(
        f"""<div class="chat-header">
            <span class="chat-header-icon">{icon}</span>
            <span class="chat-header-title">{bot_name}</span>
            <span class="chat-header-badge">AI-Powered</span>
        </div>""",
        unsafe_allow_html=True,
    )

    # ── Initialize messages ──
    if session_key not in st.session_state:
        greeting = get_greeting(persona, resident_name)
        st.session_state[session_key] = [
            {"role": "assistant", "content": greeting}
        ]

    # ── Render message history ──
    chat_container = st.container(height=420)
    with chat_container:
        for msg in st.session_state[session_key]:
            with st.chat_message(
                msg["role"],
                avatar=icon if msg["role"] == "assistant" else None,
            ):
                st.markdown(msg["content"])

    # ── Chat input ──
    family_label = f"Ask about {resident_name}'s care..." if resident_name else "Ask about your loved one's care..."
    placeholder_map = {
        "caregiver": "Ask about vitals, medications, fall risk, handoffs...",
        "family": family_label,
        "executive": "Ask about occupancy, costs, staffing, compliance...",
        "lab_admin": "Ask about TAT, specimens, critical values, infections...",
    }
    placeholder = placeholder_map.get(persona, "Type your question...")

    if prompt := st.chat_input(placeholder, key=f"input_{session_key}"):
        # Add user message
        st.session_state[session_key].append({"role": "user", "content": prompt})

        # Get AI response
        response = get_chat_response(prompt, persona=persona)
        st.session_state[session_key].append({"role": "assistant", "content": response})

        # Rerun to show new messages
        st.rerun()

    # ── Quick suggestion chips ──
    suggestions_map = {
        "caregiver": [
            "Current vitals",
            "Shift handoff",
            "Fall risk",
            "UTI alert status",
            "Sleep analysis",
            "Outstanding tasks",
        ],
        "family": [
            "How did they eat today?",
            "How did they sleep?",
            "How's their mood?",
            "Visiting info",
            "Medications",
        ],
        "executive": [
            "Occupancy trend",
            "Financial impact",
            "Staffing metrics",
            "Fall prevention ROI",
            "Compliance score",
            "Family satisfaction",
        ],
        "lab_admin": [
            "TAT performance",
            "Specimen report",
            "Critical values",
            "Infection rates",
            "Wound AI status",
        ],
    }
    suggestions = suggestions_map.get(persona, [])

    if suggestions:
        st.markdown('<div class="suggestion-chips">', unsafe_allow_html=True)
        cols = st.columns(len(suggestions))
        for i, suggestion in enumerate(suggestions):
            with cols[i]:
                if st.button(
                    suggestion,
                    key=f"chip_{session_key}_{i}",
                    use_container_width=True,
                ):
                    st.session_state[session_key].append(
                        {"role": "user", "content": suggestion}
                    )
                    response = get_chat_response(suggestion, persona=persona)
                    st.session_state[session_key].append(
                        {"role": "assistant", "content": response}
                    )
                    st.rerun()
        st.markdown("</div>", unsafe_allow_html=True)