Update app.py

app.py

@@ -1,359 +1,795 @@
 """
-DocuMind - Streamlit Frontend for HuggingFace Spaces
-Updated visually to match the specific dark-sidebar & green-chat UI mockup.
+DocuMind AI - Minimal Enterprise RAG Chat UI
+Single-file Streamlit app matching a clean document-intelligence interface.
+
+Features:
+- Minimal light UI inspired by the provided reference image
+- Left sidebar with New Chat + Upload
+- Native Streamlit sidebar open/close toggle (top-left)
+- Chat history with user/assistant bubbles
+- Optional source citations from backend
+- Works with an existing FastAPI backend
 """
 
-import streamlit as st
+from __future__ import annotations
+
+import html
 import json
 import time
+from datetime import datetime
+
 import requests
+import streamlit as st
+
 
 # ============================================================================
-# STREAMLIT CONFIG (Must be first)
+# PAGE CONFIG
 # ============================================================================
-
 st.set_page_config(
-    page_title="DocuMind AI - Digital Curator",
+    page_title="DocuMind AI",
     page_icon="🧠",
     layout="wide",
-    initial_sidebar_state="expanded"
+    initial_sidebar_state="expanded",
 )
 
+
 # ============================================================================
-# CUSTOM CSS INJECTION (Theming to match the new mockup)
+# THEME + MINIMAL CSS
 # ============================================================================
