app.py

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.")