Update main.py

main.py

@@ -68,44 +68,122 @@ except Exception as e:
     traceback.print_exc()
 # --- END Firebase Initialization ---
 
+
+# --- Helper function to ensure user profile in Realtime Database ---
+def ensure_tunasonga_profile_in_rtdb(uid, email, display_name, is_admin_flag):
+    """
+    Ensures a user profile exists in Realtime DB for Tunasonga Agri.
+    Creates or updates it with the is_admin flag and other defaults.
+    """
+    user_ref = db.reference(f'users/{uid}')
+    user_data = user_ref.get()
+
+    if not user_data:
+        print(f"Profile for UID {uid} (Email: {email}) not found in RTDB. Creating new profile.")
+        new_profile = {
+            'email': email,
+            'name': display_name,
+            'phone_number': "", # Default, can be updated by user
+            'location': "",     # Default, can be updated by user
+            'roles': {          # Default roles
+                'farmer': False,
+                'buyer': False,
+                'transporter': False
+            },
+            'role_applications': {},
+            'document_references': {},
+            'is_admin': is_admin_flag, # Set based on the script's input
+            'is_facilitator': False, # Default
+            'created_at': datetime.now(timezone.utc).isoformat(),
+            'suspended': False
+        }
+        try:
+            user_ref.set(new_profile)
+            print(f"Successfully created RTDB profile for {email} with is_admin: {is_admin_flag}")
+        except Exception as e_set:
+            print(f"Error setting RTDB profile for {email}: {e_set}")
+    else:
+        # Profile exists, ensure 'is_admin' flag is correctly set/updated
+        if user_data.get('is_admin') != is_admin_flag:
+            print(f"Updating is_admin flag for {email} in RTDB to: {is_admin_flag}")
+            try:
+                user_ref.update({'is_admin': is_admin_flag})
+                print(f"Successfully updated is_admin for {email} in RTDB.")
+            except Exception as e_update:
+                print(f"Error updating is_admin for {email} in RTDB: {e_update}")
+        else:
+            print(f"RTDB profile for {email} already exists and is_admin is correctly set to: {is_admin_flag}.")
+        
+        # Ensure other essential fields exist if the profile was created by a different process
+        # (e.g., if user signed up via app before this script ran)
+        updates_needed = {}
+        if 'name' not in user_data: updates_needed['name'] = display_name
+        if 'roles' not in user_data: updates_needed['roles'] = {'farmer': False, 'buyer': False, 'transporter': False}
+        if 'created_at' not in user_data: updates_needed['created_at'] = datetime.now(timezone.utc).isoformat()
+        # Add other default fields if necessary
+
+        if updates_needed:
+            print(f"Ensuring essential fields for existing profile of {email} in RTDB.")
+            try:
+                user_ref.update(updates_needed)
+                print(f"Successfully updated essential fields for {email} in RTDB.")
+            except Exception as e_essential_update:
+                 print(f"Error updating essential fields for {email} in RTDB: {e_essential_update}")
+
+
+# --- User Management Logic ---
 users_to_manage = [
-    # These two will now be the admin users
-    {"email": "rairorr@gmail.com", "password": "tuna2025!", "display_name": "Rairo", "is_admin": True},
+    {"email": "rairorr@gmail.com", "password": "tuna2025!", "display_name": "Rairo Admin", "is_admin": True},
     {"email": "gwatidzomisheck@gmail.com", "password": "tuna2025!", "display_name": "Gwatidzo Admin", "is_admin": True},
-    # The original user will no longer have the admin claim set by this script
-
+    # Add other users here if needed, e.g., to remove admin status from a previous admin
+    # {"email": "oldadmin@example.com", "password": "somepassword", "display_name": "Old Admin", "is_admin": False},
 ]
 
-# ─── Ensure specified users exist and set/unset admin claims ────────────────────
+print("\n--- Starting User and Admin Claim Management ---")
 for user_data in users_to_manage:
     email = user_data["email"]
-    password = user_data["password"]
+    password = user_data["password"] # Password is only used if creating a new Auth user
     display_name = user_data.get("display_name", email.split('@')[0])
-    is_admin = user_data.get("is_admin", False)
+    is_admin_flag_for_script = user_data.get("is_admin", False) # Renamed to avoid conflict
+
+    print(f"\nProcessing user: {email} (Target admin status: {is_admin_flag_for_script})")
 
     try:
-        # Attempt to retrieve the user by their email
         current_user = auth.get_user_by_email(email)
-        # print(f"User {email} already exists.") # Optional: uncomment for feedback
-    except firebase_admin._auth_utils.UserNotFoundError:
-        # If the user is not found, create a new one
-        current_user = auth.create_user(
-            email=email,
-            email_verified=True,
-            password=password,
-            display_name=display_name
-        )
-        # print(f"Created new user: {email}") # Optional: uncomment for feedback
-
-    # Set or unset custom user claims based on 'is_admin' flag
-    if is_admin:
-        auth.set_custom_user_claims(current_user.uid, {"admin": True})
-        # print(f"Set admin claim for {email}") # Optional: uncomment for feedback
-    else:
-        # If 'is_admin' is False, ensure the admin claim is removed (or not set)
-        # It's good practice to explicitly remove if it might have been set previously
-        auth.set_custom_user_claims(current_user.uid, {}) # Clear all custom claims, or set specific ones to False
-        # print(f"Removed admin claim for {email}") # Optional: uncomment for feedback
+        print(f"Firebase Auth: User {email} (UID: {current_user.uid}) already exists.")
+    except firebase_admin.auth.UserNotFoundError: # Corrected exception type
+        print(f"Firebase Auth: User {email} not found. Creating new user...")
+        try:
+            current_user = auth.create_user(
+                email=email,
+                email_verified=True, # You can set this as needed
+                password=password,
+                display_name=display_name
+            )
+            print(f"Firebase Auth: Created new user {email} (UID: {current_user.uid})")
+        except Exception as e_create:
+            print(f"Firebase Auth: Error creating user {email}: {e_create}")
+            continue # Skip to the next user if creation fails
+
+    # Set or unset custom user claims for Firebase Auth
+    current_claims = current_user.custom_claims or {}
+    if is_admin_flag_for_script:
+        if not current_claims.get("admin"):
+            auth.set_custom_user_claims(current_user.uid, {"admin": True})
+            print(f"Firebase Auth: Set admin custom claim for {email}.")
+        else:
+            print(f"Firebase Auth: Admin custom claim already set for {email}.")
+    else: # is_admin_flag_for_script is False
+        if current_claims.get("admin"):
+            auth.set_custom_user_claims(current_user.uid, {}) # Remove admin claim
+            print(f"Firebase Auth: Removed admin custom claim for {email}.")
+        else:
+            print(f"Firebase Auth: Admin custom claim already not set (or removed) for {email}.")
+
+    # Ensure/Update profile in Realtime Database
+    ensure_tunasonga_profile_in_rtdb(current_user.uid, email, display_name, is_admin_flag_for_script)
+
+print("\n--- User and Admin Claim Management Complete ---")
 
 # --- Helper Functions ---
 def verify_token(auth_header):