src/streamlit_app.py

"""
src/streamlit_app.py — PharmGPT main Streamlit application.

Features
────────
- Multi-turn chatbot powered by an Ollama text model (OpenAI-compatible API)
- Demo login with 3 user accounts and 1 admin account
- Top-right credential panel for easy demo exploration
- Persistent chat history stored in SQLite (data/pharmgpt.db)
- Rotating application logs (logs/pharmgpt.log)
- Admin-only view: browse full per-user chat history

Run locally
───────────
  streamlit run src/streamlit_app.py

Docker / Hugging Face Spaces
─────────────────────────────
  See Dockerfile and README.md for deployment instructions.
  Set MODEL_NAME, MODEL_BASE_URL, OLLAMA_API_KEY as Space secrets.
"""

# ── Make sure project root is on sys.path so "from src.xxx" works ─────────────
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import logging
import uuid

import streamlit as st

from src.agents import PharmGPTAgent
from src.auth import DEMO_CREDENTIALS, User, authenticate, is_admin
from src.config import APP_SUBTITLE, APP_TITLE, MODEL_NAME
from src.storage import (
    get_all_users_summary,
    get_platform_stats,
    get_session_messages,
    get_user_history,
    get_user_sessions,
    init_db,
    save_message,
    setup_logging,
)

# ── Bootstrap ─────────────────────────────────────────────────────────────────
setup_logging()
init_db()
logger = logging.getLogger(__name__)

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title=APP_TITLE,
    page_icon="💊",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ── CSS tweaks ...────────────────────────────────────────────────────────────────
st.markdown(
    """
    <style>
    /* Hide Streamlit top bar (Deploy button, 3-dot menu, running indicator) */
    header[data-testid="stHeader"] { display: none !important; }
    #MainMenu { visibility: hidden !important; }
    .stDeployButton { display: none !important; }
    /* Compensate for removed header height */
    .block-container { padding-top: 1.5rem !important; }
    /* Credential badge font */
    .cred-badge { font-family: monospace; font-size: 0.85rem; }
    </style>
    """,
    unsafe_allow_html=True,
)


# ─────────────────────────────────────────────────────────────────────────────
# Session-state initialisation
# ─────────────────────────────────────────────────────────────────────────────
def _init_state() -> None:
    defaults: dict = {
        "authenticated": False,
        "user": None,           # auth.User object
        "session_id": None,     # UUID string, reset on "New Chat"
        "messages": [],         # list of {"role": str, "content": str}
    }
    for key, val in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = val

_init_state()


# ─────────────────────────────────────────────────────────────────────────────
# Component: top-right demo credentials panel
# ─────────────────────────────────────────────────────────────────────────────
def render_credentials_panel() -> None:
    with st.expander("🔑 Demo Credentials", expanded=False):
        st.caption("⚠️ For demo / review use only — not for production.")
        st.markdown("| Role | Username | Password |")
        st.markdown("|------|----------|----------|")
        for c in DEMO_CREDENTIALS:
            role_badge = "🛡️ Admin" if c["role"] == "Admin" else "👤 User"
            st.markdown(f"| {role_badge} | `{c['username']}` | `{c['password']}` |")


# ─────────────────────────────────────────────────────────────────────────────
# Component: login form (centred, shown when not authenticated)
# ─────────────────────────────────────────────────────────────────────────────
def render_login() -> None:
    _, mid, _ = st.columns([1, 2, 1])
    with mid:
        st.markdown("## 🔐 Login")
        st.caption("Enter your credentials to start chatting with PharmGPT.")
        username = st.text_input("Username", key="login_username", placeholder="e.g. alice")
        password = st.text_input("Password", type="password", key="login_password")
        if st.button("Login", use_container_width=True, type="primary"):
            user = authenticate(username, password)
            if user:
                st.session_state.authenticated = True
                st.session_state.user = user
                st.session_state.session_id = str(uuid.uuid4())
                st.session_state.messages = []
                logger.info("LOGIN | user=%s role=%s", user.username, user.role)
                st.rerun()
            else:
                st.error("Invalid username or password.")


# ─────────────────────────────────────────────────────────────────────────────
# Component: sidebar (info panel — shown when authenticated)
# ─────────────────────────────────────────────────────────────────────────────
def render_sidebar() -> None:
    user: User = st.session_state.user
    with st.sidebar:
        st.markdown(f"## 💊 {APP_TITLE}")
        st.caption(APP_SUBTITLE)
        st.divider()
        role_label = "🛡️ Admin" if is_admin(user) else "👤 User"
        st.markdown(f"**{role_label}:** {user.display_name}")
        st.markdown(f"**Model:** `{MODEL_NAME}`")
        st.divider()
        st.caption("💊 PharmGPT · Educational use only.")
        st.caption("ℹ️ Information here is general and does not substitute professional medical advice.")


