main.py

import streamlit as st
from dashboard_components import dashboard_header, account_linking_form, apply_custom_css
from utils.security import clear_metrics

def logout_user():
    """Properly handle user logout."""
    try:
        # Clear metrics before session cleanup
        clear_metrics()
        # Clear session state
        st.session_state.clear()
        # Rerun the app
        st.rerun()
    except Exception as e:
        st.error(f"Logout error: {str(e)}")

def main():
    """Main application entry point."""
    # Initialize session state
    if 'page' not in st.session_state:
        st.session_state.page = 'login'
    
    # Apply custom CSS
    apply_custom_css()
    
    # Show different pages based on state
    if st.session_state.page == 'login':
        show_login_page()
    elif st.session_state.page == 'dashboard':
        # Show header with proper logout callback
        dashboard_header(logout_callback=logout_user)
        show_dashboard()
    elif st.session_state.page == 'link_accounts':
        dashboard_header(logout_callback=logout_user)
        account_linking_form()
    elif st.session_state.page == 'wallets':
        dashboard_header(logout_callback=logout_user)
        show_wallets_page()

if __name__ == "__main__":
    main()