-custom_css = """
-<style>
-    /* Global Font and Backgrounds */
-    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
-    
-    html, body, [class*="css"]  {
-        font-family: 'Inter', sans-serif;
-    }
-    
-    /* Main App Background */
-    .stApp {
-        background-color: #F8F9FA;
-    }
-    
-    /* Sidebar Background & Styling (Dark Theme) */
+st.markdown(
+    """
+    <style>
+    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');
+
+    :root {
+        --bg: #f6f8fb;
+        --panel: #ffffff;
+        --panel-soft: #f2f5f8;
+        --panel-dark: #0f4c5c;
+        --panel-dark-2: #0b3d49;
+        --text: #0f172a;
+        --muted: #667085;
+        --muted-2: #94a3b8;
+        --border: rgba(15, 23, 42, 0.08);
+        --shadow: 0 8px 30px rgba(15, 23, 42, 0.06);
+        --radius-xl: 28px;
+        --radius-lg: 22px;
+        --radius-md: 16px;
+        --radius-sm: 12px;
+        --accent: #0f4c5c;
+        --accent-soft: rgba(15, 76, 92, 0.10);
+    }
+
+    html, body, [data-testid="stAppViewContainer"] {
+        background: var(--bg) !important;
+        color: var(--text) !important;
+        font-family: 'Inter', sans-serif !important;
+    }
+
+    *, *::before, *::after { box-sizing: border-box; }
+
+    /* Hide Streamlit chrome */
+    #MainMenu, footer, header, [data-testid="stToolbar"], [data-testid="stDecoration"] {
+        display: none !important;
+    }
+
+    /* Sidebar */
     [data-testid="stSidebar"] {
-        background-color: #172430 !important;
-        border-right: none;
-    }
-    [data-testid="stSidebar"] * {
-        color: #FFFFFF !important;
-    }
-    
-    /* Hide default header and footer */
-    header {visibility: hidden;}
-    footer {visibility: hidden;}
-    
-    /* Sidebar Headers & Text */
-    .sidebar-title {
-        font-size: 1.25rem;
-        font-weight: 700;
-        color: #FFFFFF;
-        margin-bottom: 2px;
+        background: linear-gradient(180deg, #f8fafc 0%, #f5f7fa 100%) !important;
+        border-right: 1px solid var(--border) !important;
     }
-    .sidebar-subtitle {
-        font-size: 0.8rem;
-        color: #8A9CA8;
-        margin-bottom: 25px;
-    }
-    .recent-header {
-        font-size: 0.75rem;
-        font-weight: 600;
-        color: #8A9CA8;
-        margin-top: 20px;
-        margin-bottom: 15px;
-        letter-spacing: 1px;
-    }
-    .recent-item {
-        color: #D1D8DD;
-        font-size: 0.95rem;
-        padding: 8px 0;
-        display: flex;
-        align-items: center;
-        gap: 10px;
-        cursor: pointer;
+
+    [data-testid="stSidebarContent"] {
+        padding-top: 0 !important;
     }
-    
-    /* Upload Box override for Dark Sidebar */
+
+    /* Main container */
+    .main .block-container {
+        padding: 0 !important;
+        max-width: 100% !important;
+    }
+
+    /* Sidebar collapse button */
+    [data-testid="collapsedControl"] {
+        color: #0f172a !important;
+        background: rgba(255,255,255,0.75) !important;
+        border: 1px solid var(--border) !important;
+        border-radius: 999px !important;
+        box-shadow: var(--shadow) !important;
+    }
+
+    /* Buttons */
+    .stButton > button {
+        border-radius: 18px !important;
+        border: 1px solid var(--border) !important;
+        background: var(--panel) !important;
+        color: var(--text) !important;
+        font-family: 'Inter', sans-serif !important;
+        font-weight: 600 !important;
+        transition: 0.18s ease !important;
+        box-shadow: none !important;
+        width: 100% !important;
+        padding: 0.8rem 1rem !important;
+    }
+    .stButton > button:hover {
+        transform: translateY(-1px) !important;
+        border-color: rgba(15, 76, 92, 0.18) !important;
+        background: rgba(15, 76, 92, 0.04) !important;
+    }
+
+    /* Sidebar file uploader */
     [data-testid="stFileUploader"] {
-        border: 1px solid #2B3B4A;
-        border-radius: 12px;
-        padding: 15px 10px;
-        background-color: transparent;
-        text-align: center;
+        background: var(--panel) !important;
+        border: 1px dashed rgba(15, 23, 42, 0.10) !important;
+        border-radius: 18px !important;
+        padding: 10px !important;
+        box-shadow: none !important;
     }
-    [data-testid="stFileUploader"] > div > div {
-        background-color: transparent;
+    [data-testid="stFileUploader"] label,
+    [data-testid="stFileUploadDropzone"] p,
+    [data-testid="stFileUploadDropzone"] span {
+        color: var(--muted) !important;
+        font-family: 'Inter', sans-serif !important;
     }
-    
-    /* Chat bubbles */
-    [data-testid="stChatMessage"] {
-        border-radius: 12px;
-        padding: 15px 20px;
-        margin-bottom: 10px;
-        max-width: 85%;
+
+    /* Chat input */
+    [data-testid="stChatInput"] {
+        background: transparent !important;
     }
-    /* Assistant Bubble */
-    [data-testid="stChatMessage"]:has([data-testid="assistantAvatar"]) {
-        background-color: #EEF0F2;
-        color: #111827;
-        margin-right: auto;
+    [data-testid="stChatInput"] textarea {
+        background: var(--panel) !important;
+        border-radius: 24px !important;
+        border: 1px solid rgba(15, 23, 42, 0.08) !important;
+        box-shadow: var(--shadow) !important;
+        font-family: 'Inter', sans-serif !important;
+        padding: 1rem 1rem 1rem 1rem !important;
+        color: var(--text) !important;
     }
-    /* User Bubble */
-    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) {
-        background-color: #0C4A42;
-        color: white !important;
-        margin-left: auto;
+    [data-testid="stChatInput"] textarea::placeholder {
+        color: var(--muted-2) !important;
     }
-    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) p {
-        color: white !important;
+    [data-testid="stChatInput"] button {
+        background: var(--accent) !important;
+        color: #fff !important;
+        border-radius: 999px !important;
     }
 
-    /* Welcome Header */
-    .welcome-header {
-        text-align: center;
-        color: #111827;
-        font-size: 2.5rem;
-        font-weight: 700;
-        margin-top: 60px;
-        margin-bottom: 15px;
+    /* Spacers */
+    hr {
+        border: none !important;
+        border-top: 1px solid rgba(15, 23, 42, 0.06) !important;
+        margin: 0 !important;
     }
-    .welcome-sub {
+
+    /* Scrollbar */
+    ::-webkit-scrollbar { width: 8px; height: 8px; }
+    ::-webkit-scrollbar-track { background: transparent; }
+    ::-webkit-scrollbar-thumb { background: rgba(100, 116, 139, 0.25); border-radius: 999px; }
+    ::-webkit-scrollbar-thumb:hover { background: rgba(100, 116, 139, 0.45); }
+
+    /* Hero */
+    .hero-wrap {
+        display: flex;
+        flex-direction: column;
+        align-items: center;
+        justify-content: center;
         text-align: center;
-        color: #4B5563;
+        padding: 5rem 1.25rem 2rem;
+        min-height: 48vh;
+    }
+    .hero-title {
+        font-size: clamp(2rem, 3.5vw, 3.5rem);
+        line-height: 1.08;
+        font-weight: 800;
+        letter-spacing: -0.04em;
+        color: #0f172a;
+        margin: 0 0 1rem 0;
+    }
+    .hero-subtitle {
+        max-width: 760px;
+        font-size: 1.02rem;
+        line-height: 1.65;
+        color: var(--muted);
+        margin: 0;
+    }
+
+    /* Session header */
+    .session-row {
+        display: flex;
+        align-items: center;
+        gap: 0.75rem;
+        padding: 1rem 1.25rem 0.85rem 1.25rem;
+    }
+    .session-label {
+        font-size: 0.95rem;
+        font-weight: 800;
+        letter-spacing: 0.02em;
+        color: #0f4c5c;
+        text-transform: uppercase;
+    }
+    .session-badge {
+        display: inline-flex;
+        align-items: center;
+        padding: 0.28rem 0.65rem;
+        border-radius: 999px;
+        background: rgba(15, 76, 92, 0.10);
+        color: #0f4c5c;
+        font-size: 0.7rem;
+        font-weight: 800;
+        letter-spacing: 0.08em;
+    }
+
+    /* Sidebar content blocks */
+    .sidebar-brand {
+        padding: 1.1rem 1rem 1rem;
+        border-bottom: 1px solid rgba(15, 23, 42, 0.06);
+    }
+    .brand-row {
+        display: flex;
+        align-items: center;
+        gap: 0.75rem;
+    }
+    .brand-icon {
+        width: 44px;
+        height: 44px;
+        border-radius: 14px;
+        background: linear-gradient(135deg, #0f4c5c 0%, #116579 100%);
+        color: white;
+        display: flex;
+        align-items: center;
+        justify-content: center;
         font-size: 1.1rem;
-        margin-bottom: 50px;
-        line-height: 1.5;
+        box-shadow: var(--shadow);
+    }
+    .brand-title {
+        font-size: 1.05rem;
+        font-weight: 800;
+        color: #0f172a;
+        line-height: 1.1;
+    }
+    .brand-subtitle {
+        font-size: 0.78rem;
+        color: var(--muted-2);
+        margin-top: 0.15rem;
+    }
+
+    .sidebar-section {
+        padding: 1rem;
+    }
+    .sidebar-heading {
+        font-size: 0.72rem;
+        font-weight: 800;
+        letter-spacing: 0.12em;
+        text-transform: uppercase;
+        color: #64748b;
+        margin: 0 0 0.6rem 0;
+    }
+    .recents-item {
+        display: flex;
+        align-items: center;
+        gap: 0.6rem;
+        padding: 0.65rem 0.2rem;
+        color: #334155;
+        font-size: 0.92rem;
+        border-radius: 12px;
+    }
+    .recents-item span.icon {
+        color: #64748b;
+        font-size: 0.85rem;
     }
 
-    /* Top Nav Badge */
-    .top-badge-container {
+    /* Message bubbles */
+    .msg-user-wrap {
         display: flex;
         justify-content: flex-end;
-        padding: 10px 20px 0 0;
-    }
-    .status-badge {
-        background-color: #E2E8F0;
-        color: #475569;
-        padding: 6px 12px;
-        border-radius: 6px;
-        font-size: 0.75rem;
+        margin: 0.8rem 0 0.25rem;
+    }
+    .msg-user {
+        max-width: min(720px, 66%);
+        background: linear-gradient(180deg, #0f4c5c 0%, #0b3d49 100%);
+        color: #fff;
+        border-radius: 22px 22px 6px 22px;
+        padding: 1rem 1.1rem;
+        box-shadow: var(--shadow);
+        font-size: 0.98rem;
+        line-height: 1.65;
+        border: 1px solid rgba(255,255,255,0.08);
+        white-space: pre-wrap;
+        word-break: break-word;
+    }
+    .msg-bot-row {
+        display: flex;
+        gap: 0.75rem;
+        align-items: flex-start;
+        margin: 1rem 0 0.3rem;
+    }
+    .bot-avatar {
+        width: 36px;
+        height: 36px;
+        border-radius: 12px;
+        background: linear-gradient(135deg, #0f4c5c 0%, #116579 100%);
+        color: #fff;
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        font-size: 1rem;
+        flex-shrink: 0;
+        box-shadow: var(--shadow);
+    }
+    .msg-bot {
+        max-width: min(820px, 72%);
+        background: #eef2f6;
+        color: #0f172a;
+        border: 1px solid rgba(15, 23, 42, 0.06);
+        border-radius: 8px 22px 22px 22px;
+        padding: 1rem 1.1rem;
+        font-size: 0.98rem;
+        line-height: 1.7;
+        white-space: normal;
+        word-break: break-word;
+        box-shadow: 0 2px 10px rgba(15, 23, 42, 0.03);
+    }
+    .msg-time {
+        font-size: 0.72rem;
+        color: #94a3b8;
+        margin-top: 0.35rem;
+    }
+    .typing-cursor {
+        opacity: 0.85;
+    }
+
+    /* Source cards */
+    .source-card {
+        background: #fff;
+        border: 1px solid rgba(15, 23, 42, 0.07);
+        border-radius: 16px;
+        padding: 0.9rem 1rem;
+        margin: 0.45rem 0;
+        box-shadow: 0 2px 10px rgba(15, 23, 42, 0.03);
+    }
+    .source-head {
+        display: flex;
+        justify-content: space-between;
+        align-items: center;
+        margin-bottom: 0.35rem;
+        font-size: 0.8rem;
         font-weight: 700;
-        letter-spacing: 0.5px;
     }
-</style>
-"""
-st.markdown(custom_css, unsafe_allow_html=True)
+    .source-title { color: #0f4c5c; }
+    .source-score {
+        color: #64748b;
+        background: #f1f5f9;
+        padding: 0.1rem 0.45rem;
+        border-radius: 999px;
+        font-weight: 700;
+    }
+    .source-file {
+        font-size: 0.78rem;
+        color: #64748b;
+        margin-bottom: 0.4rem;
+    }
+    .source-snippet {
+        font-size: 0.86rem;
+        line-height: 1.55;
+        color: #334155;
+        border-left: 2px solid rgba(15, 76, 92, 0.18);
+        padding-left: 0.7rem;
+    }
+
+    /* Bottom spacing so chat input doesn't overlap content */
+    .bottom-space {
+        height: 1.75rem;
+    }
+    </style>
+    """,
+    unsafe_allow_html=True,
+)
 
 
 # ============================================================================
-# API CONFIGURATION & UTILITIES
+# CONFIG
 # ============================================================================
-
 API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 30  # seconds
+API_TIMEOUT = 60
 
-def check_backend_health(retries=5):
-    """Check if backend API is running"""
-    for attempt in range(retries):
-        try:
-            response = requests.get(f"{API_URL}/docs", timeout=5)
-            if response.status_code == 200:
-                return True
-        except:
-            if attempt < retries - 1:
-                time.sleep(1)
-    return False
-
-def safe_api_call(method, endpoint, **kwargs):
-    """Safely call API with error handling"""
+
+# ============================================================================
+# SESSION STATE
+# ============================================================================
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "backend_ok" not in st.session_state:
+    st.session_state.backend_ok = False
+if "backend_checked" not in st.session_state:
+    st.session_state.backend_checked = False
+if "chat_sessions" not in st.session_state:
+    st.session_state.chat_sessions = [
+        "Q3 Financial Report",
+        "Product Roadmap 2024",
+        "Research Summary",
+    ]
+if "ingested_docs" not in st.session_state:
+    st.session_state.ingested_docs = []
+if "show_sources" not in st.session_state:
+    st.session_state.show_sources = True
+
+
+# ============================================================================
+# HELPERS
+# ============================================================================
+def check_backend() -> bool:
     try:
+        r = requests.get(f"{API_URL}/docs", timeout=5)
+        return r.status_code == 200
+    except Exception:
+        return False
+
+
+def api_call(method: str, endpoint: str, **kwargs):
+    try:
+        kwargs.setdefault("timeout", API_TIMEOUT)
         url = f"{API_URL}{endpoint}"
-        kwargs.setdefault('timeout', API_TIMEOUT)
-        
-        if method == "GET":
-            response = requests.get(url, **kwargs)
-        elif method == "POST":
-            response = requests.post(url, **kwargs)
-        else:
-            return None, "Invalid method"
-        
-        return response, None
+        fn = requests.get if method.upper() == "GET" else requests.post
+        return fn(url, **kwargs), None
     except requests.exceptions.Timeout:
         return None, "⏱️ Request timed out"
     except requests.exceptions.ConnectionError:
         return None, "⚠️ Backend not responding"
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        return None, f"Error: {e}"
 
-# Initialize Session States
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-if "backend_checked" not in st.session_state:
-    st.session_state.backend_checked = False
-    st.session_state.backend_healthy = False
 
-# Health check on load
+def now_time() -> str:
+    return datetime.now().strftime("%-I:%M %p")
+
+
+def esc(text: str) -> str:
+    return html.escape(text or "")
+
+
+def render_text_block(text: str) -> str:
+    """Convert plain text to safe HTML for display in a bubble."""
+    safe = esc(text).replace("\n", "<br>")
+    return safe
+
+
+def render_assistant_text(text: str) -> str:
+    """Render assistant content with lightweight bullet styling."""
+    lines = (text or "").splitlines()
+    out = []
+    for raw in lines:
+        line = raw.strip()
+        if not line:
+            continue
+
+        if line.startswith(("- ", "• ", "* ")):
+            item = esc(line[2:])
+            if ":" in item:
+                key, val = item.split(":", 1)
+                out.append(
+                    f"<div style='display:flex;gap:0.6rem;align-items:flex-start;margin:0.35rem 0;'>"
+                    f"<span style='color:#0f4c5c;font-weight:900;line-height:1.2;'>✓</span>"
+                    f"<span><strong>{key.strip()}:</strong> {val.strip()}</span>"
+                    f"</div>"
+                )
+            else:
+                out.append(
+                    f"<div style='display:flex;gap:0.6rem;align-items:flex-start;margin:0.35rem 0;'>"
+                    f"<span style='color:#0f4c5c;font-weight:900;line-height:1.2;'>✓</span>"
+                    f"<span>{item}</span>"
+                    f"</div>"
+                )
+        else:
+            out.append(f"<div style='margin:0.22rem 0;'>{esc(line)}</div>")
+    return "".join(out)
+
+
+# ============================================================================
+# BOOTSTRAP
+# ============================================================================
 if not st.session_state.backend_checked:
-    st.session_state.backend_healthy = check_backend_health(retries=1)
+    with st.spinner("Starting up..."):
+        time.sleep(1.2)
+        st.session_state.backend_ok = check_backend()
     st.session_state.backend_checked = True
 
+try:
+    resp, _ = api_call("GET", "/sources")
+    if resp and resp.status_code == 200:
+        st.session_state.ingested_docs = resp.json().get("documents", [])
+except Exception:
+    pass
+
+
 # ============================================================================
-# UI: SIDEBAR
+# SIDEBAR
 # ============================================================================
 with st.sidebar:
-    # Logo & Title
-    st.markdown("""
-        <div style="display: flex; align-items: center; gap: 12px; margin-bottom: 10px;">
-            <div style="font-size: 2rem;">◧</div>
-            <div>
-                <div class="sidebar-title">DocuMind AI</div>
-                <div class="sidebar-subtitle">Digital Curator</div>
+    st.markdown(
+        """
+        <div class="sidebar-brand">
+            <div class="brand-row">
+                <div class="brand-icon">🧠</div>
+                <div>
+                    <div class="brand-title">DocuMind AI</div>
+                    <div class="brand-subtitle">Digital Curator</div>
+                </div>
             </div>
         </div>
-    """, unsafe_allow_html=True)
-    
-    # New Chat Button
-    if st.button("＋ New Chat", use_container_width=True, type="primary"):
+        """,
+        unsafe_allow_html=True,
+    )
+
+    st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)
+    if st.button("＋  New Chat", key="new_chat"):
         st.session_state.messages = []
         st.rerun()
-    
-    st.write("") # Spacer
-    
-    # Upload Area
-    uploaded_file = st.file_uploader("Upload Document", type=["txt", "pdf", "docx"], label_visibility="collapsed")
-    # Custom badge for the limit underneath the uploader
-    st.markdown("""
-        <div style="text-align: center; margin-top: -15px; margin-bottom: 20px;">
-            <span style="background-color: #2B3B4A; color: #8A9CA8; font-size: 0.7rem; padding: 3px 8px; border-radius: 12px; font-weight: 600;">200MB LIMIT</span>
-        </div>
-    """, unsafe_allow_html=True)
-    
-    if uploaded_file and st.button("Process Document", use_container_width=True):
-        with st.spinner("Analyzing document..."):
+    st.markdown("</div>", unsafe_allow_html=True)
+
+    st.markdown('<div class="sidebar-section" style="padding-top:0;">', unsafe_allow_html=True)
+    st.markdown('<div class="sidebar-heading">Upload</div>', unsafe_allow_html=True)
+    uploaded_file = st.file_uploader(
+        "Upload document",
+        type=["txt", "pdf", "docx"],
+        label_visibility="collapsed",
+    )
+    if uploaded_file and st.button("Ingest Document", key="ingest_btn"):
+        with st.spinner("Processing..."):
             files = {"file": (uploaded_file.name, uploaded_file.getvalue())}
-            response, error = safe_api_call("POST", "/ingest", files=files)
-            if error:
-                st.error(error)
-            elif response and response.status_code == 200:
-                st.success("✅ Added to knowledge base!")
+            resp, err = api_call("POST", "/ingest", files=files)
+            if err:
+                st.error(err)
+            elif resp and resp.status_code == 200:
+                st.success("Ingested successfully")
+                st.session_state.ingested_docs.append(uploaded_file.name)
             else:
-                st.error("❌ Upload failed.")
-    
-    st.divider()
-    
-    # Recents List
-    st.markdown("<div class='recent-header'>RECENTS</div>", unsafe_allow_html=True)
-    
-    # Try fetching real recents, otherwise show static UI
-    response, error = safe_api_call("GET", "/sources")
-    if not error and response and response.status_code == 200:
-        documents = response.json().get("documents", [])
-        for doc in documents[:5]: # Show top 5
-            st.markdown(f"<div class='recent-item'>🕒 {doc[:20]}{'...' if len(doc)>20 else ''}</div>", unsafe_allow_html=True)
-    else:
-        # Static Mockup UI items based on the image
-        st.markdown("<div class='recent-item'>🕒 Q3 Financial Report</div>", unsafe_allow_html=True)
-        st.markdown("<div class='recent-item'>🕒 Product Roadmap 2024</div>", unsafe_allow_html=True)
-        st.markdown("<div class='recent-item'>🕒 Research Summary</div>", unsafe_allow_html=True)
+                detail = "Unknown error"
+                try:
+                    detail = resp.json().get("detail", resp.text)
+                except Exception:
+                    if resp is not None:
+                        detail = resp.text
+                st.error(detail)
+    st.markdown("</div>", unsafe_allow_html=True)
+
+    st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)
+    st.markdown('<div class="sidebar-heading">Recents</div>', unsafe_allow_html=True)
+    recent_items = st.session_state.ingested_docs or st.session_state.chat_sessions
+    for item in recent_items[:6]:
+        label = item if len(item) <= 26 else item[:23] + "..."
+        st.markdown(
+            f"<div class='recents-item'><span class='icon'>🕘</span><span>{esc(label)}</span></div>",
+            unsafe_allow_html=True,
+        )
+    st.markdown("</div>", unsafe_allow_html=True)
+
 
 # ============================================================================
-# UI: MAIN INTERFACE
+# MAIN UI
 # ============================================================================
-
-# Top Navigation Badge
-st.markdown("""
-    <div class="top-badge-container">
-        <span class="status-badge">CURRENT SESSION ACTIVE</span>
+st.markdown(
+    """
+    <div class="session-row">
+        <div class="session-label">CURRENT SESSION</div>
+        <div class="session-badge">ACTIVE</div>
     </div>
-""", unsafe_allow_html=True)
+    <hr>
+    """,
+    unsafe_allow_html=True,
+)
 
-# Show backend warning cleanly if down
-if not st.session_state.backend_healthy:
-    st.warning("⚠️ Backend API is offline. Starting up or unreachable.")
+if not st.session_state.backend_ok:
+    st.markdown(
+        """
+        <div style="margin:1rem 1.25rem 0; padding:0.85rem 1rem; border-radius:16px;
+                    border:1px solid rgba(245, 158, 11, 0.18); background: rgba(245, 158, 11, 0.08);
+                    color:#b45309; font-size:0.9rem;">
+            ⚠️ Backend is starting up. If it still shows after a minute, refresh the page.
+        </div>
+        """,
+        unsafe_allow_html=True,
+    )
 
-# Welcome Screen (If no messages)
-if len(st.session_state.messages) == 0:
-    st.markdown("<div class='welcome-header'>How can I assist your<br>research today?</div>", unsafe_allow_html=True)
-    st.markdown("<div class='welcome-sub'>I can analyze documents, summarize findings, or<br>answer specific questions about your archived data.</div>", unsafe_allow_html=True)
+# Empty state
+if not st.session_state.messages:
+    st.markdown(
+        """
+        <div class="hero-wrap">
+            <h1 class="hero-title">How can I assist your research today?</h1>
+            <p class="hero-subtitle">
+                I can analyze documents, summarize findings, or answer specific questions about your archived data.
+            </p>
+        </div>
+        """,
+        unsafe_allow_html=True,
+    )
 
-# Container for chat messages
-chat_container = st.container()
+# Chat history
+for msg in st.session_state.messages:
+    role = msg.get("role", "assistant")
+    content = msg.get("content", "")
+    ts = msg.get("time", "")
+    sources = msg.get("sources", [])
 
-with chat_container:
-    for msg in st.session_state.messages:
-        is_user = msg["role"] == "user"
-        # Hide the default user avatar to match the clean dark bubble look, use logo for assistant
-        avatar_icon = "◧" if not is_user else None 
-        
-        with st.chat_message(msg["role"], avatar=avatar_icon):
-            st.markdown(msg["content"])
+    if role == "user":
+        st.markdown(
+            f"""
+            <div class="msg-user-wrap">
+                <div>
+                    <div class="msg-user">{render_text_block(content)}</div>
+                    <div class="msg-time" style="text-align:right; padding-right:0.4rem;">{esc(ts)}</div>
+                </div>
+            </div>
+            """,
+            unsafe_allow_html=True,
+        )
+    else:
+        st.markdown(
+            f"""
+            <div class="msg-bot-row">
+                <div class="bot-avatar">🧠</div>
+                <div>
+                    <div class="msg-bot">{render_assistant_text(content)}</div>
+                    <div class="msg-time">{esc(ts)}</div>
+                </div>
+            </div>
+            """,
+            unsafe_allow_html=True,
+        )
 
-# ============================================================================
-# UI: CHAT INPUT
-# ============================================================================
+        if sources and st.session_state.show_sources:
+            with st.expander(f"📚 {len(sources)} source(s) cited"):
+                for i, src in enumerate(sources, start=1):
+                    score = float(src.get("score", 0) or 0)
+                    source_name = esc(str(src.get("source", "Unknown")))
+                    snippet = esc(str(src.get("content", ""))[:400])
+                    if len(str(src.get("content", ""))) > 400:
+                        snippet += "..."
+
+                    st.markdown(
+                        f"""
+                        <div class="source-card">
+                            <div class="source-head">
+                                <span class="source-title">Source {i}</span>
+                                <span class="source-score">{score:.0%} match</span>
+                            </div>
+                            <div class="source-file">📄 {source_name}</div>
+                            <div class="source-snippet">{snippet}</div>
+                        </div>
+                        """,
+                        unsafe_allow_html=True,
+                    )
 
-st.write("") # Vertical spacing
 
+# Chat input
+st.markdown('<div class="bottom-space"></div>', unsafe_allow_html=True)
 user_input = st.chat_input("Ask a question...")
 
 if user_input:
-    # Add and display user message immediately
-    st.session_state.messages.append({"role": "user", "content": user_input})
-    with chat_container:
-        with st.chat_message("user", avatar=None):
-            st.markdown(user_input)
-    
-    # Get and display assistant response
-    with chat_container:
-        with st.chat_message("assistant", avatar="◧"):
-            placeholder = st.empty()
-            full_response = ""
-            sources = []
-            
-            response, error = safe_api_call(
-                "POST", 
-                "/query", 
-                json={"question": user_input}, 
-                stream=True
-            )
-            
-            if error:
-                st.error(error)
-                full_response = error
-            elif response:
+    user_ts = now_time()
+    st.session_state.messages.append({"role": "user", "content": user_input, "time": user_ts})
+
+    # Render user bubble immediately
+    st.markdown(
+        f"""
+        <div class="msg-user-wrap">
+            <div>
+                <div class="msg-user">{render_text_block(user_input)}</div>
+                <div class="msg-time" style="text-align:right; padding-right:0.4rem;">{esc(user_ts)}</div>
+            </div>
+        </div>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    full_response = ""
+    sources: list[dict] = []
+    assistant_placeholder = st.empty()
+
+    # Assistant shell
+    st.markdown(
+        """
+        <div class="msg-bot-row">
+            <div class="bot-avatar">🧠</div>
+            <div style="width:100%;">
+        """,
+        unsafe_allow_html=True,
+    )
+
+    resp, err = api_call("POST", "/query", json={"question": user_input}, stream=True)
+
+    if err:
+        full_response = err
+        assistant_placeholder.markdown(
+            f"""
+            <div class="msg-bot" style="border-color: rgba(239, 68, 68, 0.25); color:#b91c1c; background:#fff5f5;">
+                {esc(err)}
+            </div>
+            """,
+            unsafe_allow_html=True,
+        )
+    elif resp:
+        try:
+            buf = ""
+            for line in resp.iter_lines():
+                if not line:
+                    continue
                 try:
-                    for line in response.iter_lines():
-                        if line:
-                            try:
-                                decoded_line = line.decode('utf-8')
-                                data = json.loads(decoded_line)
-                                
-                                if data.get("type") == "sources":
-                                    sources = data.get("data", [])
-                                elif data.get("type") == "token":
-                                    full_response += data.get("content", "")
-                                    placeholder.markdown(full_response + "▌")
-                            except json.JSONDecodeError:
-                                continue
-                    
-                    # Convert bullet points to checkmarks to match mockup design style
-                    styled_response = full_response.replace("- ", "✅ ")
-                    placeholder.markdown(styled_response)
-                                
-                except Exception as e:
-                    error_msg = f"❌ Error processing response: {str(e)}"
-                    placeholder.markdown(error_msg)
-                    full_response = error_msg
-            
-    # Save to state
-    st.session_state.messages.append({
-        "role": "assistant",
-        "content": full_response,
-        "sources": sources
-    })
+                    data = json.loads(line.decode("utf-8"))
+                except json.JSONDecodeError:
+                    continue
+
+                if data.get("type") == "sources":
+                    sources = data.get("data", []) or []
+                elif data.get("type") == "token":
+                    buf += data.get("content", "")
+                    full_response = buf
+                    assistant_placeholder.markdown(
+                        f"""
+                        <div class="msg-bot">{render_assistant_text(buf)}<span class="typing-cursor">▌</span></div>
+                        """,
+                        unsafe_allow_html=True,
+                    )
+
+            assistant_placeholder.markdown(
+                f"""
+                <div class="msg-bot">{render_assistant_text(full_response)}</div>
+                """,
+                unsafe_allow_html=True,
+            )
+        except Exception as e:
+            full_response = f"Error: {e}"
+            assistant_placeholder.markdown(
+                f"""
+                <div class="msg-bot" style="border-color: rgba(239, 68, 68, 0.25); color:#b91c1c; background:#fff5f5;">
+                    {esc(full_response)}
+                </div>
+                """,
+                unsafe_allow_html=True,
+            )
+
+    st.markdown("</div></div>", unsafe_allow_html=True)
+
+    if sources and st.session_state.show_sources:
+        with st.expander(f"📚 {len(sources)} source(s) cited"):
+            for i, src in enumerate(sources, start=1):
+                score = float(src.get("score", 0) or 0)
+                source_name = esc(str(src.get("source", "Unknown")))
+                snippet = esc(str(src.get("content", ""))[:400])
+                if len(str(src.get("content", ""))) > 400:
+                    snippet += "..."
+
+                st.markdown(
+                    f"""
+                    <div class="source-card">
+                        <div class="source-head">
+                            <span class="source-title">Source {i}</span>
+                            <span class="source-score">{score:.0%} match</span>
+                        </div>
+                        <div class="source-file">📄 {source_name}</div>
+                        <div class="source-snippet">{snippet}</div>
+                    </div>
+                    """,
+                    unsafe_allow_html=True,
+                )
+
+    bot_ts = now_time()
+    st.session_state.messages.append(
+        {
+            "role": "assistant",
+            "content": full_response,
+            "sources": sources,
+            "time": bot_ts,
+        }
+    )
+
+
+# ============================================================================
+# FOOTER NOTE
+# ============================================================================
+st.markdown(
+    """
+    <div style="height:1.5rem;"></div>
+    """,
+    unsafe_allow_html=True,
+)