debug

app/main.py

@@ -1,269 +1,276 @@
-# streamlit_app.py
-import streamlit as st
-from streamlit_autorefresh import st_autorefresh # For periodic refresh
+# app/main.py
+import gradio as gr
 import httpx
-import asyncio
 import websockets
+import asyncio
 import json
-import threading
-import queue
-import logging
-import time
 import os
+import logging
+from contextlib import asynccontextmanager
 
-# Import backend components
-from app import crud, models, schemas, auth, dependencies
-from app.database import ensure_db_and_table_exist # Sync function
-from app.websocket import manager # Import the manager instance
+from fastapi import FastAPI, Depends # Import FastAPI itself
+from .database import connect_db, disconnect_db, database, metadata, users
+from .api import router as api_router
+from . import schemas, auth, dependencies
+from .websocket import manager # Import the connection manager instance
 
-# FastAPI imports for mounting
-from fastapi import FastAPI, Depends, HTTPException, status
-from fastapi.routing import Mount
-from fastapi.staticfiles import StaticFiles
-from app.api import router as api_router # Import the specific API router
+from sqlalchemy.schema import CreateTable
+from sqlalchemy.dialects import sqlite
 
-# --- Logging ---
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --- Configuration ---
-# Use environment variable or default for local vs deployed API endpoint
-# Since we are mounting FastAPI within Streamlit for the HF Space deployment:
-API_BASE_URL = "http://127.0.0.1:7860/api" # Calls within the same process
-WS_BASE_URL = API_BASE_URL.replace("http", "ws")
-
-# --- Ensure DB exists on first run ---
-# This runs once per session/process start
-ensure_db_and_table_exist()
-
-# --- FastAPI Mounting Setup ---
-# Create a FastAPI instance (separate from the Streamlit one)
-# We won't run this directly with uvicorn, but Streamlit uses it internally
-api_app = FastAPI(title="Backend API") # Can add lifespan if needed for API-specific setup later
-api_app.include_router(api_router, prefix="/api")
-
-# Mount the FastAPI app within Streamlit's internal Tornado server
-# This requires monkey-patching or using available hooks if Streamlit allows.
-# Simpler approach for HF Space: Run FastAPI separately is cleaner if possible.
-# Reverting to the idea that for this HF Space demo, API calls will be internal HTTP requests.
-
-# --- WebSocket Listener Thread ---
-stop_event = threading.Event()
-notification_queue = queue.Queue()
-
-def websocket_listener(token: str):
-    """Runs in a background thread to listen for WebSocket messages."""
-    logger.info(f"[WS Thread] Listener started for token: {token[:10]}...")
-    ws_url = f"{WS_BASE_URL}/ws/{token}"
-
-    async def listen():
-        try:
-            async with websockets.connect(ws_url, open_timeout=10.0) as ws:
-                logger.info(f"[WS Thread] Connected to {ws_url}")
-                st.session_state['ws_connected'] = True
-                while not stop_event.is_set():
-                    try:
-                        message = await asyncio.wait_for(ws.recv(), timeout=1.0) # Check stop_event frequently
-                        logger.info(f"[WS Thread] Received message: {message[:100]}...")
-                        try:
-                             data = json.loads(message)
-                             if data.get("type") == "new_user":
-                                  notification = schemas.Notification(**data)
-                                  notification_queue.put(notification.message) # Put message in queue
-                                  logger.info("[WS Thread] Put notification in queue.")
-                        except json.JSONDecodeError:
-                             logger.error("[WS Thread] Failed to decode JSON.")
-                        except Exception as e:
-                             logger.error(f"[WS Thread] Error processing message: {e}")
-
-                    except asyncio.TimeoutError:
-                        continue # No message, check stop_event again
-                    except websockets.ConnectionClosed:
-                        logger.warning("[WS Thread] Connection closed.")
-                        break # Exit loop if closed
-        except Exception as e:
-            logger.error(f"[WS Thread] Connection failed or error: {e}")
-        finally:
-            logger.info("[WS Thread] Listener loop finished.")
-            st.session_state['ws_connected'] = False
+API_BASE_URL = "http://127.0.0.1:7860/api"
 
