# /home/user/app/app.py
import streamlit as st
from pathlib import Path
from datetime import datetime
from sqlmodel import select

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
from assets.logo import get_logo_path

# --- 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 ---
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 = []
if 'user_understands_disclaimer' not in st.session_state:
    st.session_state.user_understands_disclaimer = False

# --- Disclaimer Function ---
def display_disclaimer():
    st.warning(f"**Important Disclaimer:** {settings.MAIN_DISCLAIMER_LONG}")
    st.info(settings.SIMULATION_DISCLAIMER)
    if st.button("I Understand and Agree to Proceed", key="disclaimer_agree_btn"):
        st.session_state.user_understands_disclaimer = True
        st.rerun()

# --- Authentication Logic ---
def display_login_form():
    # (This function remains the same as the last full app.py provided, which had the fix for
    #  `Instance '<ChatSession ...>' is not persistent within this Session`.
    #  Ensure that version is used here for robustness.)
    with st.form("login_form"):
        st.subheader("Login")
        username_input = st.text_input("Username", key="login_username_input_main_app")
        password_input = st.text_input("Password", type="password", key="login_password_input_main_app")
        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:
                        statement = select(User).where(User.username == username_input)
                        live_user = db_session.exec(statement).first()
                        if not live_user:
                            st.error("Auth inconsistency. User not found post-login.")
                            app_logger.error(f"CRITICAL: User '{username_input}' auth OK but NOT FOUND in DB.")
                            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} - {datetime.now().strftime('%Y-%m-%d %H:%M')}"
                        )
                        db_session.add(new_chat_session)
                        app_logger.debug("New ChatSession added to SQLModel session.")
                        db_session.flush()
                        app_logger.debug(f"Session flushed. new_chat_session ID: {new_chat_session.id}")
                        if new_chat_session.id is None:
                            app_logger.error("CRITICAL: new_chat_session.id is None after flush.")
                            raise Exception("Failed to obtain ID for new chat session after flush.")
                        db_session.refresh(new_chat_session)
                        app_logger.debug("new_chat_session refreshed.")
                        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}'.")
                    app_logger.info(f"Post-login setup complete for '{username_input}'. Rerunning app.")
                    st.rerun()
                except Exception as e:
                    app_logger.error(f"Error during post-login session setup for '{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
                    st.session_state.current_chat_session_id = None
            else:
                st.error("Invalid username or password.")
                app_logger.warning(f"Failed login attempt for username: {username_input}")

def display_signup_form():
    # (This function remains largely the same as the last full app.py provided.
    #  Ensure it correctly calls your `create_user_in_db` which handles hashing and DB commit.)
    with st.form("signup_form"):
        st.subheader("Sign Up")
        new_username = st.text_input("Choose a Username", key="signup_username_input_main_app")
        new_email = st.text_input("Email (Optional)", key="signup_email_input_main_app")
        new_password = st.text_input("Choose a Password", type="password", key="signup_password_input_main_app")
        confirm_password = st.text_input("Confirm Password", type="password", key="signup_confirm_password_input_main_app")
        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/Email might be taken, or another error occurred. Check logs or try different credentials.")
                    app_logger.warning(f"Signup failed for '{new_username}'. create_user_in_db returned None.")

# --- Main App Logic ---
if not st.session_state.get("authenticated_user_id"): # Not logged in
    if not st.session_state.user_understands_disclaimer:
        st.title(f"Welcome to {settings.APP_TITLE}")
        st.markdown("Your AI-powered partner for advanced healthcare insights.")
        st.markdown("---")
        display_disclaimer()
        st.stop() # Stop further execution until disclaimer is agreed to
    else:
        # Disclaimer agreed, show login/signup
        st.title(f"Welcome to {settings.APP_TITLE}")
        st.caption(settings.MAIN_DISCLAIMER_SHORT) # Short disclaimer always visible on login page
        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 ---
    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 sidebar logo: {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.info(settings.MAIN_DISCLAIMER_SHORT) # Constant reminder
        st.sidebar.markdown("---")

        if st.sidebar.button("Logout", key="sidebar_logout_button_main_app"):
            logged_out_username = st.session_state.get("authenticated_username", "UnknownUser")
            app_logger.info(f"User '{logged_out_username}' logging out.")
            # Clear all session state, or selectively
            for key in list(st.session_state.keys()): # Clear most session state keys on logout
                if key not in ['query_params']: # Keep essential Streamlit keys if needed
                    del st.session_state[key]
            st.session_state.user_understands_disclaimer = False # Require disclaimer agreement again
            st.success("You have been logged out.")
            st.rerun()

    # Main content area: Streamlit MPA handles this by showing the selected page
    # from the `pages/` directory.
    # If on app.py itself, give a pointer.
    current_page = st.session_state.get('page_script_hash', None)
    if not current_page or current_page == st.elements.utils.calc_page_script_hash('app.py'):
         st.sidebar.success("Select a page from the navigation.")
         st.info("Please select a page from the sidebar (e.g., Home, Consult, Reports) to begin.")


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