app.py

# /home/user/app/app.py
import streamlit as st
from pathlib import Path
from sqlmodel import select # <--- IMPORT SELECT FOR SQLMODEL QUERIES

from config.settings import settings
from models import (
    create_db_and_tables,
    get_session_context, # Your SQLModel session context manager
    User,
    ChatMessage,
    ChatSession
)
from models.user import UserCreate
from services.auth import create_user_in_db, authenticate_user
from services.logger import app_logger
from assets.logo import get_logo_path

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

# --- Database Initialization --- (Assuming this is already correct)
@st.cache_resource
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 --- (Assuming this is already correct)
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:
            user_object_from_auth = authenticate_user(username_input, password_input)

            if user_object_from_auth:
                st.success(f"Welcome back, {username_input}!")
                app_logger.info(f"User {username_input} authenticated successfully.")

                try:
                    with get_session_context() as db_session: # db_session is a SQLModel Session
                        # --- SQLMODEL QUERY ---
                        statement = select(User).where(User.username == username_input)
                        live_user = db_session.exec(statement).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.")
                            st.session_state.authenticated_user_id = None
                            st.session_state.authenticated_username = None
                            st.rerun()
                            return

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

                        new_chat_session = ChatSession(user_id=live_user.id, title=f"Session for {live_user.username} ({Path(settings.APP_TITLE).name})")
                        db_session.add(new_chat_session)
                        # db_session.commit() # Handled by get_session_context manager
                        db_session.refresh(new_chat_session) # Refresh after add, before commit if ID needed
                        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}.")
                    # Session commit and close handled by get_session_context
                    st.rerun()

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

# display_signup_form() remains largely the same as it calls create_user_in_db,
# which should internally use SQLModel syntax if it interacts with the DB.

# --- Main App Logic --- (Remains the same, using updated session state keys)
# ... (rest of app.py as previously provided, ensuring it uses authenticated_user_id and authenticated_username)
# For brevity, I'm omitting the parts of app.py that don't change due to SQLModel query syntax.
# The signup form and the main structure checking st.session_state.get("authenticated_user_id")
# and the sidebar display are correct from the previous version.
# Ensure any other direct DB queries in app.py (if any were added) are also updated.

# (Make sure the rest of your app.py from the previous good version is here)
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 len(new_password) < 6:
                 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
                )
                user = create_user_in_db(user_data)
                if user:
                    st.success(f"Account created for {new_username}. Please log in.")
                    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}")

if not st.session_state.get("authenticated_user_id"):
    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:
    with st.sidebar:
        logo_path_str_sidebar = get_logo_path()
        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)
                except Exception as e:
                    app_logger.warning(f"Could not display logo in sidebar from path '{logo_path_str_sidebar}': {e}")
        elif settings.APP_TITLE:
             st.sidebar.markdown(f"#### {settings.APP_TITLE}")
        username_for_display = st.session_state.get("authenticated_username", "User")
        st.sidebar.markdown(f"### Welcome, {username_for_display}!")
        st.sidebar.markdown("---")
        if st.sidebar.button("Logout", key="sidebar_logout_button"):
            logged_out_username = st.session_state.get("authenticated_username", "UnknownUser")
            app_logger.info(f"User {logged_out_username} logging out.")
            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()
app_logger.info(f"Streamlit app '{settings.APP_TITLE}' initialized and running.")