# ─────────────────────────────────────────────────────────────────────────────
# Component: top action bar (always visible when authenticated)
# ─────────────────────────────────────────────────────────────────────────────
def render_action_bar() -> None:
    """Persistent toolbar row: user info + New Chat + Logout — always in view."""
    user: User = st.session_state.user
    role_label = "🛡️ Admin" if is_admin(user) else "👤 User"
    info_col, _, btn1_col, btn2_col = st.columns([5, 1, 1, 1])
    with info_col:
        st.markdown(
            f"👋 **{user.display_name}** &nbsp;|&nbsp; {role_label} &nbsp;|&nbsp; "
            f"<span style='font-family:monospace;font-size:0.85rem'>model: {MODEL_NAME}</span>",
            unsafe_allow_html=True,
        )
    with btn1_col:
        if st.button("🗒️ New Chat", use_container_width=True):
            st.session_state.session_id = str(uuid.uuid4())
            st.session_state.messages = []
            logger.info("NEW_CHAT | user=%s session=%s",
                        user.username, st.session_state.session_id)
            st.rerun()
    with btn2_col:
        if st.button("🚪 Logout", use_container_width=True):
            logger.info("LOGOUT | user=%s", user.username)
            st.session_state.authenticated = False
            st.session_state.user = None
            st.session_state.session_id = None
            st.session_state.messages = []
            st.rerun()


# ─────────────────────────────────────────────────────────────────────────────
# Component: main chat interface
# ─────────────────────────────────────────────────────────────────────────────
def render_chat() -> None:
    user: User = st.session_state.user
    agent = PharmGPTAgent()

    # ── Replay all existing messages in this session ───────────────────────
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    # ── Handle new input ───────────────────────────────────────────────────
    prompt = st.chat_input("Ask about medications, symptoms, diseases…")
    if not prompt:
        return

    # Display user message immediately
    with st.chat_message("user"):
        st.markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Persist user turn
    save_message(
        session_id=st.session_state.session_id,
        user_id=user.user_id,
        username=user.username,
        user_role=user.role,
        msg_role="user",
        content=prompt,
        model_name=MODEL_NAME,
    )

    # Build conversation context (all turns before the current one)
    history = [
        {"role": m["role"], "content": m["content"]}
        for m in st.session_state.messages[:-1]
    ]

    # Get and display assistant response with streaming
    with st.chat_message("assistant"):
        full_response = st.write_stream(
            agent.stream(prompt, conversation_history=history)
        )
    status = agent.last_status

    st.session_state.messages.append({"role": "assistant", "content": full_response})

    # Persist assistant turn
    error_info = full_response if status == "error" else ""
    save_message(
        session_id=st.session_state.session_id,
        user_id=user.user_id,
        username=user.username,
        user_role=user.role,
        msg_role="assistant",
        content=full_response,
        model_name=MODEL_NAME,
        status=status,
        error_info=error_info,
    )
    
    logger.info(
        "CHAT | user=%s status=%s chars=%d",
        user.username, status, len(full_response),
    )