-    try:
-        asyncio.run(listen())
-    except Exception as e:
-        logger.error(f"[WS Thread] asyncio.run error: {e}")
-    logger.info("[WS Thread] Listener thread exiting.")
-
-
-# --- Streamlit UI ---
-
-st.set_page_config(layout="wide")
-
-# --- Initialize Session State ---
-if 'logged_in' not in st.session_state:
-    st.session_state.logged_in = False
-if 'token' not in st.session_state:
-    st.session_state.token = None
-if 'user_email' not in st.session_state:
-    st.session_state.user_email = None
-if 'notifications' not in st.session_state:
-    st.session_state.notifications = []
-if 'ws_thread' not in st.session_state:
-    st.session_state.ws_thread = None
-if 'ws_connected' not in st.session_state:
-     st.session_state.ws_connected = False
-
-# --- Notification Processing ---
-new_notifications = []
-while not notification_queue.empty():
-    try:
-        msg = notification_queue.get_nowait()
-        new_notifications.append(msg)
-    except queue.Empty:
-        break
-
-if new_notifications:
-    logger.info(f"Processing {len(new_notifications)} notifications from queue.")
-    # Prepend new notifications to the session state list
-    current_list = st.session_state.notifications
-    st.session_state.notifications = new_notifications + current_list
-    # Limit history
-    if len(st.session_state.notifications) > 15:
-        st.session_state.notifications = st.session_state.notifications[:15]
-    # No explicit rerun needed here, Streamlit should rerun due to state change (?)
-    # or due to autorefresh below.
-
-# --- Auto Refresh ---
-# Refresh every 2 seconds to check the queue and update display
-count = st_autorefresh(interval=2000, limit=None, key="notifrefresh")
-
-# --- API Client ---
-client = httpx.AsyncClient(base_url=API_BASE_URL, timeout=10.0)
-
-# --- Helper Functions for API Calls ---
-async def api_register(email, password):
-    try:
-        response = await client.post("/register", json={"email": email, "password": password})
-        response.raise_for_status()
-        return {"success": True, "data": response.json()}
-    except httpx.HTTPStatusError as e:
-        detail = e.response.json().get("detail", e.response.text)
-        logger.error(f"API Register Error: {e.response.status_code} - {detail}")
-        return {"success": False, "error": f"API Error: {detail}"}
-    except Exception as e:
-        logger.exception("Register call failed")
-        return {"success": False, "error": f"Request failed: {e}"}
-
-async def api_login(email, password):
-    try:
-        response = await client.post("/login", json={"email": email, "password": password})
-        response.raise_for_status()
-        return {"success": True, "data": response.json()}
-    except httpx.HTTPStatusError as e:
-        detail = e.response.json().get("detail", e.response.text)
-        logger.error(f"API Login Error: {e.response.status_code} - {detail}")
-        return {"success": False, "error": f"API Error: {detail}"}
-    except Exception as e:
-        logger.exception("Login call failed")
-        return {"success": False, "error": f"Request failed: {e}"}
-
-# --- UI Rendering ---
-st.title("Authentication & Notification App (Streamlit)")
-
-if not st.session_state.logged_in:
-    st.sidebar.header("Login or Register")
-    login_tab, register_tab = st.sidebar.tabs(["Login", "Register"])
-
-    with login_tab:
-        with st.form("login_form"):
-            login_email = st.text_input("Email", key="login_email")
-            login_password = st.text_input("Password", type="password", key="login_password")
-            login_button = st.form_submit_button("Login")
-
-            if login_button:
-                if not login_email or not login_password:
-                    st.error("Please enter email and password.")
-                else:
-                    result = asyncio.run(api_login(login_email, login_password)) # Run async in sync context
-                    if result["success"]:
-                        token = result["data"]["access_token"]
-                        # Attempt to get user info immediately - needs modification if /users/me requires auth header
-                        # For simplicity, just store email from login form for now
-                        st.session_state.logged_in = True
-                        st.session_state.token = token
-                        st.session_state.user_email = login_email # Store email used for login
-                        st.session_state.notifications = [] # Clear old notifications
-
-                        # Start WebSocket listener thread
-                        stop_event.clear() # Ensure stop event is clear
-                        thread = threading.Thread(target=websocket_listener, args=(token,), daemon=True)
-                        st.session_state.ws_thread = thread
-                        thread.start()
-                        logger.info("Login successful, WS thread started.")
-                        st.rerun() # Rerun immediately to switch view
-                    else:
-                        st.error(f"Login failed: {result['error']}")
-
-    with register_tab:
-        with st.form("register_form"):
-            reg_email = st.text_input("Email", key="reg_email")
-            reg_password = st.text_input("Password", type="password", key="reg_password")
-            reg_confirm = st.text_input("Confirm Password", type="password", key="reg_confirm")
-            register_button = st.form_submit_button("Register")
-
-            if register_button:
-                if not reg_email or not reg_password or not reg_confirm:
-                    st.error("Please fill all fields.")
-                elif reg_password != reg_confirm:
-                    st.error("Passwords do not match.")
-                elif len(reg_password) < 8:
-                    st.error("Password must be at least 8 characters.")
-                else:
-                    result = asyncio.run(api_register(reg_email, reg_password))
-                    if result["success"]:
-                        st.success(f"Registration successful for {result['data']['email']}! Please log in.")
-                    else:
-                        st.error(f"Registration failed: {result['error']}")
-
-else: # Logged In View
-    st.sidebar.header(f"Welcome, {st.session_state.user_email}!")
-    if st.sidebar.button("Logout"):
-        logger.info("Logout requested.")
-        # Stop WebSocket thread
-        if st.session_state.ws_thread and st.session_state.ws_thread.is_alive():
-             logger.info("Signalling WS thread to stop.")
-             stop_event.set()
-             st.session_state.ws_thread.join(timeout=2.0) # Wait briefly for thread exit
-             if st.session_state.ws_thread.is_alive():
-                  logger.warning("WS thread did not exit cleanly.")
-        # Clear session state
-        st.session_state.logged_in = False
-        st.session_state.token = None
-        st.session_state.user_email = None
-        st.session_state.notifications = []
-        st.session_state.ws_thread = None
-        st.session_state.ws_connected = False
-        logger.info("Session cleared.")
-        st.rerun()
-
-    st.header("Dashboard")
-    # Display notifications
-    st.subheader("Real-time Notifications")
-    ws_status = "Connected" if st.session_state.ws_connected else "Disconnected"
-    st.caption(f"WebSocket Status: {ws_status}")
-
-    if st.session_state.notifications:
-        for i, msg in enumerate(st.session_state.notifications):
-            st.info(f"{msg}", icon="🔔")
+# --- Lifespan (remains the same) ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # ... (same DB setup code) ...
+    logger.info("Application startup: Connecting DB...")
+    await connect_db()
+    logger.info("Application startup: DB Connected. Checking/Creating tables...")
+    if database.is_connected:
+        try:
+            check_query = "SELECT name FROM sqlite_master WHERE type='table' AND name=:table_name;"
+            table_exists = await database.fetch_one(query=check_query, values={"table_name": users.name})
+            if not table_exists:
+                logger.info(f"Table '{users.name}' not found, attempting creation using async connection...")
+                dialect = sqlite.dialect()
+                create_table_stmt = str(CreateTable(users).compile(dialect=dialect))
+                await database.execute(query=create_table_stmt)
+                logger.info(f"Table '{users.name}' created successfully via async connection.")
+                table_exists_after = await database.fetch_one(query=check_query, values={"table_name": users.name})
+                if table_exists_after: logger.info(f"Table '{users.name}' verified after creation.")
+                else: logger.error(f"Table '{users.name}' verification FAILED after creation attempt!")
+            else:
+                logger.info(f"Table '{users.name}' already exists (checked via async connection).")
+        except Exception as db_setup_err:
+            logger.exception(f"CRITICAL error during async DB table setup: {db_setup_err}")
     else:
