app.py

import streamlit as st
from pathlib import Path # Ensure Path is imported if used for settings.LOGO_PATH

from config.settings import settings
from models import create_db_and_tables, get_session_context, User, ChatMessage, ChatSession
from models.user import UserCreate # For type hinting
from services.auth import create_user_in_db, authenticate_user
from services.logger import app_logger
# Agent will be initialized and used in specific pages like Consult.py
# from agent import get_agent_executor # Import in pages where needed

# --- Page Configuration ---
st.set_page_config(
    page_title=settings.APP_TITLE,
    page_icon="⚕️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Database Initialization ---
@st.cache_resource # Ensure this runs only once
def init_db():
    app_logger.info("Initializing database and tables...")
    create_db_and_tables()
    app_logger.info("Database initialized.")

init_db()

# --- Session State Initialization ---
if 'authenticated_user' not in st.session_state:
    st.session_state.authenticated_user = None # Stores User object (ideally one that is "live" or has loaded attributes)
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: # For the current active chat
    st.session_state.chat_messages = []


# --- Authentication Logic ---
def display_login_form():
    with st.form("login_form"):
        st.subheader("Login")
        # Use unique keys for inputs if there's any chance of collision or for better state tracking
        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:
            # IMPORTANT: 'authenticate_user' should return a User object where essential attributes
            # like 'id' and 'username' are already loaded (not expired).
            # If 'authenticate_user' involves a commit (e.g., updating last_login) and
            # session.expire_on_commit=True (default), it should call db.refresh(user_object)
            # before closing its session and returning the user object.
            user_from_auth = authenticate_user(username_input, password_input)

            if user_from_auth:
                # FIX 1: For the success message, use the 'username_input' from the form directly.
                # This avoids accessing 'user_from_auth.username' which might be on a detached/expired instance.
                st.success(f"Welcome back, {username_input}!")
                app_logger.info(f"User {username_input} logged in successfully.")

                try:
                    with get_session_context() as db_session:
                        # FIX 2: Re-fetch or merge the user into the current db_session to ensure it's "live"
                        # This assumes 'user_from_auth.id' is accessible (i.e., was loaded by authenticate_user).
                        # If 'user_from_auth.id' itself is expired, 'authenticate_user' MUST be fixed.
                        try:
                            user_id = user_from_auth.id
                        except AttributeError:
                            st.error("Authentication error: User data is incomplete. Please try again.")
                            app_logger.error(f"User object from authenticate_user for '{username_input}' lacks 'id' attribute or it's inaccessible.")
                            st.session_state.authenticated_user = None # Clear potentially bad state
                            st.rerun()
                            return

                        live_user = db_session.get(User, user_id)
                        if not live_user:
                            st.error("Critical user session error. Please log out and log in again.")
                            app_logger.error(f"Failed to re-fetch user with id {user_id} (username: {username_input}) in new session.")
                            st.session_state.authenticated_user = None # Clear broken state
                            st.rerun()
                            return

                        # Store the "live" user object (attached to db_session or with loaded attributes)
                        # in session_state. This 'live_user' will be used by the sidebar.
                        st.session_state.authenticated_user = live_user

                        # Now use 'live_user' for creating the chat session
                        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()
                        db_session.refresh(new_chat_session) # Refresh to get DB-generated values for new_chat_session
                        st.session_state.current_chat_session_id = new_chat_session.id
                        st.session_state.chat_messages = [] # Clear previous messages
                    st.rerun() # Rerun to reflect login state and navigate
                except Exception as e:
                    app_logger.error(f"Error creating chat session for user {username_input}: {e}", exc_info=True)
                    st.error(f"Could not start a new session: {e}")
                    # Optionally, clear authenticated_user if session creation is critical
                    # st.session_state.authenticated_user = None
                    # st.rerun()
            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
                )
                # 'create_user_in_db' should return a User object with attributes loaded
                # (e.g., by calling session.refresh(user) before its session closes).
                user = create_user_in_db(user_data)
                if user:
                    # Use 'new_username' (from form input) for the message to avoid DetachedInstanceError
                    st.success(f"Account created for {new_username}. Please log in.")
                    app_logger.info(f"Account created for {new_username}.")
                    # Optionally log them in directly:
                    # To do this safely, 'user' returned by create_user_in_db must be "live"
                    # or its attributes fully loaded. You might need to re-fetch it in a new session here
                    # if you intend to store it in st.session_state.authenticated_user.
                    # For now, redirecting to login is simpler.
                else:
                    st.error("Username might already be taken or an error occurred during signup.")
                    app_logger.warning(f"Signup failed for username: {new_username}")

# --- Main App Logic ---
if not st.session_state.authenticated_user:
    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:
    # If authenticated
    with st.sidebar:
        # Accessing st.session_state.authenticated_user.username.
        # This relies on the object stored in authenticated_user having its username loaded.
        # The `display_login_form` now stores the 'live_user' which should have this.
        try:
            st.markdown(f"### Welcome, {st.session_state.authenticated_user.username}!")
        except AttributeError as e: # Fallback if username somehow isn't accessible
             st.markdown(f"### Welcome!")
             app_logger.error(f"Could not access username for authenticated_user in sidebar: {e}. User object: {st.session_state.authenticated_user}")


        # Example for conditional logo 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: # Fallback to title if no logo
             st.markdown(f"#### {settings.APP_TITLE}")

        st.markdown("---")

        if st.button("Logout"):
            if hasattr(st.session_state.authenticated_user, 'username'):
                app_logger.info(f"User {st.session_state.authenticated_user.username} logging out.")
            else:
                app_logger.info("User logging out (username not accessible from session state object).")
            st.session_state.authenticated_user = None
            st.session_state.current_chat_session_id = None
            st.session_state.chat_messages = []
            st.success("You have been logged out.")
            st.rerun()

    # Main area content (could be your "Home" page if no `pages/1_Home.py`)
    st.sidebar.success("Select a page from the navigation.") # More generic message
    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.")