import streamlit as st
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="⚕️", # You can use an emoji or a path to an image
    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 upon successful login
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")
        username = st.text_input("Username")
        password = st.text_input("Password", type="password")
        submit_button = st.form_submit_button("Login")

        if submit_button:
            user = authenticate_user(username, password)
            if user:
                st.session_state.authenticated_user = user
                st.success(f"Welcome back, {user.username}!")
                # Create a new chat session for the user upon login
                with get_session_context() as db_session:
                    new_chat_session = ChatSession(user_id=user.id, title=f"Session for {user.username}")
                    db_session.add(new_chat_session)
                    db_session.commit()
                    db_session.refresh(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
            else:
                st.error("Invalid username or password.")

def display_signup_form():
    with st.form("signup_form"):
        st.subheader("Sign Up")
        new_username = st.text_input("Choose a Username")
        new_email = st.text_input("Email (Optional)")
        new_password = st.text_input("Choose a Password", type="password")
        confirm_password = st.text_input("Confirm Password", type="password")
        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)
                if user:
                    st.success(f"Account created for {user.username}. Please log in.")
                    # Optionally log them in directly:
                    # st.session_state.authenticated_user = user
                    # st.rerun()
                else:
                    st.error("Username might already be taken or an error occurred.")

# --- 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, Streamlit automatically navigates to pages in the `pages/` directory.
    # The content of `app.py` typically acts as the "Home" page if no `1_Home.py` exists,
    # or it can be used for global elements like a custom sidebar if not using Streamlit's default page navigation.

    # Custom Sidebar for logged-in users (Streamlit handles page navigation automatically)
    with st.sidebar:
        st.markdown(f"### Welcome, {st.session_state.authenticated_user.username}!")
        st.image("assets/logo.png", width=100) # Display logo if available
        st.markdown("---")
        
        if st.button("Logout"):
            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()
    
    # This content will show if no other page is selected, or if you don't have a 1_Home.py
    # If you have 1_Home.py, Streamlit will show that by default after login.
    # So, this part might be redundant if 1_Home.py exists and is the intended landing.
    st.sidebar.success("Select a page above to get started.")
    st.markdown(f"# {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 see this, ensure you have a `pages/1_Home.py` or that this `app.py` is your intended landing page after login.")

app_logger.info("Streamlit app initialized and running.")