UI.py

import streamlit as st
from engine import process_question
from datetime import datetime

# =========================
# PAGE SETUP
# =========================
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")

# =========================
# SESSION STATE
# =========================
if "messages" not in st.session_state:
    st.session_state.messages = []

if "transcript" not in st.session_state:
    st.session_state.transcript = []

# =========================
# SHOW CHAT HISTORY
# =========================
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"], unsafe_allow_html=True)

# =========================
# 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 backend
    with st.spinner("Thinking..."):
        try:
            result = process_question(user_input)
        except Exception as e:
            result = {
                "status": "error",
                "message": str(e)
            }

    # =========================
    # BUILD RESPONSE
    # =========================
    reply = ""

    if result.get("status") == "ok":

        if result.get("message"):
            reply += f"❗ {result['message']}\n\n"

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

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

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

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

        if result.get("sql"):
            reply += "\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')}"

    # =========================
    # SAVE TRANSCRIPT ✅ FIXED
    # =========================
    st.session_state.transcript.append({
        "timestamp": datetime.utcnow().isoformat(),
        "question": user_input,
        "reply": reply,            # ✅ THIS WAS MISSING
        "sql": result.get("sql"),
        "result_preview": result.get("data", [])[:10],
        "error": result.get("message") if result.get("status") != "ok" else None
    })

    # =========================
    # SHOW ASSISTANT MESSAGE
    # =========================
    st.session_state.messages.append(
        {"role": "assistant", "content": reply}
    )

    with st.chat_message("assistant"):
        st.markdown(reply, unsafe_allow_html=True)

# =========================
# DOWNLOAD TRANSCRIPT
# =========================
def download_transcript_txt():
    lines = []
    for i, entry in enumerate(st.session_state.transcript, 1):
        lines.append(f"\n--- Query {i} ---")
        lines.append(f"Time: {entry['timestamp']}")
        lines.append(f"Question: {entry['question']}")

        if entry.get("reply"):
            lines.append("Reply:")
            lines.append(entry["reply"])

        if entry.get("sql"):
            lines.append(f"SQL:\n{entry['sql']}")

        if entry.get("error"):
            lines.append(f"Error: {entry['error']}")

    return "\n".join(lines)

# Show download button only if data exists
if st.session_state.transcript:
    st.divider()
    st.download_button(
        label="⬇️ Download conversation",
        data=download_transcript_txt(),
        file_name="chat_transcript.txt",
        mime="text/plain",
        use_container_width=True
    )