File size: 9,416 Bytes
6145bc0
79d193c
 
e9c724f
864b488
0b77233
79d193c
0b77233
 
864b488
0b77233
 
 
e9c724f
0b77233
 
 
864b488
0b77233
79d193c
0b77233
 
 
 
64f8a92
0b77233
0a2437f
79d193c
 
 
 
 
0b77233
 
c21fd02
0b77233
0a2437f
0b77233
 
 
 
e9c724f
 
0b77233
64f8a92
0b77233
79d193c
 
 
 
e9c724f
79d193c
 
 
e9c724f
04ae885
 
79d193c
 
 
 
c21fd02
e9c724f
04ae885
79d193c
 
 
 
 
 
04ae885
 
79d193c
 
 
 
 
e9c724f
79d193c
04ae885
 
79d193c
 
 
 
04ae885
79d193c
c21fd02
79d193c
e9c724f
c21fd02
04ae885
79d193c
 
 
 
 
0b77233
 
e9c724f
64f8a92
0b77233
 
 
e9c724f
 
 
 
0b77233
64f8a92
0b77233
 
 
 
 
 
 
 
 
 
 
79d193c
0b77233
79d193c
04ae885
0b77233
c21fd02
e9c724f
0a2437f
79d193c
 
0b77233
 
04ae885
0b77233
 
 
 
 
 
79d193c
0b77233
79d193c
 
 
e9c724f
 
 
 
 
 
 
 
 
 
 
c21fd02
04ae885
 
0b77233
04ae885
0b77233
79d193c
 
 
 
 
0b77233
 
 
 
04ae885
c21fd02
04ae885
0b77233
 
e9c724f
04ae885
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
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.")