Update app.py

app.py

@@ -130,12 +130,60 @@ def login_callback():
     except Exception as e:
         st.error(f"Login failed: {e}")
 
-# Callback for logout
-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!")
+# Callback for login
+def login_callback():
+    login_identifier = st.session_state.login_identifier
+    password = st.session_state.login_password
+    try:
+        # Try to sign in the user programmatically to check the password validity
+        url = f'https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={FIREBASE_API_KEY}'
+        data = {
+            'email': login_identifier,
+            'password': password,
+            'returnSecureToken': True
+        }
+        response = requests.post(url, json=data)
+        result = response.json()
+
+        if 'idToken' in result:
+             # If sign in was successful, then use email to fetch the user
+             user = auth.get_user_by_email(login_identifier)
+             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!")
+
+        elif 'error' in result:
+           # If sign-in fails, retrieve user using display name
+            try:
+                user_list = auth.list_users()
+                for user_info in user_list.users:
+                  if user_info.display_name == login_identifier:
+                    user = user_info
+                    # If user is found using display name, try signing in using email
+                    url = f'https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={FIREBASE_API_KEY}'
+                    data = {
+                        'email': user.email,
+                        'password': password,
+                        'returnSecureToken': True
+                    }
+                    response = requests.post(url, json=data)
+                    result = response.json()
+                    if 'idToken' in result:
+                      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!")
+                      return
+                
+                raise Exception("User not found with provided credentials.")  # if not found, raise exception.
+            except Exception as e:
+                st.error(f"Login failed: {e}") # if any error, display this message.
+        else:
+            raise Exception("Error with sign-in endpoint") # If sign-in endpoint doesn't return error or id token, then throw this error.
+
+    except Exception as e:
+        st.error(f"Login failed: {e}")
 
 # Function to get image from url
 def get_image_from_url(url):