-        st.text("No new notifications.")
-
-    # Add a button to manually check queue/refresh if needed
-    # if st.button("Check for notifications"):
-    #      st.rerun() # Force rerun which includes queue check
+        logger.error("CRITICAL: Database connection failed, skipping table setup.")
+    logger.info("Application startup: DB setup phase complete.")
+    yield
+    logger.info("Application shutdown: Disconnecting DB...")
+    await disconnect_db()
+    logger.info("Application shutdown: DB Disconnected.")
+
+
+# --- FastAPI App Setup (remains the same) ---
+app = FastAPI(lifespan=lifespan)
+app.include_router(api_router, prefix="/api")
+
+# --- Helper functions (make_api_request remains the same) ---
+async def make_api_request(method: str, endpoint: str, **kwargs):
+    async with httpx.AsyncClient() as client:
+        url = f"{API_BASE_URL}{endpoint}"
+        try: response = await client.request(method, url, **kwargs); response.raise_for_status(); return response.json()
+        except httpx.RequestError as e: logger.error(f"HTTP Request failed: {e.request.method} {e.request.url} - {e}"); return {"error": f"Network error: {e}"}
+        except httpx.HTTPStatusError as e:
+            logger.error(f"HTTP Status error: {e.response.status_code} - {e.response.text}")
+            try: detail = e.response.json().get("detail", e.response.text)
+            except json.JSONDecodeError: detail = e.response.text
+            return {"error": f"API Error: {detail}"}
+        except Exception as e: logger.error(f"Unexpected error during API call: {e}"); return {"error": f"Unexpected error: {str(e)}"}
+
+# --- WebSocket handling ---
+# <<< Pass state objects by reference >>>
+async def listen_to_websockets(token: str, notification_list_state: gr.State, notification_trigger_state: gr.State):
+    """Connects to WS and updates state list and trigger when a message arrives."""
+    ws_listener_id = f"WSListener-{os.getpid()}-{asyncio.current_task().get_name()}"
+    logger.info(f"[{ws_listener_id}] Starting WebSocket listener task.")
+
+    if not token:
+        logger.warning(f"[{ws_listener_id}] No token provided. Task exiting.")
+        return # Just exit, don't need to return state values
+
+    ws_url_base = API_BASE_URL.replace("http", "ws")
+    ws_url = f"{ws_url_base}/ws/{token}"
+    logger.info(f"[{ws_listener_id}] Attempting to connect: {ws_url}")
 
