UI.py

import streamlit as st
from engine import process_question

st.set_page_config(page_title="Hospital AI Assistant", layout="wide")

st.title("🏥 Hospital AI Assistant")
st.caption("Ask questions about patients, conditions, visits, medications, labs")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# Chat input
user_input = st.chat_input("Ask a question about hospital data...")

if user_input:
    # Show user message
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    # Call AI engine directly
    with st.spinner("Thinking..."):
        try:
            result = process_question(user_input)
        except Exception as e:
            result = {"status": "error", "message": str(e)}

    # Build assistant reply
    if result.get("status") == "ok":
        reply = ""

        # Time note (if any)
        if result.get("note"):
            reply += f"🕒 *{result['note']}*\n\n"

        # Data table
        if result.get("data"):
            columns = result.get("columns", [])
            data = result["data"]

            table_md = "| " + " | ".join(columns) + " |\n"
            table_md += "| " + " | ".join(["---"] * len(columns)) + " |\n"

            for row in data[:10]:
                table_md += "| " + " | ".join(str(x) for x in row) + " |\n"

            reply += table_md
        else:
            reply += result.get("message", "No data found.")

        # SQL toggle
        reply += "\n\n---\n"
        reply += "<details><summary><b>Generated SQL</b></summary>\n\n"
        reply += f"```sql\n{result['sql']}\n```"
        reply += "\n</details>"

    else:
        reply = f"❌ {result.get('message', 'Something went wrong')}"

    # Show assistant message
    st.session_state.messages.append({"role": "assistant", "content": reply})
    with st.chat_message("assistant"):
        st.markdown(reply)