import streamlit as st from pathlib import Path from sqlalchemy.orm.exc import DetachedInstanceError as SQLAlchemyDetachedInstanceError # For specific error catching if needed from config.settings import settings from models import create_db_and_tables, get_session_context, User, ChatMessage, ChatSession from models.user import UserCreate from services.auth import create_user_in_db, authenticate_user from services.logger import app_logger # --- Page Configuration --- st.set_page_config( page_title=settings.APP_TITLE, page_icon="⚕️", layout="wide", initial_sidebar_state="expanded" ) # --- Database Initialization --- @st.cache_resource def init_db(): app_logger.info("Initializing database and tables...") create_db_and_tables() app_logger.info("Database initialized.") init_db() # --- Session State Initialization (Using new keys) --- if 'authenticated_user_id' not in st.session_state: st.session_state.authenticated_user_id = None if 'authenticated_username' not in st.session_state: st.session_state.authenticated_username = None if 'current_chat_session_id' not in st.session_state: st.session_state.current_chat_session_id = None if 'chat_messages' not in st.session_state: st.session_state.chat_messages = [] # --- Authentication Logic --- def display_login_form(): with st.form("login_form"): st.subheader("Login") username_input = st.text_input("Username", key="login_username_input") password_input = st.text_input("Password", type="password", key="login_password_input") submit_button = st.form_submit_button("Login") if submit_button: # authenticate_user should ideally return a User object with at least ID loaded, # or be robust enough that its session handling doesn't cause immediate detachment issues # if we were to access its attributes directly (though we now avoid that for long-term storage). user_object_from_auth = authenticate_user(username_input, password_input) if user_object_from_auth: st.success(f"Welcome back, {username_input}!") # Use form input for immediate message app_logger.info(f"User {username_input} authenticated successfully.") try: with get_session_context() as db_session: # Fetch the "live" user to get definite ID and current username, # and to ensure it's attached to *this* session for creating the chat session. # The primary fix for 'authenticate_user' is still for it to return a usable object, # but re-fetching here ensures data for st.session_state is from a controlled point. live_user = db_session.query(User).filter(User.username == username_input).first() if not live_user: st.error("Authentication inconsistency. User details not found after login. Please contact support.") app_logger.error(f"CRITICAL: User '{username_input}' authenticated but then not found in DB by username.") # Clear any potentially partially set auth state from previous attempts st.session_state.authenticated_user_id = None st.session_state.authenticated_username = None st.rerun() return # IMPORTANT: Store primitive data in st.session_state # Do this *before* the commit for new_chat_session, as commit might expire 'live_user' attributes. st.session_state.authenticated_user_id = live_user.id st.session_state.authenticated_username = live_user.username # Username from DB app_logger.info(f"Stored user ID {live_user.id} and username '{live_user.username}' in session state.") # Now create the chat session using the live_user's ID new_chat_session = ChatSession(user_id=live_user.id, title=f"Session for {live_user.username}") db_session.add(new_chat_session) db_session.commit() # This commit is for new_chat_session # After commit, live_user attributes might be expired if expire_on_commit=True. # But we've already stored the primitives we need. db_session.refresh(new_chat_session) # Get ID for new_chat_session st.session_state.current_chat_session_id = new_chat_session.id st.session_state.chat_messages = [] app_logger.info(f"New chat session (ID: {new_chat_session.id}) created for user {live_user.username}.") st.rerun() # Rerun to reflect login state and navigate except Exception as e: app_logger.error(f"Error during post-login session setup for user {username_input}: {e}", exc_info=True) st.error(f"Could not complete login process: {e}") # Clear auth state on error st.session_state.authenticated_user_id = None st.session_state.authenticated_username = None else: st.error("Invalid username or password.") app_logger.warning(f"Failed login attempt for username: {username_input}") def display_signup_form(): with st.form("signup_form"): st.subheader("Sign Up") new_username = st.text_input("Choose a Username", key="signup_username_input") new_email = st.text_input("Email (Optional)", key="signup_email_input") new_password = st.text_input("Choose a Password", type="password", key="signup_password_input") confirm_password = st.text_input("Confirm Password", type="password", key="signup_confirm_password_input") submit_button = st.form_submit_button("Sign Up") if submit_button: if not new_username or not new_password: st.error("Username and password are required.") elif new_password != confirm_password: st.error("Passwords do not match.") else: user_data = UserCreate( username=new_username, password=new_password, email=new_email if new_email else None ) user = create_user_in_db(user_data) # Should handle its own session and return User or None if user: st.success(f"Account created for {new_username}. Please log in.") # Use form input app_logger.info(f"Account created for {new_username}.") else: st.error("Username might already be taken or another error occurred during signup.") app_logger.warning(f"Signup failed for username: {new_username}") # --- Main App Logic (Checks for authenticated_user_id) --- if not st.session_state.get("authenticated_user_id"): # Check if user_id is set st.title(f"Welcome to {settings.APP_TITLE}") st.markdown("Your AI-powered partner for advanced healthcare insights.") login_tab, signup_tab = st.tabs(["Login", "Sign Up"]) with login_tab: display_login_form() with signup_tab: display_signup_form() else: # User is authenticated (authenticated_user_id is present) with st.sidebar: # Use the stored primitive username username_for_display = st.session_state.get("authenticated_username", "User") st.markdown(f"### Welcome, {username_for_display}!") logo_path_str = getattr(settings, "LOGO_PATH", None) if logo_path_str: logo_path = Path(logo_path_str) if logo_path.exists(): try: st.image(str(logo_path), width=100) except Exception as e: app_logger.warning(f"Could not load logo from {logo_path_str}: {e}") else: app_logger.warning(f"Logo path specified but does not exist: {logo_path_str}") elif settings.APP_TITLE: st.markdown(f"#### {settings.APP_TITLE}") st.markdown("---") if st.button("Logout"): logged_out_username = st.session_state.get("authenticated_username", "UnknownUser") app_logger.info(f"User {logged_out_username} logging out.") # Clear authentication state st.session_state.authenticated_user_id = None st.session_state.authenticated_username = None st.session_state.current_chat_session_id = None st.session_state.chat_messages = [] st.success("You have been logged out.") st.rerun() st.sidebar.success("Select a page from the navigation.") st.header(f"Dashboard - {settings.APP_TITLE}") st.markdown("Navigate using the sidebar to consult with the AI or view your reports.") st.markdown("---") st.info("This is the main application area. If you have pages in a `pages/` directory, Streamlit will show the selected page here. Otherwise, this content is shown.") app_logger.info(f"Streamlit app '{settings.APP_TITLE}' initialized and running.")