Update main.py

main.py

@@ -60,6 +60,16 @@ def verify_admin_and_get_uid(auth_header):
         raise PermissionError('Admin access required')
     return uid
 
+# Helper function for migration
+# Place this helper function somewhere accessible, e.g., before the admin endpoints section.
+def copy_collection(source_ref, dest_ref):
+    """Recursively copies documents and sub-collections."""
+    docs = source_ref.stream()
+    for doc in docs:
+        dest_ref.document(doc.id).set(doc.to_dict())
+        for sub_coll_ref in doc.reference.collections():
+            copy_collection(sub_coll_ref, dest_ref.document(doc.id).collection(sub_coll_ref.id))
+
 # -----------------------------------------------------------------------------
 # 3. AUTHENTICATION & USER MANAGEMENT ENDPOINTS
 # -----------------------------------------------------------------------------
@@ -271,7 +281,8 @@ def get_user_dashboard():
 def update_user_profile():
     """
     Allows a user to update their profile. If the phone number is changed,
-    it validates for uniqueness and then stages the change for migration.
+    it validates for uniqueness and stages the change based on user's choice
+    (migrate or start fresh).
     """
     uid = verify_token(request.headers.get('Authorization'))
     if not uid: return jsonify({'error': 'Invalid or expired token'}), 401
@@ -290,36 +301,39 @@ def update_user_profile():
     if new_display_name and new_display_name.strip() != current_user_data.get('displayName'):
         update_data['displayName'] = new_display_name.strip()
 
-    # --- VALIDATION AND MIGRATION LOGIC START ---
+    # --- REWORKED LOGIC FOR PHONE CHANGE ---
     new_phone = data.get('phone')
     if new_phone:
         new_phone_stripped = new_phone.strip()
         current_phone = current_user_data.get('phone')
 
         if new_phone_stripped != current_phone:
-            # --- THE FIX IS HERE: VALIDATE UNIQUENESS ---
-            # Query to see if any OTHER user already has this phone number.
+            # Get the user's choice: 'migrate' or 'start_fresh'
+            action = data.get('phoneChangeAction')
+            if action not in ['migrate', 'start_fresh']:
+                return jsonify({'error': "A choice ('migrate' or 'start_fresh') is required when changing phone numbers."}), 400
+
+            # Validate that the new phone number is not already in use
             existing_user_query = db.collection('users').where('phone', '==', new_phone_stripped).limit(1).stream()
-            
-            # Convert stream to a list to check its contents
-            conflicting_users = list(existing_user_query)
-
-            if len(conflicting_users) > 0:
-                # IMPORTANT: We need to make sure the conflicting user isn't the current user themselves
-                # This can happen in rare edge cases.
-                if conflicting_users[0].id != uid:
-                    logger.warning(f"User {uid} tried to claim phone number {new_phone_stripped}, but it's already used by user {conflicting_users[0].id}.")
-                    return jsonify({'error': 'This phone number is already registered to another account.'}), 409 # 409 Conflict is a good status code here
-            # --- END OF VALIDATION ---
-
-            # If validation passes, stage the phone number for migration.
+            if len(list(existing_user_query)) > 0:
+                return jsonify({'error': 'This phone number is already registered to another account.'}), 409
+
+            # Stage the change based on the user's chosen action
             update_data['migration_data'] = {
                 'from_phone': current_phone,
-                'to_phone': new_phone_stripped
+                'to_phone': new_phone_stripped,
+                'action': action  # Store the user's choice
             }
-            update_data['phoneStatus'] = 'pending_migration'
-            logging.info(f"User {uid} initiated phone migration from {current_phone} to {new_phone_stripped}. Awaiting admin approval.")
-    # --- VALIDATION AND MIGRATION LOGIC END ---
+            
+            # Use a descriptive status for the admin
+            if action == 'migrate':
+                update_data['phoneStatus'] = 'pending_migration'
+            else: # action == 'start_fresh'
+                update_data['phoneStatus'] = 'pending_fresh_start'
+            
+            logging.info(f"User {uid} initiated phone change to {new_phone_stripped} with action '{action}'. Awaiting admin approval.")
+            
+    # --- END OF REWORKED LOGIC ---
     
     if not update_data:
         return jsonify({'message': 'No changes detected.'}), 200
@@ -332,7 +346,7 @@ def update_user_profile():
         
         updated_user_doc = user_ref.get()
         response_message = 'Profile updated successfully.'
