app.py

# /home/user/app/app.py
import streamlit as st
from pathlib import Path
from datetime import datetime # For default chat session title
from sqlmodel import select # For SQLModel queries

from config.settings import settings
# Ensure models are imported correctly based on your __init__.py in models/
from models import (
    create_db_and_tables,
    get_session_context, # Your SQLModel session context manager
    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
from assets.logo import get_logo_path # For consistent logo handling

# --- 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() # This uses SQLModel.metadata.create_all(engine)
    app_logger.info("Database initialized.")

init_db()

# --- Session State Initialization (Using primitive types for user auth) ---
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: # For the current active chat in Consult page
    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_main_app") # Unique key
        password_input = st.text_input("Password", type="password", key="login_password_input_main_app") # Unique key
        submit_button = st.form_submit_button("Login")

        if submit_button:
            user_object_from_auth = authenticate_user(username_input, password_input)

            if user_object_from_auth: # Indicates successful authentication
                st.success(f"Welcome back, {username_input}!")
                app_logger.info(f"User '{username_input}' authenticated successfully by authenticate_user.")

                try:
                    with get_session_context() as db_session: # db_session is a SQLModel Session
                        # Re-fetch the user to ensure we have a live object for this session
                        # and to get definitive ID/username if authenticate_user returned a detached object.
                        statement = select(User).where(User.username == username_input)
                        live_user = db_session.exec(statement).first()

                        if not live_user:
                            # This is a critical inconsistency if authenticate_user said OK.
                            st.error("Authentication inconsistency. User details not found after login. Please contact support or try again.")
                            app_logger.error(f"CRITICAL: User '{username_input}' authenticated but then NOT FOUND in DB by username query.")
                            # Clear potentially corrupted auth state
                            st.session_state.authenticated_user_id = None
                            st.session_state.authenticated_username = None
                            st.rerun() # Force rerun to show login page
                            return

                        app_logger.info(f"Live user object for '{live_user.username}' (ID: {live_user.id}) obtained in session.")
                        # Store primitive data in st.session_state
                        st.session_state.authenticated_user_id = live_user.id
                        st.session_state.authenticated_username = live_user.username
                        app_logger.info(f"Stored user ID {live_user.id} and username '{live_user.username}' in Streamlit session state.")

                        # Create a new chat session for the user upon login
                        app_logger.debug(f"Attempting to create new chat session for user_id: {live_user.id}")
                        new_chat_session = ChatSession(
                            user_id=live_user.id,
                            title=f"Session for {live_user.username} - {datetime.now().strftime('%Y-%m-%d %H:%M')}"
                            # start_time is default_factory=datetime.utcnow in the model
                        )
                        db_session.add(new_chat_session)
                        app_logger.debug("New ChatSession ORM object added to SQLModel session.")

                        # Flush the session to persist new_chat_session and get its ID
                        app_logger.debug("Flushing SQLModel session to assign ID to new_chat_session...")
                        db_session.flush()
                        app_logger.debug(f"Session flushed. new_chat_session potential ID: {new_chat_session.id}")

                        # Check if ID was assigned; if not, something is wrong with DB/model setup for auto-increment PK.
                        if new_chat_session.id is None:
                            app_logger.error("CRITICAL: new_chat_session.id is None after flush. Cannot proceed with session creation.")
                            raise Exception("Failed to obtain ID for new chat session after database flush.")

                        # Refresh the object to load any server-side defaults or confirm state
                        app_logger.debug(f"Refreshing new_chat_session (ID: {new_chat_session.id}) from database...")
                        db_session.refresh(new_chat_session)
                        app_logger.debug("new_chat_session ORM object refreshed.")

                        # Store the new chat session ID for use in the Consult page
                        st.session_state.current_chat_session_id = new_chat_session.id
                        st.session_state.chat_messages = [] # Clear any old messages for the new session
                        app_logger.info(f"New chat session (ID: {new_chat_session.id}) created and assigned for user '{live_user.username}'.")
                        # The commit will be handled by the get_session_context manager upon successful exit of this 'with' block.

                    # If all operations within the 'with' block succeed, the context manager will commit.
                    app_logger.info(f"Post-login setup complete for '{username_input}'. Rerunning app.")
                    st.rerun() # Rerun to reflect login state and navigate to the default page (e.g., Home)

                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 potentially partial auth state to force re-login or prevent inconsistent state
                    st.session_state.authenticated_user_id = None
                    st.session_state.authenticated_username = None
                    st.session_state.current_chat_session_id = None
                    # No rerun here, let user see the error and try again or contact support.
            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_main_app") # Unique key
        new_email = st.text_input("Email (Optional)", key="signup_email_input_main_app") # Unique key
        new_password = st.text_input("Choose a Password", type="password", key="signup_password_input_main_app") # Unique key
        confirm_password = st.text_input("Confirm Password", type="password", key="signup_confirm_password_input_main_app") # Unique key
        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 len(new_password) < 6: # Example: Minimum password length
                 st.error("Password must be at least 6 characters long.")
            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
                    # full_name can be added if part of UserCreate and signup form
                )
                # create_user_in_db should handle its own session and return a User object or None.
                # It's responsible for hashing password and committing the new user.
                user = create_user_in_db(user_data)
                if user:
                    # Use 'new_username' (from form input) for the message
                    st.success(f"Account created for {new_username}. Please log in.")
                    app_logger.info(f"Account created for '{new_username}'.")
                else:
                    # create_user_in_db returns None if user exists or another error occurred.
                    # The function itself should log the specific reason.
                    st.error("Username or Email might already be taken, or another error occurred during signup. Please check logs or try different credentials.")
                    app_logger.warning(f"Signup failed for username: '{new_username}'. create_user_in_db returned None.")