-
-# --- Final Cleanup ---
-# Ensure httpx client is closed if script exits abnormally
-# (This might not always run depending on how Streamlit terminates)
-# Ideally handled within context managers if used more extensively
-# asyncio.run(client.aclose())
+    try:
+        async with websockets.connect(ws_url, open_timeout=15.0) as websocket:
+            logger.info(f"[{ws_listener_id}] WebSocket connected successfully.")
+            while True:
+                try:
+                    message_str = await asyncio.wait_for(websocket.recv(), timeout=300.0)
+                    logger.info(f"[{ws_listener_id}] Received: {message_str[:100]}...")
+                    try:
+                         message_data = json.loads(message_str)
+                         if message_data.get("type") == "new_user":
+                            notification = schemas.Notification(**message_data)
+                            logger.info(f"[{ws_listener_id}] Processing 'new_user': {notification.message}")
+                            # --- Modify state values directly ---
+                            current_list = notification_list_state.value.copy() # Operate on copy
+                            current_list.insert(0, notification.message)
+                            if len(current_list) > 10: current_list.pop()
+                            notification_list_state.value = current_list # Assign back modified copy
+                            notification_trigger_state.value += 1 # Increment trigger
+                            # --- Log the update ---
+                            logger.info(f"[{ws_listener_id}] States updated: list len={len(notification_list_state.value)}, trigger={notification_trigger_state.value}")
+                         else:
+                              logger.warning(f"[{ws_listener_id}] Unknown message type: {message_data.get('type')}")
+                    # ... (error handling for parsing) ...
+                    except json.JSONDecodeError: logger.error(f"[{ws_listener_id}] JSON Decode Error: {message_str}")
+                    except Exception as parse_err: logger.error(f"[{ws_listener_id}] Message Processing Error: {parse_err}")
+                # ... (error handling for websocket recv/connection) ...
+                except asyncio.TimeoutError: logger.debug(f"[{ws_listener_id}] WebSocket recv timed out."); continue
+                except websockets.ConnectionClosedOK: logger.info(f"[{ws_listener_id}] WebSocket closed normally."); break
+                except websockets.ConnectionClosedError as e: logger.error(f"[{ws_listener_id}] WebSocket closed with error: {e}"); break
+                except Exception as e: logger.error(f"[{ws_listener_id}] Listener loop error: {e}"); await asyncio.sleep(1); break
+    # ... (error handling for websocket connect) ...
+    except asyncio.TimeoutError: logger.error(f"[{ws_listener_id}] WebSocket initial connection timed out.")
+    except websockets.exceptions.InvalidURI: logger.error(f"[{ws_listener_id}] Invalid WebSocket URI.")
+    except websockets.exceptions.WebSocketException as e: logger.error(f"[{ws_listener_id}] WebSocket connection failed: {e}")
+    except Exception as e: logger.error(f"[{ws_listener_id}] Unexpected error in listener task: {e}")
+
+    logger.info(f"[{ws_listener_id}] Listener task finished.")
+    # No need to return state values when task ends
+
+# --- Gradio Interface ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    # State variables
+    auth_token = gr.State(None)
+    user_info = gr.State(None)
+    notification_list = gr.State([])
+    websocket_task = gr.State(None)
+    # Add trigger states
+    notification_trigger = gr.State(0)
+    last_polled_trigger = gr.State(0) # State to track last seen trigger
+
+    # --- UI Components ---
+    with gr.Tabs() as tabs:
+        # --- Registration/Login Tabs (remain the same) ---
+        with gr.TabItem("Register", id="register_tab"):
+            gr.Markdown("## Create a new account")
+            reg_email = gr.Textbox(label="Email", type="email")
+            reg_password = gr.Textbox(label="Password (min 8 chars)", type="password")
+            reg_confirm_password = gr.Textbox(label="Confirm Password", type="password")
+            reg_button = gr.Button("Register")
+            reg_status = gr.Textbox(label="Status", interactive=False)
+        with gr.TabItem("Login", id="login_tab"):
+            gr.Markdown("## Login to your account")
+            login_email = gr.Textbox(label="Email", type="email")
+            login_password = gr.Textbox(label="Password", type="password")
+            login_button = gr.Button("Login")
+            login_status = gr.Textbox(label="Status", interactive=False)
+
+        # --- Welcome Tab ---
+        with gr.TabItem("Welcome", id="welcome_tab", visible=False) as welcome_tab:
+            gr.Markdown("## Welcome!", elem_id="welcome_header")
+            welcome_message = gr.Markdown("", elem_id="welcome_message")
+            logout_button = gr.Button("Logout")
+            gr.Markdown("---")
+            gr.Markdown("## Real-time Notifications")
+            notification_display = gr.Textbox( # Visible display
+                label="New User Alerts", lines=5, max_lines=10, interactive=False
+            )
+            # <<< Hidden component for polling >>>
+            dummy_poller = gr.Number(label="poller", value=0, visible=False, every=1)
+
+
+    # --- Event Handlers ---
+
+    # Registration Logic (remains the same)
+    async def handle_register(email, password, confirm_password):
+        if not email or not password or not confirm_password: return gr.update(value="Please fill fields.")
+        if password != confirm_password: return gr.update(value="Passwords mismatch.")
+        if len(password) < 8: return gr.update(value="Password >= 8 chars.")
+        payload = {"email": email, "password": password}
+        result = await make_api_request("post", "/register", json=payload)
+        if "error" in result: return gr.update(value=f"Register failed: {result['error']}")
+        else: return gr.update(value=f"Register success: {result.get('email')}! Log in.")
+    reg_button.click(handle_register, inputs=[reg_email, reg_password, reg_confirm_password], outputs=[reg_status])
+
+    # Login Logic
+    # <<< MODIFIED: Pass/Reset both trigger states >>>
+    async def handle_login(email, password, current_task, current_trigger_val, current_last_poll_val):
+        # Define failure outputs matching the outputs list
+        fail_outputs = (gr.update(value="..."), None, None, None, gr.update(visible=False), None, current_task, current_trigger_val, current_last_poll_val)
+
+        if not email or not password: return fail_outputs[:1] + (gr.update(value="Enter email/password"),) + fail_outputs[2:]
+
+        payload = {"email": email, "password": password}
+        result = await make_api_request("post", "/login", json=payload)
+
+        if "error" in result: return (gr.update(value=f"Login failed: {result['error']}"),) + fail_outputs[1:]
+        else:
+            token = result.get("access_token")
+            user_data = await dependencies.get_optional_current_user(token)
+            if not user_data: return (gr.update(value="Login ok but user fetch failed."),) + fail_outputs[1:]
+
+            if current_task and not current_task.done():
+                current_task.cancel()
+                try: await current_task
+                except asyncio.CancelledError: logger.info("Previous WebSocket task cancelled.")
+
+            # <<< Pass state objects to listener >>>
+            new_task = asyncio.create_task(listen_to_websockets(token, notification_list, notification_trigger))
+
+            welcome_msg = f"Welcome, {user_data.email}!"
+            # Reset triggers on successful login
+            return (
+                gr.update(value="Login successful!"), token, user_data.model_dump(), # status, token, user_info
+                gr.update(selected="welcome_tab"), gr.update(visible=True), gr.update(value=welcome_msg), # UI changes
+                new_task, 0, 0 # websocket_task, notification_trigger, last_polled_trigger
+            )
+
+    # <<< MODIFIED: Add last_polled_trigger to inputs/outputs >>>
+    login_button.click(
+        handle_login,
+        inputs=[login_email, login_password, websocket_task, notification_trigger, last_polled_trigger],
+        outputs=[login_status, auth_token, user_info, tabs, welcome_tab, welcome_message, websocket_task, notification_trigger, last_polled_trigger]
+    )
+
+
+    # <<< Polling function >>>
+    def poll_and_update(current_trigger_value, last_known_trigger, current_notif_list):
+        """ Checks trigger state change and updates UI if needed. """
+        if current_trigger_value != last_known_trigger:
+            logger.info(f"Polling detected trigger change ({last_known_trigger} -> {current_trigger_value}). Updating UI.")
+            new_value = "\n".join(current_notif_list)
+            # Return new display value AND update last_known_trigger state value
+            return gr.update(value=new_value), current_trigger_value
+        else:
+            # No change, return NoUpdate for display AND existing last_known_trigger value
+            return gr.NoUpdate(), last_known_trigger
+
+    # <<< Attach polling function to the dummy component's change event >>>
+    dummy_poller.change(
+        fn=poll_and_update,
+        inputs=[notification_trigger, last_polled_trigger, notification_list],
+        outputs=[notification_display, last_polled_trigger], # Update display & last polled state
+        queue=False # Try immediate update
+    )
+
+
+    # Logout Logic
+    # <<< MODIFIED: Reset both trigger states >>>
+    async def handle_logout(current_task):
+        if current_task and not current_task.done():
+            current_task.cancel()
+            try: await current_task
+            except asyncio.CancelledError: logger.info("WebSocket task cancelled on logout.")
+        # Reset all relevant states
+        return ( None, None, [], None, # auth_token, user_info, notification_list, websocket_task
+                 gr.update(selected="login_tab"), gr.update(visible=False), # tabs, welcome_tab
+                 gr.update(value=""), gr.update(value=""), # welcome_message, login_status
+                 0, 0 ) # notification_trigger, last_polled_trigger (reset)
+
+    # <<< MODIFIED: Add last_polled_trigger to outputs >>>
+    logout_button.click(
+        handle_logout,
+        inputs=[websocket_task],
+        outputs=[
+            auth_token, user_info, notification_list, websocket_task,
+            tabs, welcome_tab, welcome_message, login_status,
+            notification_trigger, last_polled_trigger # Add trigger states here
+        ]
+    )
+
+# Mount Gradio App (remains the same)
+app = gr.mount_gradio_app(app, demo, path="/")
+
+# Run Uvicorn (remains the same)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app.main:app", host="0.0.0.0", port=7860, reload=True)