-        if 'phoneStatus' in update_data and update_data['phoneStatus'] == 'pending_migration':
+        if 'phoneStatus' in update_data:
             response_message += ' Your request to change your phone number has been submitted for approval.'
         
         return jsonify({
@@ -517,77 +531,81 @@ def approve_user_phone():
         logging.error(f"Admin approval failed for user {data.get('uid')}: {e}")
         return jsonify({'error': 'An internal error occurred'}), 500
 
-# Helper function for migration
-def copy_collection(source_ref, dest_ref):
-    """Recursively copies documents and sub-collections."""
-    docs = source_ref.stream()
-    for doc in docs:
-        # Copy the document data
-        dest_ref.document(doc.id).set(doc.to_dict())
-        # Recursively copy sub-collections
-        for sub_coll_ref in doc.reference.collections():
-            copy_collection(sub_coll_ref, dest_ref.document(doc.id).collection(sub_coll_ref.id))
 
-@app.route('/api/admin/users/approve-migration', methods=['POST'])
-def approve_phone_migration():
-    """Admin: Approves a phone number change and migrates all associated bot data."""
+@app.route('/api/admin/users/approve-phone-change', methods=['POST'])
+def approve_phone_change():
+    """
+    Admin: Approves a phone number change. Handles both 'migrate' and 'start_fresh'
+    scenarios based on the user's stored choice.
+    """
     try:
         verify_admin_and_get_uid(request.headers.get('Authorization'))
         data = request.get_json()
         target_uid = data.get('uid')
         if not target_uid:
-            return jsonify({'error': 'User UID is required for migration approval'}), 400
+            return jsonify({'error': 'User UID is required for approval'}), 400
 
         user_ref = db.collection('users').document(target_uid)
         user_doc = user_ref.get()
         if not user_doc.exists: return jsonify({'error': 'User not found'}), 404
 
         user_data = user_doc.to_dict()
-        if user_data.get('phoneStatus') != 'pending_migration':
-            return jsonify({'error': 'User is not awaiting a phone migration'}), 400
+        status = user_data.get('phoneStatus')
+        if status not in ['pending_migration', 'pending_fresh_start']:
+            return jsonify({'error': f'User is not awaiting a phone change (status is {status})'}), 400
 
         migration_data = user_data.get('migration_data')
-        if not migration_data or 'from_phone' not in migration_data or 'to_phone' not in migration_data:
+        if not migration_data or 'to_phone' not in migration_data or 'action' not in migration_data:
             return jsonify({'error': 'Invalid migration data in user profile'}), 500
             
-        from_phone_id = migration_data['from_phone'].lstrip('+')
-        to_phone_id = migration_data['to_phone'].lstrip('+')
-
-        source_ref = db.collection('users').document(from_phone_id)
-        dest_ref = db.collection('users').document(to_phone_id)
-
-        # --- DATA MIGRATION ---
-        # This can be a long operation. For a production app, this should be
-        # offloaded to a background task (e.g., Cloud Task or Pub/Sub).
-        logging.info(f"Starting data migration from {from_phone_id} to {to_phone_id}")
-        source_doc = source_ref.get()
-        if source_doc.exists:
-             # Copy the main document fields (like 'status')
-            dest_ref.set(source_doc.to_dict())
-            # Copy all sub-collections
-            for coll_ref in source_ref.collections():
-                copy_collection(coll_ref, dest_ref.collection(coll_ref.id))
-
-        # --- FINALIZATION ---
-        # Update the user's profile with the new phone and set status to 'approved'
-        # Also, remove the temporary migration data field.
+        action = migration_data['action']
+        to_phone = migration_data['to_phone']
+        to_phone_id = to_phone.lstrip('+')
+        
+        # --- EXECUTE ACTION ---
+        if action == 'migrate':
+            from_phone = migration_data.get('from_phone')
+            if not from_phone:
+                return jsonify({'error': 'Cannot migrate: Original phone number is missing.'}), 400
+            
+            from_phone_id = from_phone.lstrip('+')
+            source_ref = db.collection('users').document(from_phone_id)
+            dest_ref = db.collection('users').document(to_phone_id)
+            
+            logging.info(f"Admin approved MIGRATION for user {target_uid}. Copying from {from_phone_id} to {to_phone_id}")
+            source_doc = source_ref.get()
+            if source_doc.exists:
+                dest_ref.set(source_doc.to_dict())
+                for coll_ref in source_ref.collections():
+                    copy_collection(coll_ref, dest_ref.collection(coll_ref.id))
+            else:
+                 # If old document doesn't exist, just create a new empty one
+                 dest_ref.set({'status': 'approved', 'ownerUid': target_uid})
+
+        elif action == 'start_fresh':
+            logging.info(f"Admin approved FRESH START for user {target_uid} on new number {to_phone_id}")
+            # Simply create a new, approved document for the bot data.
+            db.collection('users').document(to_phone_id).set({
+                'status': 'approved',
+                'ownerUid': target_uid
+            })
+        
+        # --- FINALIZE ---
+        # Update the user's profile to complete the change.
         user_ref.update({
-            'phone': migration_data['to_phone'],
+            'phone': to_phone,
             'phoneStatus': 'approved',
             'migration_data': firestore.DELETE_FIELD
         })
         
-        # Optional: Delete the old data document after successful migration
-        # source_ref.delete() # Be cautious with this in production. Archival might be better.
-
-        logging.info(f"Successfully migrated data and updated phone for user {target_uid}")
-        return jsonify({'success': True, 'message': 'Phone number migration completed successfully.'}), 200
+        return jsonify({'success': True, 'message': f"Phone change action '{action}' completed successfully."}), 200
 
     except PermissionError as e:
         return jsonify({'error': str(e)}), 403
     except Exception as e:
-        logging.error(f"Admin migration approval failed for user {data.get('uid')}: {e}", exc_info=True)
-        return jsonify({'error': 'An internal error occurred during migration'}), 500
+        logging.error(f"Admin phone change approval failed for user {data.get('uid')}: {e}", exc_info=True)
+        return jsonify({'error': 'An internal error occurred during the process'}), 500
+        
 # -----------------------------------------------------------------------------
 # 6. ORGANIZATION MANAGEMENT (FULL CRUD)
 # -----------------------------------------------------------------------------