Update main.py

main.py

@@ -83,7 +83,49 @@ except Exception as e:
     api_key_error = True
     app.logger.error(f"Error initializing Gemini API Client: {e}")
 
-# --- Authentication Decorator ---
+
+# --- UID to Phone Mapping Functions (NEW for email auth) ---
+def get_uid_phone_mapping_path(uid):
+    """Get the Firebase path for storing user's phone number by UID"""
+    return f'users_by_uid/{uid}'
+
+def get_user_phone_by_uid(uid):
+    """Look up a user's phone number from their UID"""
+    if firebase_error or not firebase_db_ref:
+        return None, firebase_error or "Firebase not initialized"
+    try:
+        path = get_uid_phone_mapping_path(uid)
+        data = firebase_db_ref.child(path).get()
+        if data and isinstance(data, dict):
+            return data.get('phone_number'), None
+        return None, "Phone number not linked for this user"
+    except Exception as e:
+        return None, f"Error looking up phone: {e}"
+
+def set_user_phone_by_uid(uid, phone_number, display_name=None, email=None):
+    """Store a user's phone number mapping by UID"""
+    if firebase_error or not firebase_db_ref:
+        return False, firebase_error or "Firebase not initialized"
+    norm_phone = normalize_phone(phone_number)
+    if not norm_phone:
+        return False, "Invalid phone number format"
+    try:
+        path = get_uid_phone_mapping_path(uid)
+        user_data = {
+            'phone_number': norm_phone,
+            'linked_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        }
+        if display_name:
+            user_data['display_name'] = display_name
+        if email:
+            user_data['email'] = email
+        firebase_db_ref.child(path).set(user_data)
+        return True, None
+    except Exception as e:
+        return False, f"Error saving phone mapping: {e}"
+
+
+# --- Authentication Decorator (UPDATED for email auth) ---
 def token_required(f):
     @wraps(f)
     def decorated_function(*args, **kwargs):
@@ -97,20 +139,43 @@ def token_required(f):
         id_token = auth_header.split('Bearer ')[1]
         try:
             decoded_token = firebase_auth.verify_id_token(id_token)
-            flask_g.user = {
-                "uid": decoded_token.get('uid'),
-                "phone_number": decoded_token.get('phone_number') # This is crucial
-            }
-            if not flask_g.user["phone_number"]:
-                 app.logger.error(f"Token for UID {flask_g.user['uid']} does not contain a phone number.")
-                 return jsonify({"error": "Authenticated user has no phone number associated."}), 403
+            uid = decoded_token.get('uid')
+            email = decoded_token.get('email')
+            
+            # First check if phone is in the token (legacy phone auth)
+            phone_from_token = decoded_token.get('phone_number')
+            
+            if phone_from_token:
+                # Legacy phone authentication
+                phone_number = normalize_phone(phone_from_token)
+                if not phone_number:
+                    app.logger.error(f"Failed to normalize phone number from token for UID {uid}.")
+                    return jsonify({"error": "Invalid phone number format in token."}), 403
+                flask_g.user = {
+                    "uid": uid,
+                    "email": email,
+                    "phone_number": phone_number,
+                    "phone_linked": True
+                }
+            else:
+                # Email authentication - look up phone from database
+                phone_number, lookup_error = get_user_phone_by_uid(uid)
+                if not phone_number:
+                    # Phone not yet linked - allow access but mark as unlinked
+                    flask_g.user = {
+                        "uid": uid,
+                        "email": email,
+                        "phone_number": None,
+                        "phone_linked": False
+                    }
+                else:
+                    flask_g.user = {
+                        "uid": uid,
+                        "email": email,
+                        "phone_number": phone_number,
+                        "phone_linked": True
+                    }
             
-            # Normalize the phone number from the token immediately
-            flask_g.user["phone_number"] = normalize_phone(flask_g.user["phone_number"])
-            if not flask_g.user["phone_number"]:
-                app.logger.error(f"Failed to normalize phone number from token for UID {flask_g.user['uid']}.")
-                return jsonify({"error": "Invalid phone number format in token."}), 403
-
         except firebase_auth.ExpiredIdTokenError:
             return jsonify({"error": "Token has expired"}), 401
         except firebase_auth.InvalidIdTokenError:
@@ -123,6 +188,20 @@ def token_required(f):
     return decorated_function
 
 
+def phone_required(f):
+    """Decorator for endpoints that require a linked phone number"""
+    @wraps(f)
+    def decorated_function(*args, **kwargs):
+        if not flask_g.user.get("phone_linked") or not flask_g.user.get("phone_number"):
+            return jsonify({
+                "error": "Phone number not linked",
+                "message": "Please link your phone number to access this feature.",
+                "action_required": "link_phone"
+            }), 403
+        return f(*args, **kwargs)
+    return decorated_function
+
+
 # --- Helper Functions (Adapted for API context) ---
 def normalize_phone(phone):
     if not phone: return None
@@ -688,116 +767,200 @@ def generate_diff_summary(current_data, proposed_data):
             for mid in deleted_members:
                 summary.append(f"- {current_members[mid].get('name', 'Unnamed')} (ID: ...{mid[-6:]})")
 
-        modified_member_summary = []
+        # Check for edited members
+        edited_summary = []
         for mid in common_members:
-            member_diff = list(diff(current_members[mid], proposed_members[mid], ignore={'last_edited_at', 'last_edited_by', 'created_at', 'created_by', 'stories'})) # Ignore stories here, handle separately
-            if member_diff:
-                changes = []
-                for change_type, path, values in member_diff:
-                    field = path if isinstance(path, str) else path # Get field name
-                    if change_type == 'change': changes.append(f"{field}: '{values}' -> '{values}'")
-                    elif change_type == 'add': changes.append(f"{field}: added '{values}'") # Should be rare at top level
-                    elif change_type == 'remove': changes.append(f"{field}: removed '{values}'") # Should be rare at top level
-                if changes:
-                    modified_member_summary.append(f"- **{current_members[mid].get('name', 'Unnamed')}**: {'; '.join(changes)}")
-        if modified_member_summary:
-            summary.append("**Members Modified:**")
-            summary.extend(modified_member_summary)
-
-        # Process Relationships (Simplified comparison)
-        current_rels = {(r.get('from_id'), r.get('to_id'), r.get('type')) for r in current_data.get('relationships', []) if isinstance(r, dict)}
-        proposed_rels = {(r.get('from_id'), r.get('to_id'), r.get('type')) for r in proposed_data.get('relationships', []) if isinstance(r, dict)}
-        added_rels = proposed_rels - current_rels
-        deleted_rels = current_rels - proposed_rels
-
-        id_to_name_current = {m['id']: m.get('name', 'Unknown') for m in current_data.get('family_members', [])}
-        id_to_name_proposed = {m['id']: m.get('name', 'Unknown') for m in proposed_data.get('family_members', [])}
-        def format_rel(rel_tuple, names):
-            f_id, t_id, r_type = rel_tuple
-            f_name = names.get(f_id, f"...{f_id[-6:]}")
-            t_name = names.get(t_id, f"...{t_id[-6:]}")
-            return f"{f_name} {r_type} {t_name}"
+            curr_m, prop_m = current_members[mid], proposed_members[mid]
+            changes = []
+            for key in ['name', 'dob', 'dod', 'gender', 'totem', 'phone']:
+                if curr_m.get(key) != prop_m.get(key):
+                    changes.append(f"{key}: '{curr_m.get(key)}' -> '{prop_m.get(key)}'")
+            # Story changes
+            curr_stories = set(s.get('timestamp') for s in curr_m.get('stories', []) if isinstance(s, dict))
+            prop_stories = set(s.get('timestamp') for s in prop_m.get('stories', []) if isinstance(s, dict))
+            if curr_stories != prop_stories:
+                added_s = len(prop_stories - curr_stories)
+                removed_s = len(curr_stories - prop_stories)
+                if added_s > 0: changes.append(f"{added_s} storie(s) added")
+                if removed_s > 0: changes.append(f"{removed_s} storie(s) removed")
+            if changes:
+                edited_summary.append(f"- {curr_m.get('name', 'Unnamed')} (ID: ...{mid[-6:]}): {', '.join(changes)}")
+        if edited_summary:
+            summary.append("**Members Edited:**")
+            summary.extend(edited_summary)
+
+        # Process relationships
+        curr_rels = set(tuple(sorted([r['from_id'],r['to_id']])) + (r['type'],) for r in current_data.get('relationships', []) if isinstance(r, dict) and all(k in r for k in ['from_id', 'to_id', 'type']))
+        prop_rels = set(tuple(sorted([r['from_id'],r['to_id']])) + (r['type'],) for r in proposed_data.get('relationships', []) if isinstance(r, dict) and all(k in r for k in ['from_id', 'to_id', 'type']))
+        added_rels = prop_rels - curr_rels
+        deleted_rels = curr_rels - prop_rels
 
         if added_rels:
             summary.append("**Relationships Added:**")
-            for rel in added_rels: summary.append(f"- {format_rel(rel, id_to_name_proposed)}")
+            for rel_tuple in added_rels:
+                name1 = proposed_members.get(rel_tuple[0], {}).get('name', f'...{rel_tuple[0][-6:]}')
+                name2 = proposed_members.get(rel_tuple[1], {}).get('name', f'...{rel_tuple[1][-6:]}')
+                summary.append(f"- {name1} <--> {name2} ({rel_tuple[2]})")
         if deleted_rels:
             summary.append("**Relationships Deleted:**")
-            for rel in deleted_rels: summary.append(f"- {format_rel(rel, id_to_name_current)}")
-
-        # Process Stories (Simplified: check if stories list changed for any common member)
-        story_changes = []
-        for mid in common_members:
-            current_stories = current_members[mid].get('stories', [])
-            proposed_stories = proposed_members[mid].get('stories', [])
-            # Simple check: compare lengths or string representations (not perfect but gives an idea)
-            if json.dumps(current_stories, sort_keys=True) != json.dumps(proposed_stories, sort_keys=True):
-                 story_changes.append(f"- Stories for **{current_members[mid].get('name', 'Unnamed')}**")
-        if story_changes:
-             summary.append("**Stories Modified:**")
-             summary.extend(story_changes)
+            for rel_tuple in deleted_rels:
+                name1 = current_members.get(rel_tuple[0], {}).get('name', f'...{rel_tuple[0][-6:]}')
+                name2 = current_members.get(rel_tuple[1], {}).get('name', f'...{rel_tuple[1][-6:]}')
+                summary.append(f"- {name1} <--> {name2} ({rel_tuple[2]})")
+        
+        # Settings/Metadata changes
+        settings_changes = []
+        if current_data.get('settings', {}).get('theme') != proposed_data.get('settings', {}).get('theme'):
+            settings_changes.append(f"Theme: '{current_data.get('settings', {}).get('theme')}' -> '{proposed_data.get('settings', {}).get('theme')}'")
+        if current_data.get('metadata', {}).get('tree_name') != proposed_data.get('metadata', {}).get('tree_name'):
+            settings_changes.append(f"Tree Name: '{current_data.get('metadata', {}).get('tree_name')}' -> '{proposed_data.get('metadata', {}).get('tree_name')}'")
+        if settings_changes:
+            summary.append("**Settings/Metadata Changed:**")
+            summary.extend([f"- {c}" for c in settings_changes])
+
+        if not summary: return "No significant changes detected."
+        return "\n".join(summary)
     except Exception as e:
-        summary.append(f"Error generating diff: {e}")
-    return "\n".join(summary) if summary else "No significant changes detected."
+        app.logger.error(f"Error generating diff summary: {e}")
+        return f"Error generating diff summary: {e}"
 
-# --- API Endpoints ---
 
-# Helper to check for Firebase/API key errors (used by non-authed endpoints)
+# --- System Health Check Helper ---
 def check_system_health():
-    if firebase_error: return jsonify({"error": "Firebase system error", "detail": firebase_error}), 503
+    if firebase_error:
+        return jsonify({"error": "Firebase not available", "detail": firebase_error}), 503
+    if not firebase_db_ref:
+        return jsonify({"error": "Firebase database reference not set."}), 503
     return None
 
-@app.route("/health", methods=["GET"])
+
+# ======================= API Endpoints =======================
+
+@app.route('/health', methods=['GET'])
 def health_check():
-    system_error_response = check_system_health()
-    if system_error_response: return system_error_response
-    if api_key_error:
-        return jsonify({"status": "Firebase OK, Gemini API Key Error"}), 200
-    return jsonify({"status": "Firebase OK, Gemini OK"}), 200
+    status = {
+        "status": "ok",
+        "firebase": "ok" if not firebase_error and firebase_db_ref else f"error: {firebase_error or 'DB ref not set'}",
+        "gemini_api": "ok" if not api_key_error and genai_client else "error or not configured"
+    }
+    overall_ok = status["firebase"] == "ok"
+    return jsonify(status), 200 if overall_ok else 503
+
 
-# --- User/Auth related endpoints ---
+# --- NEW: User Management Endpoints for Email Auth ---
 @app.route('/user/me', methods=['GET'])
 @token_required
-def get_my_user_info():
-    """
-    Returns basic info about the authenticated user and their primary tree.
-    This also serves as a way to "register" or ensure the user's tree exists.
-    """
-    current_user_phone = flask_g.user["phone_number"]
+def get_current_user_api():
+    """Get the authenticated user's information"""
+    uid = flask_g.user["uid"]
+    email = flask_g.user.get("email")
+    phone_number = flask_g.user.get("phone_number")
+    phone_linked = flask_g.user.get("phone_linked", False)
     
-    # Load (or initialize) the user's primary tree data
-    tree_data, error = load_tree_data(current_user_phone)
-    if error:
-        # Specific check for initialization failure, which might be a server-side issue
-        if "Failed to initialize and save new tree" in error:
-            app.logger.error(f"Critical error initializing tree for {current_user_phone}: {error}")
-            return jsonify({"error": "Could not initialize user tree data", "detail": error}), 500
-        # Other load errors might be less critical or indicate no data (which is fine if new)
-        app.logger.warning(f"Could not load tree data for {current_user_phone} during /user/me: {error}")
-        # If tree_data is None and it's not an init error, it implies a load issue post-init.
-        # However, load_tree_data is designed to create if None, so this path needs careful thought.
-        # For now, assume if tree_data is None after load_tree_data, it's a significant issue.
-        if tree_data is None:
-             return jsonify({"error": "Could not load user tree data", "detail": error}), 404 # Or 500
+    response = {
+        "uid": uid,
+        "email": email,
+        "phone_number": phone_number,
+        "phone_linked": phone_linked
+    }
+    
+    # If phone is linked, get additional user info and primary tree
+    if phone_linked and phone_number:
+        tree_data, _ = load_tree_data(phone_number)
+        if tree_data:
+            response["display_name"] = tree_data.get("profile", {}).get("name")
+            response["primary_tree"] = {
+                "owner_phone": phone_number,
+                "tree_name": tree_data.get("metadata", {}).get("tree_name", "My Family Tree"),
+                "created_at": tree_data.get("metadata", {}).get("created_at")
+            }
+    else:
+        # Get display name from uid mapping if available
+        try:
+            path = get_uid_phone_mapping_path(uid)
+            user_data = firebase_db_ref.child(path).get()
+            if user_data and isinstance(user_data, dict):
+                response["display_name"] = user_data.get("display_name")
+        except Exception:
+            pass
+    
+    return jsonify(response), 200
 
+
+@app.route('/user/link_phone', methods=['POST'])
+@token_required
+def link_phone_api():
+    """Link a phone number to the authenticated user's account"""
+    uid = flask_g.user["uid"]
+    email = flask_g.user.get("email")
+    
+    # Check if phone is already linked
+    if flask_g.user.get("phone_linked"):
+        return jsonify({
+            "message": "Phone number already linked",
+            "phone_number": flask_g.user.get("phone_number")
+        }), 200
+    
+    req_data = request.json
+    if not req_data or "phone_number" not in req_data:
+        return jsonify({"error": "Missing phone_number in request body"}), 400
+    
+    phone_number = req_data["phone_number"]
+    display_name = req_data.get("display_name")
+    
+    # Validate phone format
+    norm_phone = normalize_phone(phone_number)
+    if not norm_phone:
+        return jsonify({"error": "Invalid phone number format. Use +263... format."}), 400
+    
+    # Check if this phone is already linked to another user
+    try:
+        all_users = firebase_db_ref.child('users_by_uid').get()
+        if all_users:
+            for existing_uid_key, user_data in all_users.items():
+                if isinstance(user_data, dict) and user_data.get('phone_number') == norm_phone:
+                    if existing_uid_key != uid:
+                        return jsonify({
+                            "error": "Phone number already linked",
+                            "message": "This phone number is already linked to another account."
+                        }), 409
+    except Exception as e:
+        app.logger.warning(f"Could not check for existing phone links: {e}")
+    
+    # Store the phone mapping
+    success, error = set_user_phone_by_uid(uid, norm_phone, display_name, email)
+    if not success:
+        return jsonify({"error": "Failed to link phone number", "detail": error}), 500
+    
+    # Initialize user's tree if it doesn't exist
+    tree_data, load_error = load_tree_data(norm_phone)
+    if load_error and "Failed to initialize" in load_error:
+        return jsonify({"error": "Phone linked but tree initialization failed", "detail": load_error}), 500
+    
+    # Update profile name if display_name provided
+    if display_name and tree_data:
+        tree_data["profile"]["name"] = display_name
+        me_node = find_person_by_id(tree_data, "Me")
+        if me_node:
+            me_node["name"] = display_name
+            me_node["last_edited_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            me_node["last_edited_by"] = norm_phone
+        save_tree_data(norm_phone, tree_data, tree_data)
+    
     return jsonify({
-        "uid": flask_g.user["uid"],
-        "phone_number": current_user_phone,
-        "primary_tree_metadata": tree_data.get("metadata") if tree_data else None
+        "message": "Phone number linked successfully",
+        "phone_number": norm_phone,
+        "tree_initialized": tree_data is not None
     }), 200
 
 
 # --- Tree Data Endpoints ---
-# Get a specific tree (can be own or another if part of it - auth handled by client if needed for "view")
-# For simplicity, this endpoint is public for GET, but write operations on trees are protected.
-@app.route('/tree/<string:owner_phone_of_tree_to_view>', methods=['GET'])
-def get_tree_data_api(owner_phone_of_tree_to_view):
-    system_error_response = check_system_health() # Basic check
+@app.route('/tree/<string:owner_phone>', methods=['GET'])
+def get_tree_data_api(owner_phone):
+    system_error_response = check_system_health()
     if system_error_response: return system_error_response
-
-    norm_phone = normalize_phone(owner_phone_of_tree_to_view)
-    if not norm_phone:
-        return jsonify({"error": "Invalid owner_phone format"}), 400
+    norm_phone = normalize_phone(owner_phone)
+    if not norm_phone: return jsonify({"error": "Invalid owner_phone format"}), 400
 
     tree_data, error = load_tree_data(norm_phone)
     if error:
@@ -811,7 +974,15 @@ def get_tree_data_api(owner_phone_of_tree_to_view):
 @token_required
 def get_linked_trees_api():
     """Gets trees the authenticated user is a member of (excluding their own primary tree)."""
-    current_user_phone = flask_g.user["phone_number"]
+    current_user_phone = flask_g.user.get("phone_number")
+    
+    # Handle unlinked phone users gracefully
+    if not current_user_phone:
+        return jsonify({
+            "message": "Link your phone number to see linked trees",
+            "trees": [],
+            "action_required": "link_phone"
+        }), 200
     
     linked_owner_phones_map, error = find_linked_trees(current_user_phone)
     if error:
@@ -859,6 +1030,7 @@ def get_tree_graph_api(owner_phone_of_tree):
 # --- Member Management Endpoints (Protected) ---
 @app.route('/tree/<string:owner_phone_of_tree>/members', methods=['POST'])
 @token_required
+@phone_required
 def add_member_api(owner_phone_of_tree):
     current_user_phone = flask_g.user["phone_number"] # Authenticated user
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -923,6 +1095,7 @@ def add_member_api(owner_phone_of_tree):
 
 @app.route('/tree/<string:owner_phone_of_tree>/members/<string:member_id>', methods=['PUT'])
 @token_required
+@phone_required
 def edit_member_api(owner_phone_of_tree, member_id):
     current_user_phone = flask_g.user["phone_number"]
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -1003,6 +1176,7 @@ def edit_member_api(owner_phone_of_tree, member_id):
 
 @app.route('/tree/<string:owner_phone_of_tree>/members/<string:member_id>', methods=['DELETE'])
 @token_required
+@phone_required
 def delete_member_api(owner_phone_of_tree, member_id):
     current_user_phone = flask_g.user["phone_number"]
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -1042,6 +1216,7 @@ def delete_member_api(owner_phone_of_tree, member_id):
 # --- Relationship Management Endpoints (Protected) ---
 @app.route('/tree/<string:owner_phone_of_tree>/relationships', methods=['POST'])
 @token_required
+@phone_required
 def add_relationship_api(owner_phone_of_tree):
     current_user_phone = flask_g.user["phone_number"]
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -1092,6 +1267,7 @@ def add_relationship_api(owner_phone_of_tree):
 
 @app.route('/tree/<string:owner_phone_of_tree>/relationships/delete', methods=['POST'])
 @token_required
+@phone_required
 def delete_relationship_api(owner_phone_of_tree):
     current_user_phone = flask_g.user["phone_number"]
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -1138,61 +1314,58 @@ def delete_relationship_api(owner_phone_of_tree):
 # --- Story Management Endpoints (Protected) ---
 @app.route('/tree/<string:owner_phone_of_tree>/members/<string:member_id>/stories', methods=['POST'])
 @token_required
+@phone_required
 def add_story_api(owner_phone_of_tree, member_id):
-    current_user_phone = flask_g.user["phone_number"] # Authenticated user adding the story
+    current_user_phone = flask_g.user["phone_number"]
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
     if not norm_owner_phone_of_tree: return jsonify({"error": "Invalid owner_phone_of_tree format"}), 400
 
-    # Anyone authenticated can add a story to any member of any tree they can see/access.
-    # The "added_by" field will track who added it.
-    # If stricter control is needed (e.g., only tree owner or proposer), add checks here.
-
     req_data = request.json
-    if not req_data or "text" not in req_data:
-        return jsonify({"error": "Missing story text in request body"}), 400
+    if not req_data or "text" not in req_data or not req_data["text"].strip():
+        return jsonify({"error": "Missing or empty 'text' for story"}), 400
     story_text = req_data["text"].strip()
-    if not story_text: return jsonify({"error": "Story text cannot be empty"}), 400
 
     tree_data, load_error = load_tree_data(norm_owner_phone_of_tree)
     if load_error: return jsonify({"error": "Could not load tree data", "detail": load_error}), 500
-    
+
     original_tree_data_for_save = deepcopy(tree_data) # For save if current_user is owner
     member = find_person_by_id(tree_data, member_id)
-    if not member: return jsonify({"error": "Member not found to add story to"}), 404
+    if not member: return jsonify({"error": "Member not found"}), 404
 
-    if "stories" not in member or not isinstance(member["stories"], list):
-        member["stories"] = []
-    
-    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    new_story_entry = deepcopy(DEFAULT_STORY_STRUCTURE)
-    new_story_entry.update({
-        "timestamp": timestamp, "text": story_text, "added_by": current_user_phone
-    })
-    member["stories"].append(new_story_entry)
-    member["last_edited_at"] = timestamp # Also update member's last edit
-    member["last_edited_by"] = current_user_phone # User who added story is last editor of member for this change
+    new_story = {
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "text": story_text,
+        "added_by": current_user_phone
+    }
+    member.setdefault("stories", []).append(new_story)
+    member["last_edited_at"] = new_story["timestamp"]
+    member["last_edited_by"] = current_user_phone
 
-    # If the authenticated user is the owner of the tree, save directly.
-    # Otherwise, this change should be part of a proposal.
     if current_user_phone == norm_owner_phone_of_tree:
+        # Owner adds story directly
         save_success, save_error = save_tree_data(norm_owner_phone_of_tree, tree_data, original_tree_data_for_save)
         if not save_success:
-            return jsonify({"error": "Failed to save new story", "detail": save_error}), 500
-        response_message = "Story added and saved successfully."
+            return jsonify({"error": "Failed to save story", "detail": save_error}), 500
+        response_message = "Story added successfully."
         if save_error: response_message += f" Warning: {save_error}"
-        return jsonify({"message": response_message, "story": new_story_entry}), 201
+        return jsonify({"message": response_message, "story": new_story}), 201
     else:
-        # This endpoint is now for direct save by owner. Collaborators use /proposals.
-        # So, if not owner, this path shouldn't be hit for adding stories directly.
-        # However, if we allow anyone to add stories to any tree (public edit), this logic would change.
-        # For now, sticking to owner-edit or proposal.
-        # This implies that if a collaborator wants to add a story, they modify their copy of tree_data
-        # and then submit the entire tree_data via the /proposals endpoint.
-        return jsonify({"error": "Forbidden: Use proposals to add stories to trees you do not own."}), 403
+        # Collaborator adds story: For simplicity, let's allow collaborators to add stories directly for now.
+        # A more strict approach would be to use proposals for all changes.
+        # Here, we allow stories to be added by non-owners as a contribution.
+        # The story will have `added_by` set to the collaborator's phone.
+        # Tree owner can delete it later if needed.
+        save_success, save_error = save_tree_data(norm_owner_phone_of_tree, tree_data, original_tree_data_for_save)
+        if not save_success:
+            return jsonify({"error": "Failed to save story (as collaborator)", "detail": save_error}), 500
+        response_message = "Story added successfully (as collaborator)."
+        if save_error: response_message += f" Warning: {save_error}"
+        return jsonify({"message": response_message, "story": new_story}), 201
 
 
-@app.route('/tree/<string:owner_phone_of_tree>/members/<string:member_id>/stories/<path:story_timestamp>', methods=['DELETE'])
+@app.route('/tree/<string:owner_phone_of_tree>/members/<string:member_id>/stories/<string:story_timestamp>', methods=['DELETE'])
 @token_required
+@phone_required
 def delete_story_api(owner_phone_of_tree, member_id, story_timestamp):
     current_user_phone = flask_g.user["phone_number"]
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -1238,6 +1411,7 @@ def delete_story_api(owner_phone_of_tree, member_id, story_timestamp):
 # --- Collaboration Endpoints (Protected) ---
 @app.route('/tree/<string:owner_phone_of_tree>/proposals', methods=['POST'])
 @token_required
+@phone_required
 def propose_changes_api(owner_phone_of_tree):
     proposer_phone_from_token = flask_g.user["phone_number"] # This is the authenticated user making the proposal
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -1261,6 +1435,7 @@ def propose_changes_api(owner_phone_of_tree):
 
 @app.route('/tree/my_proposals/pending', methods=['GET']) # Get pending proposals FOR the authenticated user's tree(s)
 @token_required
+@phone_required
 def get_my_pending_changes_api():
     owner_phone_from_token = flask_g.user["phone_number"] # Authenticated user is the owner
     
@@ -1272,6 +1447,7 @@ def get_my_pending_changes_api():
 
 @app.route('/tree/my_proposals/<string:proposer_phone_to_manage>/accept', methods=['POST'])
 @token_required
+@phone_required
 def accept_changes_api(proposer_phone_to_manage):
     owner_phone_from_token = flask_g.user["phone_number"] # Authenticated user is the owner
     norm_proposer_phone = normalize_phone(proposer_phone_to_manage)
@@ -1285,6 +1461,7 @@ def accept_changes_api(proposer_phone_to_manage):
 
 @app.route('/tree/my_proposals/<string:proposer_phone_to_manage>/reject', methods=['POST'])
 @token_required
+@phone_required
 def reject_changes_api(proposer_phone_to_manage):
     owner_phone_from_token = flask_g.user["phone_number"] # Authenticated user is the owner
     norm_proposer_phone = normalize_phone(proposer_phone_to_manage)
@@ -1297,6 +1474,7 @@ def reject_changes_api(proposer_phone_to_manage):
 
 @app.route('/tree/my_proposals/<string:proposer_phone_to_review>/diff', methods=['GET'])
 @token_required
+@phone_required
 def get_proposal_diff_api(proposer_phone_to_review):
     owner_phone_from_token = flask_g.user["phone_number"] # Authenticated user is the owner
     norm_proposer_phone = normalize_phone(proposer_phone_to_review)
@@ -1340,6 +1518,7 @@ def ai_build_tree_api():
 
 @app.route('/tree/<string:owner_phone_of_tree>/ai_merge_suggestions', methods=['POST'])
 @token_required
+@phone_required
 def ai_merge_tree_api(owner_phone_of_tree):
     current_user_phone = flask_g.user["phone_number"]
     norm_owner_phone_of_tree = normalize_phone(owner_phone_of_tree)
@@ -1382,7 +1561,7 @@ def ai_merge_tree_api(owner_phone_of_tree):
         ai_name = person_ai.get("name")
         if not ai_name: continue
         existing_people_with_name = find_person_by_name(tree_data, ai_name)
-        existing_person = existing_people_with_name if existing_people_with_name else None
+        existing_person = existing_people_with_name[0] if existing_people_with_name else None
         person_updated = False
         if existing_person:
             id_map[ai_name] = existing_person['id']
@@ -1474,6 +1653,7 @@ def generate_quiz_api(owner_phone_of_tree):
 # --- Settings & Profile Endpoints (Protected, for authenticated user's own tree) ---
 @app.route('/user/me/settings', methods=['PUT'])
 @token_required
+@phone_required
 def update_my_settings_api():
     current_user_phone = flask_g.user["phone_number"] # Settings are for the authenticated user's own tree
     
@@ -1508,6 +1688,7 @@ def update_my_settings_api():
 
 @app.route('/user/me/profile', methods=['PUT'])
 @token_required
+@phone_required
 def update_my_profile_api():
     current_user_phone = flask_g.user["phone_number"] # Profile is for the authenticated user's own tree
     
@@ -1592,4 +1773,4 @@ def get_timeline_api(owner_phone_of_tree):
 
 # --------- Run the App ---------
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860, debug=True)
+    app.run(host="0.0.0.0", port=7860, debug=True)