# ─────────────────────────────────────────────────────────────────────────────
# Component: admin — all-users history viewer
# ─────────────────────────────────────────────────────────────────────────────
def render_admin_history() -> None:
    st.subheader("🗂️ Admin — User Chat Audit")

    # ── Platform KPIs ───────────────────────────────────────────────────────
    stats = get_platform_stats()
    k1, k2, k3, k4, k5, k6 = st.columns(6)
    k1.metric("👥 Total Users",     stats.get("total_users", 0))
    k2.metric("💬 Total Sessions",  stats.get("total_sessions", 0))
    k3.metric("✉️ Total Messages",  stats.get("total_messages", 0))
    k4.metric("📝 User Messages",   stats.get("user_messages", 0))
    k5.metric("⚠️ Errors",          stats.get("errors", 0))
    k6.metric("🟢 Active Today",    stats.get("active_today", 0))

    last_ts = stats.get("last_activity", "")
    if last_ts:
        st.caption(f"Last activity: {last_ts[:19].replace('T', ' ')} UTC")

    st.divider()

    # ── User list check ───────────────────────────────────────────────────────
    all_users = get_all_users_summary()
    if not all_users:
        st.info("No chat history yet. Users will appear here once they start chatting.")
        return

    # ── Step 1: Pick a user ───────────────────────────────────────────────────
    user_labels = [
        f"{u['username']}  —  {u['message_count']} msgs · "
        f"{u['session_count']} sessions · last: {u['last_active'][:16].replace('T', ' ')}"
        for u in all_users
    ]
    filter_col, _ = st.columns([2, 3])
    with filter_col:
        selected_user_idx = st.selectbox(
            "👤 Select user",
            range(len(user_labels)),
            format_func=lambda i: user_labels[i],
            key="admin_user_select",
        )
    selected_user = all_users[selected_user_idx]

    # Inline user stat pills
    uc1, uc2, uc3, uc4 = st.columns(4)
    uc1.metric("Messages",   selected_user["message_count"])
    uc2.metric("Sessions",   selected_user["session_count"])
    uc3.metric("Role",       selected_user["user_role"].capitalize())
    uc4.metric("Last Active", selected_user["last_active"][:10])

    st.divider()

    # ── Step 2: Pick a session ───────────────────────────────────────────────
    sessions = get_user_sessions(selected_user["user_id"])
    if not sessions:
        st.info("No sessions found for this user.")
        return

    # "All sessions" option prepended
    ALL = "__all__"
    session_labels = [f"📂 All sessions  ({len(sessions)} total)"]
    session_ids    = [ALL]
    for s in sessions:
        started  = s["started_at"][:16].replace("T", " ")
        last_msg = s["last_msg_at"][:16].replace("T", " ")
        flag     = " ⚠️" if s["flagged"] else ""
        session_labels.append(
            f"💬 {started} → {last_msg}  ·  {s['message_count']} msgs{flag}  "
            f"[{s['session_id'][:8]}…]"
        )
        session_ids.append(s["session_id"])

    sess_col, _ = st.columns([3, 2])
    with sess_col:
        selected_sess_idx = st.selectbox(
            "💬 Select session",
            range(len(session_labels)),
            format_func=lambda i: session_labels[i],
            key="admin_session_select",
        )
    selected_sid = session_ids[selected_sess_idx]

    st.divider()

    # ── Step 3: Display messages ─────────────────────────────────────────────
    if selected_sid == ALL:
        msgs = get_user_history(selected_user["user_id"], limit=300)
        # newest-first → reverse for chronological display
        msgs = list(reversed(msgs))
        st.caption(
            f"Showing all messages for **{selected_user['username']}** "
            f"across {len(sessions)} session(s) — chronological order."
        )
    else:
        msgs = get_session_messages(selected_sid)
        st.caption(
            f"Session `{selected_sid[:8]}…` — "
            f"{len(msgs)} messages for **{selected_user['username']}**."
        )

    if not msgs:
        st.info("No messages in this view.")
        return

    # Render as a clean chat transcript, grouped by session
    current_session = None
    for msg in msgs:
        # Session divider when browsing all sessions
        if selected_sid == ALL and msg["session_id"] != current_session:
            current_session = msg["session_id"]
            ts_start = msg["timestamp"][:16].replace("T", " ")
            st.markdown(
                f"<div style='text-align:center; color:gray; font-size:0.8rem; "
                f"margin:1rem 0 0.4rem;'>─── Session "
                f"<code>{current_session[:8]}…</code> started {ts_start} UTC ───</div>",
                unsafe_allow_html=True,
            )

        role      = msg["msg_role"]
        ts        = msg["timestamp"][11:16]  # HH:MM
        status    = msg.get("status", "ok")
        flag_icon = "🚨 " if status == "blocked_emergency" else "🚫 " if status == "blocked_scope" else "⚠️ " if status == "error" else ""

        with st.chat_message(role):
            if flag_icon:
                st.caption(f"{flag_icon}`{status}` · {ts} UTC")
            else:
                st.caption(f"{ts} UTC")
            st.markdown(msg["content"])


# ─────────────────────────────────────────────────────────────────────────────
# Main layout
# ─────────────────────────────────────────────────────────────────────────────
def main() -> None:
    # ── Top bar: title (left) + credentials panel (right) ─────────────────
    title_col, cred_col = st.columns([5, 2])
    with title_col:
        st.markdown(f"# 💊 {APP_TITLE}")
        st.caption(APP_SUBTITLE)
    with cred_col:
        render_credentials_panel()

    st.divider()

    # ── Not logged in → show login form ───────────────────────────────────
    if not st.session_state.authenticated:
        render_login()
        return

    # ── Logged in → sidebar + action bar + main content ──────────────────
    render_sidebar()
    render_action_bar()
    st.divider()
    user: User = st.session_state.user

    if is_admin(user):
        chat_tab, history_tab = st.tabs(["💬 My Chat", "🗂️ All Users History"])
        with chat_tab:
            render_chat()
        with history_tab:
            render_admin_history()
    else:
        render_chat()


main()