# --- Main App Logic (Checks for authenticated_user_id) ---
if not st.session_state.get("authenticated_user_id"): # Check if user_id is set (i.e., user is logged in)
    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 - Display Sidebar and Page Content ---
    with st.sidebar:
        # Display Logo using get_logo_path for consistency
        logo_path_str_sidebar = get_logo_path() # Assumes assets/logo.py is set up
        if logo_path_str_sidebar:
            logo_file_sidebar = Path(logo_path_str_sidebar)
            if logo_file_sidebar.exists():
                try:
                    st.image(str(logo_file_sidebar), width=100) # Adjust width as needed
                except Exception as e:
                    app_logger.warning(f"Could not display logo in sidebar from path '{logo_path_str_sidebar}': {e}")
            # else: # get_logo_path should log if file doesn't exist
            #    app_logger.warning(f"Sidebar logo path from get_logo_path() does not exist: {logo_path_str_sidebar}")
        elif settings.APP_TITLE: # Fallback to title if no logo path or logo load fails
             st.sidebar.markdown(f"#### {settings.APP_TITLE}")

        # Use the stored primitive username
        username_for_display = st.session_state.get("authenticated_username", "User") # Fallback to "User"
        st.sidebar.markdown(f"### Welcome, {username_for_display}!")
        st.sidebar.markdown("---") # Separator

        # Navigation is handled by Streamlit's multi-page app feature (pages/ directory)
        # This app.py acts as the entry point.
        # Streamlit automatically lists pages from the `pages/` directory here in the sidebar.

        if st.sidebar.button("Logout", key="sidebar_logout_button_main_app"): # Unique key
            logged_out_username = st.session_state.get("authenticated_username", "UnknownUser")
            app_logger.info(f"User '{logged_out_username}' logging out.")
            # Clear all relevant session state variables
            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 = []
            # Potentially clear other page-specific states if needed
            # for key in list(st.session_state.keys()):
            #     if key not in ['rerun_requested', ...]: # preserve essential Streamlit keys
            #         del st.session_state[key]
            st.success("You have been logged out.")
            st.rerun()

    # Main content area for app.py when logged in.
    # If you have a `pages/1_Home.py`, Streamlit will render that page's content here by default
    # when the user navigates to "Home" via the sidebar (or if it's the first page alphabetically).
    # The content below in app.py might only be visible briefly or if no page is explicitly selected,
    # or if `app.py` is configured as a page itself.

    # For a typical MPA setup, you might not need much content directly in app.py's main area
    # once the user is logged in, as page files will take over.
    # A simple instruction or pointer could be useful.
    st.sidebar.success("Select a page from the navigation.")
    # Optionally, add a default landing message in the main area if no page is loaded yet,
    # though Streamlit usually picks the first page in `pages/`.
    # st.info("Main application area. Select a page from the sidebar to begin.")

app_logger.info(f"Streamlit app '{settings.APP_TITLE}' (app.py) processed.")