Update app.py

app.py

@@ -25,6 +25,9 @@ if "logged_in" not in st.session_state:
     st.session_state.logged_in = False
 if "current_user" not in st.session_state:
     st.session_state.current_user = None
+if "display_name" not in st.session_state:
+    st.session_state.display_name = None
+
 
 def send_verification_email(id_token):
     url = f'https://identitytoolkit.googleapis.com/v1/accounts:sendOobCode?key={FIREBASE_API_KEY}'
@@ -41,15 +44,22 @@ def send_verification_email(id_token):
         return {'status': 'error', 'message': result['error']['message']}
     else:
         return {'status': 'success', 'email': result['email']}
-        
+
 # Callback for registration
 def register_callback():
     email = st.session_state.reg_email
     password = st.session_state.reg_password
+    display_name = st.session_state.reg_display_name
     try:
+        # Step 1: Create a new user in Firebase
         user = auth.create_user(email=email, password=password)
+        
+        # Step 2: Update the user profile with the display name
+        auth.update_user(user.uid, display_name=display_name)
+        
         st.success("Registration successful! Sending verification email...")
 
+        # Step 3: Sign in the user programmatically to get the id_token
         url = f'https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={FIREBASE_API_KEY}'
         data = {
             'email': email,
@@ -75,12 +85,25 @@ def register_callback():
 
 # Callback for login
 def login_callback():
-    email = st.session_state.login_email
+    login_identifier = st.session_state.login_identifier
     password = st.session_state.login_password
     try:
-        user = auth.get_user_by_email(email)
+        # Try to get the user by email first
+        try:
+          user = auth.get_user_by_email(login_identifier)
+        except:
+            # If getting by email fails, try to get the user by display name
+            user_list = auth.list_users()
+            for user_info in user_list.users:
+                if user_info.display_name == login_identifier:
+                    user = user_info
+                    break
+            else:
+                raise Exception("User not found")
+       
         st.session_state.logged_in = True
         st.session_state.current_user = user.uid
+        st.session_state.display_name = user.display_name # Store the display name
         st.success("Logged in successfully!")
     except Exception as e:
         st.error(f"Login failed: {e}")
@@ -89,6 +112,7 @@ def login_callback():
 def logout_callback():
     st.session_state.logged_in = False
     st.session_state.current_user = None
+    st.session_state.display_name = None
     st.info("Logged out successfully!")
 
 # Function to store text in Firebase Realtime Database
@@ -109,6 +133,7 @@ def registration_form():
     with st.form("Registration"):
         st.subheader("Register")
         email = st.text_input("Email", key="reg_email")
+        display_name = st.text_input("Display Name", key="reg_display_name")
         password = st.text_input("Password (min 6 characters)", type="password", key="reg_password")
         submit_button = st.form_submit_button("Register", on_click=register_callback)
 
@@ -116,13 +141,13 @@ def registration_form():
 def login_form():
     with st.form("Login"):
         st.subheader("Login")
-        email = st.text_input("Email", key="login_email")
+        login_identifier = st.text_input("Email or Username", key="login_identifier")
         password = st.text_input("Password", type="password", key="login_password")
         submit_button = st.form_submit_button("Login", on_click=login_callback)
 
 # Main app screen (after login)
 def main_app():
-    st.subheader(f"Welcome, {st.session_state.current_user}!")
+    st.subheader(f"Welcome, {st.session_state.display_name}!")
     st.write("Enter a name below to get a greeting.")
 
     # Input field for name