app/streamlit_app.py

@@ -0,0 +1,269 @@
+# streamlit_app.py
+import streamlit as st
+from streamlit_autorefresh import st_autorefresh # For periodic refresh
+import httpx
+import asyncio
+import websockets
+import json
+import threading
+import queue
+import logging
+import time
+import os
+
+# Import backend components
+from app import crud, models, schemas, auth, dependencies
+from app.database import ensure_db_and_table_exist # Sync function
+from app.websocket import manager # Import the manager instance
+
+# FastAPI imports for mounting
+from fastapi import FastAPI, Depends, HTTPException, status
+from fastapi.routing import Mount
+from fastapi.staticfiles import StaticFiles
+from app.api import router as api_router # Import the specific API router
+
+# --- Logging ---
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# --- Configuration ---
+# Use environment variable or default for local vs deployed API endpoint
+# Since we are mounting FastAPI within Streamlit for the HF Space deployment:
+API_BASE_URL = "http://127.0.0.1:7860/api" # Calls within the same process
+WS_BASE_URL = API_BASE_URL.replace("http", "ws")
+
+# --- Ensure DB exists on first run ---
+# This runs once per session/process start
+ensure_db_and_table_exist()
+
+# --- FastAPI Mounting Setup ---
+# Create a FastAPI instance (separate from the Streamlit one)
+# We won't run this directly with uvicorn, but Streamlit uses it internally
+api_app = FastAPI(title="Backend API") # Can add lifespan if needed for API-specific setup later
+api_app.include_router(api_router, prefix="/api")
+
+# Mount the FastAPI app within Streamlit's internal Tornado server
+# This requires monkey-patching or using available hooks if Streamlit allows.
+# Simpler approach for HF Space: Run FastAPI separately is cleaner if possible.
+# Reverting to the idea that for this HF Space demo, API calls will be internal HTTP requests.
+
+# --- WebSocket Listener Thread ---
+stop_event = threading.Event()
+notification_queue = queue.Queue()
+
+def websocket_listener(token: str):
+    """Runs in a background thread to listen for WebSocket messages."""
+    logger.info(f"[WS Thread] Listener started for token: {token[:10]}...")
+    ws_url = f"{WS_BASE_URL}/ws/{token}"
+
+    async def listen():
+        try:
+            async with websockets.connect(ws_url, open_timeout=10.0) as ws:
+                logger.info(f"[WS Thread] Connected to {ws_url}")
+                st.session_state['ws_connected'] = True
+                while not stop_event.is_set():
+                    try:
+                        message = await asyncio.wait_for(ws.recv(), timeout=1.0) # Check stop_event frequently
+                        logger.info(f"[WS Thread] Received message: {message[:100]}...")
+                        try:
+                             data = json.loads(message)
+                             if data.get("type") == "new_user":
+                                  notification = schemas.Notification(**data)
+                                  notification_queue.put(notification.message) # Put message in queue
+                                  logger.info("[WS Thread] Put notification in queue.")
+                        except json.JSONDecodeError:
+                             logger.error("[WS Thread] Failed to decode JSON.")
+                        except Exception as e:
+                             logger.error(f"[WS Thread] Error processing message: {e}")
+
+                    except asyncio.TimeoutError:
+                        continue # No message, check stop_event again
+                    except websockets.ConnectionClosed:
+                        logger.warning("[WS Thread] Connection closed.")
+                        break # Exit loop if closed
+        except Exception as e:
+            logger.error(f"[WS Thread] Connection failed or error: {e}")
+        finally:
+            logger.info("[WS Thread] Listener loop finished.")
+            st.session_state['ws_connected'] = False
+
+    try:
+        asyncio.run(listen())
+    except Exception as e:
+        logger.error(f"[WS Thread] asyncio.run error: {e}")
+    logger.info("[WS Thread] Listener thread exiting.")
+
+
+# --- Streamlit UI ---
+
+st.set_page_config(layout="wide")
+
+# --- Initialize Session State ---
+if 'logged_in' not in st.session_state:
+    st.session_state.logged_in = False
+if 'token' not in st.session_state:
+    st.session_state.token = None
+if 'user_email' not in st.session_state:
+    st.session_state.user_email = None
+if 'notifications' not in st.session_state:
+    st.session_state.notifications = []
+if 'ws_thread' not in st.session_state:
+    st.session_state.ws_thread = None
+if 'ws_connected' not in st.session_state:
+     st.session_state.ws_connected = False
+
+# --- Notification Processing ---
+new_notifications = []
+while not notification_queue.empty():
+    try:
+        msg = notification_queue.get_nowait()
+        new_notifications.append(msg)
+    except queue.Empty:
+        break
+
+if new_notifications:
+    logger.info(f"Processing {len(new_notifications)} notifications from queue.")
+    # Prepend new notifications to the session state list
+    current_list = st.session_state.notifications
+    st.session_state.notifications = new_notifications + current_list
+    # Limit history
+    if len(st.session_state.notifications) > 15:
+        st.session_state.notifications = st.session_state.notifications[:15]
+    # No explicit rerun needed here, Streamlit should rerun due to state change (?)
+    # or due to autorefresh below.
+
+# --- Auto Refresh ---
+# Refresh every 2 seconds to check the queue and update display
+count = st_autorefresh(interval=2000, limit=None, key="notifrefresh")
+
+# --- API Client ---
+client = httpx.AsyncClient(base_url=API_BASE_URL, timeout=10.0)
+
+# --- Helper Functions for API Calls ---
+async def api_register(email, password):
+    try:
+        response = await client.post("/register", json={"email": email, "password": password})
+        response.raise_for_status()
+        return {"success": True, "data": response.json()}
+    except httpx.HTTPStatusError as e:
+        detail = e.response.json().get("detail", e.response.text)
+        logger.error(f"API Register Error: {e.response.status_code} - {detail}")
+        return {"success": False, "error": f"API Error: {detail}"}
+    except Exception as e:
+        logger.exception("Register call failed")
+        return {"success": False, "error": f"Request failed: {e}"}
+
+async def api_login(email, password):
+    try:
+        response = await client.post("/login", json={"email": email, "password": password})
+        response.raise_for_status()
+        return {"success": True, "data": response.json()}
+    except httpx.HTTPStatusError as e:
+        detail = e.response.json().get("detail", e.response.text)
+        logger.error(f"API Login Error: {e.response.status_code} - {detail}")
+        return {"success": False, "error": f"API Error: {detail}"}
+    except Exception as e:
+        logger.exception("Login call failed")
+        return {"success": False, "error": f"Request failed: {e}"}
+
+# --- UI Rendering ---
+st.title("Authentication & Notification App (Streamlit)")
+
+if not st.session_state.logged_in:
+    st.sidebar.header("Login or Register")
+    login_tab, register_tab = st.sidebar.tabs(["Login", "Register"])
+
+    with login_tab:
+        with st.form("login_form"):
+            login_email = st.text_input("Email", key="login_email")
+            login_password = st.text_input("Password", type="password", key="login_password")
+            login_button = st.form_submit_button("Login")
+
+            if login_button:
+                if not login_email or not login_password:
+                    st.error("Please enter email and password.")
+                else:
+                    result = asyncio.run(api_login(login_email, login_password)) # Run async in sync context
+                    if result["success"]:
+                        token = result["data"]["access_token"]
+                        # Attempt to get user info immediately - needs modification if /users/me requires auth header
+                        # For simplicity, just store email from login form for now
+                        st.session_state.logged_in = True
+                        st.session_state.token = token
+                        st.session_state.user_email = login_email # Store email used for login
+                        st.session_state.notifications = [] # Clear old notifications
+
+                        # Start WebSocket listener thread
+                        stop_event.clear() # Ensure stop event is clear
+                        thread = threading.Thread(target=websocket_listener, args=(token,), daemon=True)
+                        st.session_state.ws_thread = thread
+                        thread.start()
+                        logger.info("Login successful, WS thread started.")
+                        st.rerun() # Rerun immediately to switch view
+                    else:
+                        st.error(f"Login failed: {result['error']}")
+
+    with register_tab:
+        with st.form("register_form"):
+            reg_email = st.text_input("Email", key="reg_email")
+            reg_password = st.text_input("Password", type="password", key="reg_password")
+            reg_confirm = st.text_input("Confirm Password", type="password", key="reg_confirm")
+            register_button = st.form_submit_button("Register")
+
+            if register_button:
+                if not reg_email or not reg_password or not reg_confirm:
+                    st.error("Please fill all fields.")
+                elif reg_password != reg_confirm:
+                    st.error("Passwords do not match.")
+                elif len(reg_password) < 8:
+                    st.error("Password must be at least 8 characters.")
+                else:
+                    result = asyncio.run(api_register(reg_email, reg_password))
+                    if result["success"]:
+                        st.success(f"Registration successful for {result['data']['email']}! Please log in.")
+                    else:
+                        st.error(f"Registration failed: {result['error']}")
+
+else: # Logged In View
+    st.sidebar.header(f"Welcome, {st.session_state.user_email}!")
+    if st.sidebar.button("Logout"):
+        logger.info("Logout requested.")
+        # Stop WebSocket thread
+        if st.session_state.ws_thread and st.session_state.ws_thread.is_alive():
+             logger.info("Signalling WS thread to stop.")
+             stop_event.set()
+             st.session_state.ws_thread.join(timeout=2.0) # Wait briefly for thread exit
+             if st.session_state.ws_thread.is_alive():
+                  logger.warning("WS thread did not exit cleanly.")
+        # Clear session state
+        st.session_state.logged_in = False
+        st.session_state.token = None
+        st.session_state.user_email = None
+        st.session_state.notifications = []
+        st.session_state.ws_thread = None
+        st.session_state.ws_connected = False
+        logger.info("Session cleared.")
+        st.rerun()
+
+    st.header("Dashboard")
+    # Display notifications
+    st.subheader("Real-time Notifications")
+    ws_status = "Connected" if st.session_state.ws_connected else "Disconnected"
+    st.caption(f"WebSocket Status: {ws_status}")
+
+    if st.session_state.notifications:
+        for i, msg in enumerate(st.session_state.notifications):
+            st.info(f"{msg}", icon="🔔")
+    else:
+        st.text("No new notifications.")
+
+    # Add a button to manually check queue/refresh if needed
+    # if st.button("Check for notifications"):
+    #      st.rerun() # Force rerun which includes queue check
+
+
+# --- Final Cleanup ---
+# Ensure httpx client is closed if script exits abnormally
+# (This might not always run depending on how Streamlit terminates)
+# Ideally handled within context managers if used more extensively
+# asyncio.run(client.aclose())