Add logic to handle budget _id in get_recommendation_for_category - extract categories and generate recommendations

app/smart_recommendation.py

@@ -327,7 +327,7 @@ class SmartBudgetRecommender:
         
         Args:
             user_id: User identifier
-            category_id: Category ID to get recommendations for
+            category_id: Category ID to get recommendations for (can also be a budget _id)
             month: Target month (1-12)
             year: Target year
             budget_amount: Optional current budget amount to use for recommendations
@@ -335,6 +335,73 @@ class SmartBudgetRecommender:
         Returns:
             List of budget recommendations for the specific category
         """
+        # FIRST: Check if category_id is actually a budget _id
+        # If so, extract the budget's data and categories
+        try:
+            try:
+                budget_id_objid = ObjectId(category_id)
+                budget_by_id = self.db.budgets.find_one({"_id": budget_id_objid})
+                if budget_by_id:
+                    budget_created_by = budget_by_id.get("createdBy")
+                    # Check if this budget belongs to the user
+                    budget_user_match = False
+                    if budget_created_by:
+                        if isinstance(budget_created_by, ObjectId):
+                            budget_user_match = (str(budget_created_by) == str(user_id) or budget_created_by == ObjectId(user_id))
+                        else:
+                            budget_user_match = (str(budget_created_by) == str(user_id))
+                    
+                    if budget_user_match:
+                        # Extract budget amount and categories
+                        budget_max_amount = float(budget_by_id.get("maxAmount", 0) or 0)
+                        budget_spend_amount = float(budget_by_id.get("spendAmount", 0) or 0)
+                        budget_amount_from_budget = budget_spend_amount if budget_spend_amount > 0 else budget_max_amount
+                        
+                        # Extract categories from headCategories
+                        head_categories = budget_by_id.get("headCategories", [])
+                        category_ids_in_budget = []
+                        for hc in head_categories:
+                            if isinstance(hc, dict):
+                                hc_id = hc.get("headCategory")
+                                if hc_id:
+                                    category_ids_in_budget.append(str(hc_id))
+                                nested_cats = hc.get("categories", [])
+                                for nc in nested_cats:
+                                    if isinstance(nc, dict):
+                                        nc_id = nc.get("category")
+                                        if nc_id:
+                                            category_ids_in_budget.append(str(nc_id))
+                        
+                        # If budget has categories and amount, generate recommendation
+                        if category_ids_in_budget and budget_amount_from_budget > 0:
+                            # Use the first category ID found in the budget
+                            actual_category_id = category_ids_in_budget[0]
+                            category_name = self._get_category_name(actual_category_id)
+                            
+                            # Use budget amount if no budget_amount provided
+                            if budget_amount is None:
+                                budget_amount = budget_amount_from_budget
+                            
+                            # Generate recommendation for this category using the budget amount
+                            return self._generate_single_recommendation(
+                                user_id, actual_category_id, category_name, 
+                                month, year, budget_amount, using_budget_amount_only=False
+                            )
+                        elif budget_amount_from_budget > 0:
+                            # Budget exists but no categories - use budget name and amount
+                            budget_name = budget_by_id.get("name", "Budget")
+                            if budget_amount is None:
+                                budget_amount = budget_amount_from_budget
+                            # Generate recommendation using budget name
+                            return self._generate_single_recommendation(
+                                user_id, category_id, budget_name,
+                                month, year, budget_amount, using_budget_amount_only=False
+                            )
+            except (ValueError, TypeError):
+                pass  # category_id is not a valid ObjectId, continue with normal check
+        except Exception:
+            pass  # Silently continue if budget check fails
+        
         # Get all recommendations for the user
         all_recommendations = self.get_recommendations(user_id, month, year)
         print(f"🔍 get_recommendation_for_category: Found {len(all_recommendations)} total recommendations for user {user_id}")