app/smart_recommendation.py

import json
import math
import os
from collections import defaultdict
from datetime import datetime, timedelta
from typing import Dict, List, Optional

import requests
from dotenv import load_dotenv
from bson import ObjectId

from app.models import BudgetRecommendation, CategoryExpense

load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

class SmartBudgetRecommender:
    """
    Smart Budget Recommendation Engine
    
    Analyzes past spending behavior and recommends personalized budgets
    for each category based on historical data.
    """
    
    def __init__(self, db):
        self.db = db
    
    def get_recommendations(self, user_id: str, month: int, year: int) -> List[BudgetRecommendation]:
        """
        Get budget recommendations for all categories based on past behavior.
        
        Args:
            user_id: User identifier
            month: Target month (1-12)
            year: Target year
            
        Returns:
            List of budget recommendations for each category
        """
        # 1) Try to build stats from existing budgets for this user (createdBy)
        category_data = self._get_category_stats_from_budgets(user_id, month, year)

        # 2) Only return recommendations for actual budgets - do NOT use expenses history
        # This ensures we only show recommendations for budgets the user actually created
        if not category_data:
            print(f"No budgets found for user_id: {user_id}, returning empty recommendations")
            return []
        
        recommendations: List[BudgetRecommendation] = []

        for category_key, data in category_data.items():
            # Extract category_name and category_id from data
            # category_key format: "user_id|category_name|category_id"
            key_parts = category_key.split("|")
            if len(key_parts) >= 3:
                # Skip user_id (first part), get category_name (second part)
                category_name = data.get("category_name", key_parts[1])
            elif len(key_parts) >= 2:
                category_name = data.get("category_name", key_parts[1])
            else:
                category_name = data.get("category_name", category_key)
            category_id = data.get("category_id")
            avg_expense = data["average_monthly"]
            confidence = self._calculate_confidence(data)

            # Always try OpenAI first (primary source of recommendation)
            ai_result = self._get_ai_recommendation(category_name, data, avg_expense)
            if ai_result and ai_result.get("recommended_budget"):
                recommended_budget = ai_result.get("recommended_budget")
                reason = ai_result.get("reason", f"AI recommendation for {category_name}")
                action = ai_result.get("action")
                
                # Validate OpenAI recommendation (same logic as in get_recommendation_for_category)
                if recommended_budget == avg_expense and action == "keep":
                    std_dev = data.get("std_dev", 0.0)
                    monthly_values = data.get("monthly_values", [])
                    has_trend = len(monthly_values) > 1 and (monthly_values[-1] != monthly_values[0])
                    
                    if has_trend or std_dev > avg_expense * 0.05:
                        # Override with intelligent recommendation
                        if has_trend and monthly_values[-1] > monthly_values[0]:
                            recommended_budget = avg_expense * 1.15
                            action = "increase"
                        elif std_dev > avg_expense * 0.05:
                            recommended_budget = avg_expense * 1.20
                            action = "increase"
                        else:
                            recommended_budget = avg_expense * 1.05
                            action = "increase"
                
                print(f"✅ OpenAI recommendation for {category_name}: {recommended_budget} (action: {action})")
            else:
                # Fallback to rule-based recommendation if OpenAI fails
                recommended_budget = self._calculate_recommended_budget(avg_expense, data)
                reason = self._generate_reason(category_name, avg_expense, recommended_budget)
                action = None
                if not ai_result:
                    print(f"❌ OpenAI unavailable (no API key or error), using rule-based for {category_name}: {recommended_budget}")
                else:
                    print(f"⚠️ OpenAI returned invalid data, using rule-based for {category_name}: {recommended_budget}")

            recommendations.append(BudgetRecommendation(
                category=category_name,
                category_id=category_id,
                average_expense=round(avg_expense, 2),
                recommended_budget=round(recommended_budget or 0, 2),
                reason=reason,
                confidence=confidence,
                action=action
            ))
        
        # Sort by average expense (highest first)
        recommendations.sort(key=lambda x: x.average_expense, reverse=True)
        
        return recommendations
    
    def check_user_has_category_data(self, user_id: str, category_id: str) -> bool:
        """
        Check if user has previous budget or expense data for a specific category.
        
        Args:
            user_id: User identifier
            category_id: Category ID to check
            
        Returns:
            True if user has previous data for this category, False otherwise
        """
        # FIRST: Check if category_id is actually a budget _id
        # If so, find the budget and check if it belongs to the user and has categories
        try:
            try:
                budget_id_objid = ObjectId(category_id)
                # Try to find a budget with this _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 (handle both ObjectId and string comparisons)
                    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 all category IDs from this budget's headCategories
                        head_categories = budget_by_id.get("headCategories", [])
                        category_ids_in_budget = []
                        for hc in head_categories:
                            if isinstance(hc, dict):
                                # Check headCategory itself
                                hc_id = hc.get("headCategory")
                                if hc_id:
                                    category_ids_in_budget.append(str(hc_id))
                                # Check nested categories
                                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, consider it as having previous data
                        if category_ids_in_budget:
                            return True
            except (ValueError, TypeError):
                pass  # category_id is not a valid ObjectId, continue with normal check
        except Exception as e:
            pass  # Silently continue if budget check fails
        
        # Build comprehensive user query
        user_conditions = []
        try:
            user_objid = ObjectId(user_id)
            user_conditions = [
                {"createdBy": user_objid},
                {"createdBy": user_id},
                {"user_id": user_objid},
                {"user_id": user_id}
            ]
        except (ValueError, TypeError):
            user_conditions = [
                {"createdBy": user_id},
                {"user_id": user_id}
            ]
        
        # Build comprehensive category query - check all possible fields
        category_conditions = []
        try:
            category_objid = ObjectId(category_id)
            # Try as ObjectId
            category_conditions = [
                {"category": category_objid},
                {"categoryId": category_objid},
                {"headCategory": category_objid},
                {"headCategories.headCategory": category_objid},
                {"headCategories.categories.category": category_objid},
                # Also try as string
                {"category": category_id},
                {"categoryId": category_id},
                {"headCategory": category_id},
                {"headCategories.headCategory": category_id},
                {"headCategories.categories.category": category_id},
            ]
        except (ValueError, TypeError):
            # category_id is not a valid ObjectId, try as string only
            category_conditions = [
                {"category": category_id},
                {"categoryId": category_id},
                {"headCategory": category_id},
                {"headCategories.headCategory": category_id},
                {"headCategories.categories.category": category_id},
            ]
        
        # SECOND: Check nested structure (headCategories[].categories[].category) - most common case
        try:
            try:
                category_objid = ObjectId(category_id)
            except (ValueError, TypeError):
                category_objid = category_id
            
            # Try multiple nested query patterns
            nested_queries = [
                {
                    "$and": [
                        {"$or": user_conditions},
                        {
                            "$or": [
                                {"headCategories": {"$elemMatch": {"categories": {"$elemMatch": {"category": category_objid}}}}},
                                {"headCategories": {"$elemMatch": {"categories": {"$elemMatch": {"category": category_id}}}}},
                                {"headCategories": {"$elemMatch": {"headCategory": category_objid}}},
                                {"headCategories": {"$elemMatch": {"headCategory": category_id}}},
                            ]
                        }
                    ]
                },
                {
                    "$and": [
                        {"$or": user_conditions},
                        {
                            "$or": [
                                {"headCategories.categories.category": category_objid},
                                {"headCategories.categories.category": category_id},
                                {"headCategories.headCategory": category_objid},
                                {"headCategories.headCategory": category_id},
                            ]
                        }
                    ]
                }
            ]
            
            for nested_query in nested_queries:
                try:
                    if self.db.budgets.count_documents(nested_query) > 0:
                        return True
                except Exception:
                    continue
        except Exception:
            pass
        
        # THIRD: Check direct category fields
        for user_cond in user_conditions:
            for cat_cond in category_conditions:
                try:
                    if self.db.budgets.count_documents({**user_cond, **cat_cond}) > 0:
                        return True
                except Exception:
                    continue
        
        # THIRD: Try comprehensive query using $and and $or
        try:
            comprehensive_query = {
                "$and": [
                    {"$or": user_conditions},
                    {"$or": category_conditions}
                ]
            }
            if self.db.budgets.count_documents(comprehensive_query) > 0:
                return True
        except Exception:
            pass
        
        # FOURTH: Check expenses collection as fallback
        try:
            try:
                category_objid = ObjectId(category_id)
                expense_user_conditions = [
                    {"user_id": ObjectId(user_id)},
                    {"user_id": user_id},
                    {"createdBy": ObjectId(user_id)},
                    {"createdBy": user_id}
                ]
                expense_category_conditions = [
                    {"category": category_objid},
                    {"category": category_id},
                    {"categoryId": category_objid},
                    {"categoryId": category_id},
                    {"headCategory": category_objid},
                    {"headCategory": category_id}
                ]
            except (ValueError, TypeError):
                expense_user_conditions = [
                    {"user_id": user_id},
                    {"createdBy": user_id}
                ]
                expense_category_conditions = [
                    {"category": category_id},
                    {"categoryId": category_id},
                    {"headCategory": category_id}
                ]
            
            for user_cond in expense_user_conditions:
                for cat_cond in expense_category_conditions:
                    try:
                        if self.db.expenses.count_documents({**user_cond, **cat_cond}) > 0:
                            return True
                    except Exception:
                        continue
        except Exception:
            pass
        
        return False
    
    def get_recommendation_for_category(self, user_id: str, category_id: str, month: int, year: int, budget_amount: Optional[float] = None) -> List[BudgetRecommendation]:
        """
        Get budget recommendations for a specific category for a user.
        
        Args:
            user_id: User identifier
            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
            
        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 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, generate recommendation
                        if category_ids_in_budget:
                            # 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)
                            
                            # PRIORITY: Use provided budget_amount if available, otherwise use budget's maxAmount
                            # IMPORTANT: If user provided budget_amount, use it as the base for recommendation
                            original_budget_amount = budget_amount  # Store original for logging
                            if budget_amount is None or budget_amount <= 0:
                                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 = budget_spend_amount if budget_spend_amount > 0 else budget_max_amount
                                print(f"📊 Using budget's maxAmount/spendAmount: {budget_amount:,.2f} (user did not provide budget_amount)")
                            else:
                                print(f"✅ Using user-provided budget_amount: {budget_amount:,.2f} (ignoring budget's maxAmount)")
                            
                            # If we have a valid budget amount, generate recommendation
                            if budget_amount and budget_amount > 0:
                                # CRITICAL: Use the provided budget_amount as average_expense
                                # This is what the user wants to set, so recommendations should be based on this
                                avg_expense = budget_amount
                                print(f"💰 Setting average_expense = {avg_expense:,.2f} (from user's budget_amount)")
                                monthly_values = [avg_expense]
                                std_dev = avg_expense * 0.05
                                months_analyzed = 1
                                
                                data = {
                                    "average_monthly": avg_expense,
                                    "total": avg_expense,
                                    "count": 1,
                                    "months_analyzed": months_analyzed,
                                    "std_dev": std_dev,
                                    "monthly_values": monthly_values,
                                }
                                
                                confidence = self._calculate_confidence(data)
                                
                                # Get AI recommendation - pass the user's budget_amount so OpenAI knows what they set
                                ai_result = self._get_ai_recommendation(category_name, data, avg_expense)
                                if ai_result and ai_result.get("recommended_budget"):
                                    recommended_budget = ai_result.get("recommended_budget")
                                    reason = ai_result.get("reason", f"AI recommendation for {category_name}")
                                    action = ai_result.get("action")
                                else:
                                    recommended_budget = avg_expense * 1.10
                                    reason = f"Based on your budget of {budget_amount:,.0f}, I recommend {recommended_budget:,.0f} to account for variability and inflation."
                                    action = "increase"
                                
                                return [BudgetRecommendation(
                                    category=category_name,
                                    category_id=actual_category_id,
                                    average_expense=round(avg_expense, 2),
                                    recommended_budget=round(recommended_budget or 0, 2),
                                    reason=reason,
                                    confidence=confidence,
                                    action=action
                                )]
                        
                        # Budget exists but no categories - use budget name and amount
                        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
                        
                        if 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
                            avg_expense = budget_amount
                            monthly_values = [avg_expense]
                            std_dev = avg_expense * 0.05
                            months_analyzed = 1
                            
                            data = {
                                "average_monthly": avg_expense,
                                "total": avg_expense,
                                "count": 1,
                                "months_analyzed": months_analyzed,
                                "std_dev": std_dev,
                                "monthly_values": monthly_values,
                            }
                            
                            confidence = self._calculate_confidence(data)
                            
                            # Get AI recommendation
                            ai_result = self._get_ai_recommendation(budget_name, data, avg_expense)
                            if ai_result and ai_result.get("recommended_budget"):
                                recommended_budget = ai_result.get("recommended_budget")
                                reason = ai_result.get("reason", f"AI recommendation for {budget_name}")
                                action = ai_result.get("action")
                            else:
                                recommended_budget = avg_expense * 1.10
                                reason = f"Based on your budget of {budget_amount:,.0f}, I recommend {recommended_budget:,.0f} to account for variability."
                                action = "increase"
                            
                            return [BudgetRecommendation(
                                category=budget_name,
                                category_id=category_id,
                                average_expense=round(avg_expense, 2),
                                recommended_budget=round(recommended_budget or 0, 2),
                                reason=reason,
                                confidence=confidence,
                                action=action
                            )]
            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}")
        
        # Filter to only include recommendations for the specified category_id
        filtered_recommendations = [
            rec for rec in all_recommendations 
            if rec.category_id == category_id or str(rec.category_id) == str(category_id)
        ]
        print(f"🔍 get_recommendation_for_category: Filtered to {len(filtered_recommendations)} recommendations for category_id {category_id}")
        
        # If we found a recommendation, use it (or regenerate with budget_amount if provided)
        if filtered_recommendations:
            # If budget_amount is provided, regenerate recommendation with it
            if budget_amount and budget_amount > 0:
                # Get the first recommendation (should be the one for this category)
                original_rec = filtered_recommendations[0]
                
                # Get category name
                category_name = original_rec.category
                
                # Use the provided budget_amount as average_expense for comparison
                avg_expense = budget_amount
                
                # Create data structure for recommendation calculation
                data = {
                    "average_monthly": avg_expense,
                    "total": avg_expense,
                    "count": 1,
                    "months_analyzed": 1,
                    "std_dev": 0.0,
                    "monthly_values": [avg_expense],
                }
                
                confidence = self._calculate_confidence(data)
                
                # Always try OpenAI first (primary source of recommendation)
                ai_result = self._get_ai_recommendation(category_name, data, avg_expense)
                if ai_result and ai_result.get("recommended_budget"):
                    recommended_budget = ai_result.get("recommended_budget")
                    reason = ai_result.get("reason", f"AI recommendation for {category_name} based on your budget of {budget_amount:,.2f}")
                    action = ai_result.get("action")
                    
                    # Validate OpenAI recommendation
                    if recommended_budget == avg_expense and action == "keep":
                        # For budget_amount only, always add buffer
                        recommended_budget = avg_expense * 1.10
                        action = "increase"
                        reason = f"Based on your budget amount, I recommend increasing by 10% to {recommended_budget:,.0f} to account for variability and inflation."
                    
                    print(f"✅ OpenAI recommendation for {category_name} (budget: {budget_amount}): {recommended_budget} (action: {action})")
                else:
                    # Fallback to rule-based recommendation if OpenAI fails
                    recommended_budget = self._calculate_recommended_budget(avg_expense, data)
                    reason = self._generate_reason(category_name, avg_expense, recommended_budget)
                    action = None
                    if not ai_result:
                        print(f"❌ OpenAI unavailable (no API key or error), using rule-based for {category_name} (budget: {budget_amount}): {recommended_budget}")
                    else:
                        print(f"⚠️ OpenAI returned invalid data, using rule-based for {category_name} (budget: {budget_amount}): {recommended_budget}")
                
                # Create new recommendation based on provided budget_amount
                filtered_recommendations = [BudgetRecommendation(
                    category=category_name,
                    category_id=category_id,
                    average_expense=round(avg_expense, 2),
                    recommended_budget=round(recommended_budget or 0, 2),
                    reason=reason,
                    confidence=confidence,
                    action=action
                )]
            return filtered_recommendations
        
        # If no recommendations found, try to generate one
        print(f"🔍 get_recommendation_for_category: No recommendations found, trying to generate one")
        print(f"🔍 get_recommendation_for_category: budget_amount = {budget_amount}")
        
        # If budget_amount is provided, use it
        using_budget_amount_only = False
        if budget_amount and budget_amount > 0:
            # Check for data corruption in budget_amount
            if budget_amount > 1e15:  # Unreasonably large number
                print(f"🚨 DATA CORRUPTION: budget_amount is {budget_amount:,.2e} - too large, using safe fallback")
                # Use a reasonable default based on category or cap at 1 billion
                budget_amount = 1e9  # 1 billion as safe maximum
                print(f"   Capped budget_amount to {budget_amount:,.0f}")
            
            print(f"🔍 get_recommendation_for_category: Using provided budget_amount: {budget_amount:,.0f}")
            # First, get category name
            category_name = self._get_category_name(category_id)
            avg_expense = budget_amount
            # Mark that we're using budget_amount only (no historical data)
            using_budget_amount_only = True
        else:
            # Try to get budget data for this category
            print(f"🔍 get_recommendation_for_category: No budget_amount provided, trying to get budget data")
            # First, get category name
            category_name = self._get_category_name(category_id)
            try:
                # Try to find budgets with this category_id
                try:
                    category_objid = ObjectId(category_id)
                except (ValueError, TypeError):
                    category_objid = category_id
                
                # Build user query
                user_query = {
                    "$or": [
                        {"createdBy": ObjectId(user_id) if len(user_id) == 24 and all(c in '0123456789abcdefABCDEF' for c in user_id) else user_id},
                        {"createdBy": user_id},
                        {"user_id": ObjectId(user_id) if len(user_id) == 24 and all(c in '0123456789abcdefABCDEF' for c in user_id) else user_id},
                        {"user_id": user_id}
                    ]
                }
                
                # Build category query
                category_query = {
                    "$or": [
                        {"category": category_objid},
                        {"categoryId": category_objid},
                        {"headCategory": category_objid},
                        {"headCategories.headCategory": category_objid},
                        {"headCategories.categories.category": category_objid}
                    ]
                }
                
                # Combine queries
                budget_query = {"$and": [user_query, category_query]}
                
                budgets = list(self.db.budgets.find(budget_query).limit(10))
                print(f"🔍 get_recommendation_for_category: Found {len(budgets)} budgets for user {user_id} and category {category_id}")
                
                if budgets:
                    # Calculate average from budgets
                    total_amount = 0
                    count = 0
                    for budget in budgets:
                        try:
                            max_amount = float(budget.get("maxAmount", 0) or budget.get("max_amount", 0) or budget.get("amount", 0) or 0)
                            spend_amount = float(budget.get("spendAmount", 0) or budget.get("spend_amount", 0) or budget.get("spent", 0) or 0)
                            budget_amount_val = float(budget.get("budget", 0) or budget.get("budgetAmount", 0) or 0)
                            
                            base_amount = spend_amount if spend_amount > 0 else (max_amount if max_amount > 0 else budget_amount_val)
                            if base_amount > 0:
                                total_amount += base_amount
                                count += 1
                        except (ValueError, TypeError):
                            continue
                    
                    if count > 0:
                        avg_expense = total_amount / count
                    else:
                        # No valid amounts found, can't generate recommendation
                        return []
                else:
                    # No budgets found for this category
                    return []
            except Exception as e:
                print(f"Error getting budget data for category: {e}")
                return []
        
        # Generate recommendation
        print(f"🔍 get_recommendation_for_category: Generating recommendation for category_name={category_name}, avg_expense={avg_expense}")
        
        # If we only have budget_amount (no historical data), use it directly
        # DO NOT create fake/simulated data - be honest with OpenAI that this is a new budget
        if using_budget_amount_only:
            # Use the budget_amount as a single data point
            # Don't create fake trends - OpenAI should recommend based on:
            # 1. The budget amount provided
            # 2. Category-specific knowledge
            # 3. General inflation and best practices
            monthly_values = [avg_expense]  # Single data point - no fake history
            std_dev = avg_expense * 0.05  # Assume 5% typical variation for new budgets
            months_analyzed = 1  # Only one month of data (the provided budget_amount)
        else:
            # We have historical budget data - calculate statistics from budgets
            # Try to get monthly values from budgets by analyzing dates
            monthly_values = []
            try:
                # Get budgets again to calculate monthly statistics
                try:
                    category_objid = ObjectId(category_id) if len(category_id) == 24 and all(c in '0123456789abcdefABCDEF' for c in category_id) else category_id
                except (ValueError, TypeError):
                    category_objid = category_id
                
                user_query = {
                    "$or": [
                        {"createdBy": ObjectId(user_id) if len(user_id) == 24 and all(c in '0123456789abcdefABCDEF' for c in user_id) else user_id},
                        {"createdBy": user_id},
                        {"user_id": ObjectId(user_id) if len(user_id) == 24 and all(c in '0123456789abcdefABCDEF' for c in user_id) else user_id},
                        {"user_id": user_id}
                    ]
                }
                
                category_query = {
                    "$or": [
                        {"category": category_objid},
                        {"categoryId": category_objid},
                        {"headCategory": category_objid},
                        {"headCategories.headCategory": category_objid},
                        {"headCategories.categories.category": category_objid}
                    ]
                }
                
                budget_query = {"$and": [user_query, category_query]}
                budgets = list(self.db.budgets.find(budget_query).sort("createdAt", -1).limit(12))
                
                # Group by month and calculate monthly totals
                monthly_totals = defaultdict(float)
                monthly_counts = defaultdict(int)
                
                for budget in budgets:
                    try:
                        # Get date from budget
                        budget_date = budget.get("createdAt") or budget.get("date") or budget.get("startDate")
                        if budget_date:
                            if isinstance(budget_date, str):
                                budget_date = datetime.fromisoformat(budget_date.replace('Z', '+00:00'))
                            elif not isinstance(budget_date, datetime):
                                continue
                            
                            month_key = f"{budget_date.year}-{budget_date.month:02d}"
                            
                            max_amount = float(budget.get("maxAmount", 0) or budget.get("max_amount", 0) or budget.get("amount", 0) or 0)
                            spend_amount = float(budget.get("spendAmount", 0) or budget.get("spend_amount", 0) or budget.get("spent", 0) or 0)
                            budget_amount_val = float(budget.get("budget", 0) or budget.get("budgetAmount", 0) or 0)
                            
                            base_amount = spend_amount if spend_amount > 0 else (max_amount if max_amount > 0 else budget_amount_val)
                            if base_amount > 0:
                                monthly_totals[month_key] += base_amount
                                monthly_counts[month_key] += 1
                    except (ValueError, TypeError, AttributeError):
                        continue
                
                # Convert to monthly values list
                if monthly_totals:
                    # Sort by month key and calculate averages
                    sorted_months = sorted(monthly_totals.keys())
                    monthly_values = [monthly_totals[month] / monthly_counts[month] for month in sorted_months]
                    months_analyzed = len(monthly_values)
                    
                    # Calculate std_dev
                    if len(monthly_values) > 1:
                        mean = sum(monthly_values) / len(monthly_values)
                        variance = sum((x - mean) ** 2 for x in monthly_values) / len(monthly_values)
                        std_dev = variance ** 0.5
                    else:
                        std_dev = avg_expense * 0.05  # Default 5% variation
                else:
                    # Fallback: use avg_expense as single data point
                    monthly_values = [avg_expense]
                    std_dev = avg_expense * 0.05
                    months_analyzed = 1
            except Exception as e:
                print(f"Error calculating monthly statistics: {e}")
                # Fallback: use avg_expense as single data point
                monthly_values = [avg_expense]
                std_dev = avg_expense * 0.05
                months_analyzed = 1
        
        data = {
            "average_monthly": avg_expense,
            "total": sum(monthly_values),
            "count": months_analyzed,
            "months_analyzed": months_analyzed,
            "std_dev": std_dev,
            "monthly_values": monthly_values,
        }
        
        confidence = self._calculate_confidence(data)
        print(f"🔍 get_recommendation_for_category: Confidence calculated: {confidence}")
        
        # Always try OpenAI first
        ai_result = self._get_ai_recommendation(category_name, data, avg_expense)
        if ai_result and ai_result.get("recommended_budget"):
            recommended_budget = ai_result.get("recommended_budget")
            reason = ai_result.get("reason", f"AI recommendation for {category_name}")
            action = ai_result.get("action")
            
            # VALIDATION: Check if OpenAI returned a lazy "keep" recommendation
            # If recommended_budget equals avg_expense and action is "keep", validate if it's justified
            monthly_values = data.get("monthly_values", [])
            std_dev = data.get("std_dev", 0.0)
            
            if recommended_budget == avg_expense and action == "keep":
                # Check if this is justified
                has_trend = False
                if len(monthly_values) > 1:
                    # Check for upward trend
                    if monthly_values[-1] > monthly_values[0]:
                        has_trend = True
                    # Check for downward trend
                    elif monthly_values[-1] < monthly_values[0]:
                        has_trend = True
                
                # If there's a trend or high variation, force a better recommendation
                if has_trend or std_dev > avg_expense * 0.05:
                    print(f"⚠️ OpenAI returned 'keep' but data shows trend/variation - overriding with intelligent recommendation")
                    # Force increase with buffer
                    if has_trend and monthly_values[-1] > monthly_values[0]:
                        # Upward trend - increase by 15%
                        recommended_budget = avg_expense * 1.15
                        action = "increase"
                        reason = f"Your spending shows an upward trend. I recommend increasing your budget by 15% to {recommended_budget:,.0f} to accommodate this growth pattern and provide a buffer for continued increases."
                    elif has_trend and monthly_values[-1] < monthly_values[0]:
                        # Downward trend - decrease by 10%
                        recommended_budget = avg_expense * 0.90
                        action = "decrease"
                        reason = f"Your spending shows a downward trend. I recommend decreasing your budget by 10% to {recommended_budget:,.0f} to reflect this reduction pattern."
                    elif std_dev > avg_expense * 0.05:
                        # High variation - increase by 20%
                        recommended_budget = avg_expense * 1.20
                        action = "increase"
                        reason = f"Your spending shows high variability (std_dev: {std_dev:,.0f}). I recommend increasing your budget by 20% to {recommended_budget:,.0f} to create a safety buffer for unpredictable expenses."
                else:
                    # Even for stable spending, add inflation buffer
                    recommended_budget = avg_expense * 1.05
                    action = "increase"
                    reason = f"While your spending is stable, I recommend adding a 5% buffer ({recommended_budget:,.0f}) to account for inflation and unexpected expenses."
            
            print(f"✅ OpenAI recommendation for {category_name}: {recommended_budget:,.0f} (action: {action}, avg: {avg_expense:,.0f})")
        else:
            # Fallback to rule-based recommendation if OpenAI fails
            recommended_budget = self._calculate_recommended_budget(avg_expense, data)
            reason = self._generate_reason(category_name, avg_expense, recommended_budget)
            action = None
            if not ai_result:
                print(f"❌ OpenAI unavailable (no API key or error), using rule-based for {category_name}: {recommended_budget}")
            else:
                print(f"⚠️ OpenAI returned invalid data, using rule-based for {category_name}: {recommended_budget}")
        
        # 🚨 AGGRESSIVE FINAL VALIDATION: ALWAYS prevent lazy "keep" recommendations
        # Check if recommended_budget is essentially the same as avg_expense (within 2% tolerance)
        tolerance_percent = 0.02  # 2% tolerance
        tolerance = abs(avg_expense * tolerance_percent)
        difference = abs(recommended_budget - avg_expense)
        
        if action == "keep" and difference <= tolerance:
            # They're essentially the same - FORCE an increase
            print(f"🚨 AGGRESSIVE VALIDATION: Overriding lazy 'keep' recommendation")
            print(f"   avg_expense={avg_expense:,.2f}, recommended_budget={recommended_budget:,.2f}, difference={difference:,.2f}, tolerance={tolerance:,.2f}")
            recommended_budget = avg_expense * 1.10  # Force 10% increase
            action = "increase"
            reason = f"Based on your spending pattern, I recommend increasing your budget by 10% to {recommended_budget:,.0f} to account for inflation, variability, and unexpected expenses. This provides a safety buffer for better financial planning."
        
        # Check for data corruption (unreasonably large numbers)
        if recommended_budget > 1e15 or avg_expense > 1e15:
            print(f"🚨 DATA CORRUPTION DETECTED: Numbers are unreasonably large!")
            print(f"   avg_expense={avg_expense:,.2e}, recommended_budget={recommended_budget:,.2e}")
            # Use a safe fallback - cap at reasonable maximum
            if avg_expense > 1e15:
                print(f"   Capping avg_expense from {avg_expense:,.2e} to 1,000,000,000")
                avg_expense = 1e9  # 1 billion as safe maximum
            recommended_budget = avg_expense * 1.10
            action = "increase"
            reason = f"Based on your spending data, I recommend a budget of {recommended_budget:,.0f} (10% increase from average) to account for variability and inflation."
        
        # EXTRA SAFETY: If action is still "keep" and values are very close, force change
        if action == "keep" and difference < (avg_expense * 0.05):  # Within 5%
            print(f"🚨 EXTRA SAFETY CHECK: Forcing change from 'keep' (difference is {difference:,.2f}, which is < 5% of average)")
            recommended_budget = avg_expense * 1.10
            action = "increase"
            reason = f"I recommend increasing your budget by 10% to {recommended_budget:,.0f} to provide a buffer for inflation and unexpected expenses."
        
        # Create recommendation
        return [BudgetRecommendation(
            category=category_name,
            category_id=category_id,
            average_expense=round(avg_expense, 2),
            recommended_budget=round(recommended_budget or 0, 2),
            reason=reason,
            confidence=confidence,
            action=action
        )]
    
    def _calculate_category_statistics(self, expenses: List[Dict], start_date: datetime, end_date: datetime) -> Dict:
        """Calculate statistics for each category"""
        category_data = defaultdict(lambda: {
            "total": 0,
            "count": 0,
            "months": set(),
            "monthly_totals": defaultdict(float)
        })
        
        for expense in expenses:
            category = expense.get("category", "Uncategorized")
            amount = expense.get("amount", 0)
            date = expense.get("date")
            
            # Handle date conversion - skip if date is None or invalid
            if date is None:
                continue
            
            if isinstance(date, str):
                try:
                    date = datetime.fromisoformat(date.replace('Z', '+00:00'))
                except (ValueError, AttributeError):
                    continue
            elif not isinstance(date, datetime):
                # If date is not a string or datetime, skip this expense
                continue
            
            category_data[category]["total"] += amount
            category_data[category]["count"] += 1
            
            # Track monthly totals
            month_key = (date.year, date.month)
            category_data[category]["months"].add(month_key)
            category_data[category]["monthly_totals"][month_key] += amount
        
        # Calculate averages
        result = {}
        for category, data in category_data.items():
            num_months = len(data["months"]) or 1
            avg_monthly = data["total"] / num_months
            
            # Calculate standard deviation for variability
            monthly_values = list(data["monthly_totals"].values())
            if len(monthly_values) > 1:
                mean = sum(monthly_values) / len(monthly_values)
                variance = sum((x - mean) ** 2 for x in monthly_values) / len(monthly_values)
                std_dev = math.sqrt(variance)
            else:
                std_dev = 0
            
            result[category] = {
                "average_monthly": avg_monthly,
                "total": data["total"],
                "count": data["count"],
                "months_analyzed": num_months,
                "std_dev": std_dev,
                "monthly_values": monthly_values
            }
        
        return result
    
    def _calculate_recommended_budget(self, avg_expense: float, data: Dict) -> float:
        """
        Calculate recommended budget based on average expense.
        
        Strategy:
        - Base: Average monthly expense
        - Add 5% buffer for variability
        - Round to nearest 100 for cleaner numbers
        """
        # Add 5% buffer to handle variability
        buffer = avg_expense * 0.05
        
        # If there's high variability (std_dev > 20% of mean), add more buffer
        if data["std_dev"] > 0:
            coefficient_of_variation = data["std_dev"] / avg_expense if avg_expense > 0 else 0
            if coefficient_of_variation > 0.2:
                buffer = avg_expense * 0.10  # 10% buffer for high variability
        
        recommended = avg_expense + buffer
        
        # Round to nearest 100 for cleaner budget numbers
        recommended = round(recommended / 100) * 100
        
        # Ensure minimum of 100 if there was any expense
        if recommended < 100 and avg_expense > 0:
            recommended = 100
        
        return recommended
    
    def _calculate_confidence(self, data: Dict) -> float:
        """
        Calculate confidence score (0-1) based on data quality.
        
        Factors:
        - Number of months analyzed (more = higher confidence)
        - Number of transactions (more = higher confidence)
        - Consistency of spending (lower std_dev = higher confidence)
        """
        months_score = min(data["months_analyzed"] / 6, 1.0)  # Max at 6 months
        count_score = min(data["count"] / 10, 1.0)  # Max at 10 transactions
        
        # Consistency score (inverse of coefficient of variation)
        if data["average_monthly"] > 0:
            cv = data["std_dev"] / data["average_monthly"]
            consistency_score = max(0, 1 - min(cv, 1.0))  # Lower CV = higher score
        else:
            consistency_score = 0.5
        
        # Weighted average
        confidence = (months_score * 0.4 + count_score * 0.3 + consistency_score * 0.3)
        
        return round(confidence, 2)
    
    def _generate_reason(self, category: str, avg_expense: float, recommended_budget: float) -> str:
        """Generate human-readable reason for the recommendation"""
        # Format amounts with currency symbol
        avg_formatted = f"Rs.{avg_expense:,.0f}"
        budget_formatted = f"Rs.{recommended_budget:,.0f}"
        
        if recommended_budget > avg_expense:
            buffer = recommended_budget - avg_expense
            buffer_pct = (buffer / avg_expense * 100) if avg_expense > 0 else 0
            return (
                f"Your average monthly {category.lower()} expense is {avg_formatted}. "
                f"We suggest setting your budget to {budget_formatted} for next month "
                f"(includes a {buffer_pct:.0f}% buffer for variability)."
            )
        else:
            return (
                f"Your average monthly {category.lower()} expense is {avg_formatted}. "
                f"We recommend a budget of {budget_formatted} for next month."
            )
    
    def get_category_averages(self, user_id: str, months: int = 3) -> List[CategoryExpense]:
        """Get average expenses by category for the past N months"""
        end_date = datetime.now()
        start_date = end_date - timedelta(days=months * 30)
        
        expenses = list(self.db.expenses.find({
            "user_id": user_id,
            "date": {"$gte": start_date, "$lte": end_date},
            "type": "expense"
        }))
        
        if not expenses:
            return []
        
        category_data = self._calculate_category_statistics(expenses, start_date, end_date)
        
        result = []
        for category, data in category_data.items():
            result.append(CategoryExpense(
                category=category,
                average_monthly_expense=round(data["average_monthly"], 2),
                total_expenses=data["count"],
                months_analyzed=data["months_analyzed"]
            ))
        
        result.sort(key=lambda x: x.average_monthly_expense, reverse=True)
        return result

    def _get_category_id_by_name(self, category_name: str) -> Optional[str]:
        """
        Find category_id by category name from headCategories and categories collections.
        Returns the first matching category_id found.
        """
        if not category_name:
            return None
        
        try:
            # Search in headCategories collection by name
            head_category = self.db.headCategories.find_one({
                "$or": [
                    {"name": {"$regex": category_name, "$options": "i"}},
                    {"headCategoryName": {"$regex": category_name, "$options": "i"}},
                    {"categoryName": {"$regex": category_name, "$options": "i"}}
                ]
            })
            
            if head_category:
                # Return _id as string
                category_id = str(head_category.get("_id"))
                print(f"✅ Found category_id in headCategories: '{category_name}' -> {category_id}")
                return category_id
            
            # Search in categories collection by name
            category = self.db.categories.find_one({
                "$or": [
                    {"name": {"$regex": category_name, "$options": "i"}},
                    {"categoryName": {"$regex": category_name, "$options": "i"}}
                ]
            })
            
            if category:
                category_id = str(category.get("_id"))
                print(f"✅ Found category_id in categories: '{category_name}' -> {category_id}")
                return category_id
            
            print(f"⚠️ Category name not found: '{category_name}'")
            return None
            
        except Exception as e:
            print(f"Error looking up category_id by name '{category_name}': {e}")
            return None
    
    def _get_category_name(self, category_id) -> str:
        """
        Look up category name from headCategories and categories collections.
        Checks headCategories first, then categories collection.
        """
        if not category_id:
            return "Uncategorized"
        
        try:
            # Convert to ObjectId if it's a string
            if isinstance(category_id, str):
                try:
                    category_id_obj = ObjectId(category_id)
                except (ValueError, TypeError):
                    category_id_obj = category_id
            else:
                category_id_obj = category_id
            
            # First, try to find in headCategories collection
            head_category_doc = None
            if isinstance(category_id_obj, ObjectId):
                head_category_doc = self.db.headcategories.find_one({"_id": category_id_obj})
            else:
                try:
                    head_category_doc = self.db.headcategories.find_one({"_id": ObjectId(category_id)})
                except (ValueError, TypeError):
                    head_category_doc = self.db.headcategories.find_one({"_id": category_id})
            
            if head_category_doc:
                category_name = head_category_doc.get("name") or head_category_doc.get("title")
                if category_name:
                    print(f"✅ Found category name in headCategories: {category_id} -> {category_name}")
                    return category_name
            
            # If not found in headCategories, try categories collection
            category_doc = None
            if isinstance(category_id_obj, ObjectId):
                category_doc = self.db.categories.find_one({"_id": category_id_obj})
            else:
                try:
                    category_doc = self.db.categories.find_one({"_id": ObjectId(category_id)})
                except (ValueError, TypeError):
                    category_doc = self.db.categories.find_one({"_id": category_id})
            
            if category_doc:
                category_name = category_doc.get("name") or category_doc.get("title")
                if category_name:
                    print(f"✅ Found category name in categories: {category_id} -> {category_name}")
                    return category_name
            
            # Try searching by category ID as string in other fields (fallback)
            if isinstance(category_id, str):
                # Try searching in headCategories by other fields
                fallback_head = self.db.headcategories.find_one({"$or": [
                    {"_id": category_id},
                    {"categoryId": category_id},
                    {"id": category_id}
                ]})
                if fallback_head:
                    category_name = fallback_head.get("name") or fallback_head.get("title")
                    if category_name:
                        print(f"✅ Found category name in headCategories (fallback): {category_id} -> {category_name}")
                        return category_name
                
                # Try searching in categories by other fields
                fallback_cat = self.db.categories.find_one({"$or": [
                    {"_id": category_id},
                    {"categoryId": category_id},
                    {"id": category_id}
                ]})
                if fallback_cat:
                    category_name = fallback_cat.get("name") or fallback_cat.get("title")
                    if category_name:
                        print(f"✅ Found category name in categories (fallback): {category_id} -> {category_name}")
                        return category_name
            
            # If not found, log a warning
            print(f"⚠️ Category ID not found in headCategories or categories: {category_id}")
        except Exception as e:
            print(f"❌ Error looking up category name for {category_id}: {e}")
            pass
        
        # If not found in either collection, log and return the ID as string
        print(f"⚠️ Category ID not found in headCategories or categories collections: {category_id}")
        # Try one more time with string search (in case ID is stored as string in a different field)
        try:
            if isinstance(category_id, str):
                # Try searching by name field containing the ID (unlikely but worth trying)
                head_cat_by_name = self.db.headcategories.find_one({"name": category_id})
                if head_cat_by_name:
                    return head_cat_by_name.get("name") or str(category_id)
                cat_by_name = self.db.categories.find_one({"name": category_id})
                if cat_by_name:
                    return cat_by_name.get("name") or str(category_id)
        except Exception as e:
            print(f"Final fallback search failed: {e}")
        
        return str(category_id) if category_id else "Uncategorized"

    def _get_category_stats_from_budgets(
        self, user_id: str, month: int, year: int
    ) -> Dict:
        """
        Build category stats from existing budgets for this user.

        We treat each budget document (e.g. \"Office Maintenance\", \"LOGICGO\")
        as a spending category and derive an \"average\" from its amounts.
        Also extracts categories from headCategories array.
        """
        budgets = []
        
        print(f"Searching for budgets with user_id: {user_id} (type: {type(user_id).__name__})")
        
        # Try multiple query patterns to find budgets (include both OPEN and CLOSE status)
        # Pattern 1: Try with ObjectId (most common in WalletSync) - no status filter
        try:
            query_objid = {"createdBy": ObjectId(user_id)}
            budgets_objid = list(self.db.budgets.find(query_objid))
            print(f"Pattern 1 (createdBy ObjectId): Found {len(budgets_objid)} budgets")
            if budgets_objid:
                budgets.extend(budgets_objid)
        except (ValueError, TypeError) as e:
            print(f"Pattern 1 failed: {e}")
            pass
        
        # Pattern 2: Try with string user_id - no status filter
        try:
            query_str = {"createdBy": user_id}
            budgets_str = list(self.db.budgets.find(query_str))
            print(f"Pattern 2 (createdBy string): Found {len(budgets_str)} budgets")
            if budgets_str:
                budgets.extend(budgets_str)
        except Exception as e:
            print(f"Pattern 2 failed: {e}")
            pass
        
        # Pattern 3: Try with user_id field (alternative field name) - no status filter
        try:
            query_userid = {"user_id": user_id}
            budgets_userid = list(self.db.budgets.find(query_userid))
            print(f"Pattern 3 (user_id string): Found {len(budgets_userid)} budgets")
            if budgets_userid:
                budgets.extend(budgets_userid)
        except Exception as e:
            print(f"Pattern 3 failed: {e}")
            pass
        
        # Pattern 4: Try ObjectId with user_id field - no status filter
        try:
            query_objid_userid = {"user_id": ObjectId(user_id)}
            budgets_objid_userid = list(self.db.budgets.find(query_objid_userid))
            print(f"Pattern 4 (user_id ObjectId): Found {len(budgets_objid_userid)} budgets")
            if budgets_objid_userid:
                budgets.extend(budgets_objid_userid)
        except (ValueError, TypeError) as e:
            print(f"Pattern 4 failed: {e}")
            pass
        
        # Pattern 5: Check if user_id is actually a budget _id, then get createdBy from it
        try:
            budget_by_id = self.db.budgets.find_one({"_id": ObjectId(user_id)})
            if budget_by_id:
                print(f"Pattern 5: user_id is a budget _id, found budget: {budget_by_id.get('name', 'Unknown')}")
                created_by = budget_by_id.get("createdBy")
                if created_by:
                    # Now find all budgets for this createdBy
                    query_by_creator = {"createdBy": created_by}
                    budgets_by_creator = list(self.db.budgets.find(query_by_creator))
                    print(f"Pattern 5: Found {len(budgets_by_creator)} budgets for createdBy: {created_by}")
                    if budgets_by_creator:
                        budgets.extend(budgets_by_creator)
        except (ValueError, TypeError) as e:
            print(f"Pattern 5 failed: {e}")
            pass
        
        # Pattern 6: Try finding by budget _id as string
        try:
            budget_by_id_str = self.db.budgets.find_one({"_id": user_id})
            if budget_by_id_str:
                print(f"Pattern 6: Found budget by _id as string")
                budgets.append(budget_by_id_str)
        except Exception as e:
            print(f"Pattern 6 failed: {e}")
            pass
        
        # Remove duplicates based on _id
        seen_ids = set()
        unique_budgets = []
        for b in budgets:
            budget_id = str(b.get("_id", ""))
            if budget_id not in seen_ids:
                seen_ids.add(budget_id)
                unique_budgets.append(b)
        
        budgets = unique_budgets

        if not budgets:
            print(f"No budgets found for user_id: {user_id}")
            print(f"Tried all query patterns. Checking sample budget structure...")
            # Get a sample budget to see the structure
            sample = self.db.budgets.find_one()
            if sample:
                print(f"Sample budget structure - createdBy type: {type(sample.get('createdBy')).__name__}, value: {sample.get('createdBy')}")
                print(f"Sample budget has user_id field: {'user_id' in sample}")
            return {}

        print(f"Found {len(budgets)} budgets for user_id: {user_id}")

        result: Dict[str, Dict] = {}
        for b in budgets:
            # Extract category ID from budget (could be in category, categoryId, headCategory fields)
            category_id = b.get("category") or b.get("categoryId") or b.get("headCategory") or b.get("category_id")
            
            # Also check if category is nested in headCategories array
            if not category_id:
                head_categories = b.get("headCategories", [])
                if head_categories and isinstance(head_categories, list):
                    # Try to get category from first headCategory's categories array
                    for head_cat in head_categories:
                        if isinstance(head_cat, dict):
                            nested_categories = head_cat.get("categories", [])
                            if nested_categories and isinstance(nested_categories, list):
                                # Get first category ID from nested categories
                                for nested_cat in nested_categories:
                                    if isinstance(nested_cat, dict):
                                        category_id = nested_cat.get("category")
                                        if category_id:
                                            break
                            if category_id:
                                break
            
            # Get category name from headCategories or categories collection using category ID
            if category_id:
                print(f"🔍 Looking up category ID: {category_id} (type: {type(category_id).__name__})")
                
                # Check if category_id is already a string name (not a valid ObjectId)
                if isinstance(category_id, str):
                    # Check if it's a valid ObjectId format (24 hex characters)
                    is_valid_objectid = len(category_id) == 24 and all(c in '0123456789abcdefABCDEF' for c in category_id)
                    if not is_valid_objectid:
                        # It's already a category name, use it directly
                        category_name = category_id
                        print(f"✅ Using category name directly (not an ObjectId): '{category_name}'")
                    else:
                        # It's a valid ObjectId string, try to look it up
                        category_name = self._get_category_name(category_id)
                        if category_name == str(category_id):
                            # Category name lookup failed, still showing ID
                            print(f"⚠️ Category ID not resolved: {category_id} (not found in headCategories or categories collections)")
                            print(f"   This means the category ID doesn't exist in the database. Please check if the category exists.")
                        else:
                            print(f"✅ Found category ID: {category_id} -> Name: '{category_name}'")
                else:
                    # It's an ObjectId object, look it up
                    category_name = self._get_category_name(category_id)
                    if category_name == str(category_id):
                        # Category name lookup failed, still showing ID
                        print(f"⚠️ Category ID not resolved: {category_id} (not found in headCategories or categories collections)")
                        print(f"   This means the category ID doesn't exist in the database. Please check if the category exists.")
                    else:
                        print(f"✅ Found category ID: {category_id} -> Name: '{category_name}'")
            else:
                # Fallback to budget name if no category ID found
                category_name = b.get("name", "Uncategorized")
                if not category_name or category_name == "Uncategorized":
                    category_name = b.get("title") or "Uncategorized"
                print(f"⚠️ No category ID found, using budget name: '{category_name}'")
            
            # Skip if category name is still Uncategorized or empty
            if not category_name or category_name == "Uncategorized" or category_name.strip() == "":
                print(f"⚠️ Skipping budget with invalid category name: {b.get('_id')}")
                continue

            print(f"✅ Processing budget: '{category_name}' (budget id: {b.get('_id')}, category id: {category_id})")

            # Derive a base amount from WalletSync fields
            try:
                max_amount = float(b.get("maxAmount", 0) or b.get("max_amount", 0) or b.get("amount", 0) or 0)
                spend_amount = float(b.get("spendAmount", 0) or b.get("spend_amount", 0) or b.get("spent", 0) or 0)
                budget_amount = float(b.get("budget", 0) or b.get("budgetAmount", 0) or 0)
            except (ValueError, TypeError):
                max_amount = 0
                spend_amount = 0
                budget_amount = 0

            # Priority: spendAmount > maxAmount > budgetAmount > budget
            if spend_amount > 0:
                base_amount = spend_amount
            elif max_amount > 0:
                base_amount = max_amount
            elif budget_amount > 0:
                base_amount = budget_amount
            else:
                base_amount = 0

            # Only add budget if it has an amount - use category name as key
            # Store both category_name and category_id in the result
            if base_amount > 0:
                # Use a unique key that includes user_id and category_id to ensure user-specific grouping
                # This prevents budgets from different users with same category name from being mixed
                category_id_str = str(category_id) if category_id else "none"
                result_key = f"{user_id}|{category_name}|{category_id_str}"
                
                if result_key not in result:
                    result[result_key] = {
                        "category_name": category_name,
                        "category_id": str(category_id) if category_id else None,
                    "average_monthly": base_amount,
                    "total": base_amount,
                    "count": 1,
                    "months_analyzed": 1,
                    "std_dev": 0.0,
                    "monthly_values": [base_amount],
                }
            else:
                    result[result_key]["total"] += base_amount
                    result[result_key]["count"] += 1
                    result[result_key]["months_analyzed"] = result[result_key]["count"]
                    result[result_key]["average_monthly"] = (
                        result[result_key]["total"] / result[result_key]["count"]
                    )
                    result[result_key]["monthly_values"].append(base_amount)

        # Extract category names for logging (skip user_id part in key)
        category_names = []
        for key, data in result.items():
            key_parts = key.split("|")
            if len(key_parts) >= 2:
                category_names.append(key_parts[1])  # Get category_name (second part after user_id)
            else:
                category_names.append(data.get("category_name", key))
        print(f"✅ Processed {len(result)} budget categories for user {user_id}: {category_names}")
        return result

    def _get_ai_recommendation(self, category: str, data: Dict, avg_expense: float):
        """Use OpenAI to refine the budget recommendation."""
        if not OPENAI_API_KEY:
            print(f"⚠️ OpenAI API key not found in environment variables for category: {category}")
            return None

        print(f"🔄 Calling OpenAI API for category: {category}...")

        # Handle empty monthly_values
        if not data.get("monthly_values") or len(data["monthly_values"]) == 0:
            history = f"{avg_expense:.0f}"
        else:
            history = ", ".join(f"{value:.0f}" for value in data["monthly_values"])
        
        # Build comprehensive data summary for OpenAI to analyze
        monthly_values = data.get("monthly_values", [])
        
        # Calculate trend analysis
        trend_analysis = ""
        if len(monthly_values) > 1:
            first_half = monthly_values[:len(monthly_values)//2]
            second_half = monthly_values[len(monthly_values)//2:]
            first_avg = sum(first_half) / len(first_half) if first_half else avg_expense
            second_avg = sum(second_half) / len(second_half) if second_half else avg_expense
            
            if second_avg > first_avg * 1.05:
                trend_analysis = f"UPWARD TREND: Early period average ({first_avg:,.0f}) vs Recent period average ({second_avg:,.0f}) - spending is increasing by {((second_avg/first_avg - 1) * 100):.1f}%"
            elif second_avg < first_avg * 0.95:
                trend_analysis = f"DOWNWARD TREND: Early period average ({first_avg:,.0f}) vs Recent period average ({second_avg:,.0f}) - spending is decreasing by {((1 - second_avg/first_avg) * 100):.1f}%"
            else:
                trend_analysis = f"STABLE TREND: Early period average ({first_avg:,.0f}) vs Recent period average ({second_avg:,.0f}) - spending is relatively stable"
        else:
            trend_analysis = "INSUFFICIENT DATA: Only one data point available"
        
        # Calculate coefficient of variation
        cv = (data['std_dev'] / avg_expense * 100) if avg_expense > 0 else 0
        variability_level = ""
        if cv > 20:
            variability_level = "HIGH VARIABILITY - spending is very unpredictable"
        elif cv > 10:
            variability_level = "MODERATE VARIABILITY - some unpredictability"
        elif cv > 5:
            variability_level = "LOW VARIABILITY - relatively predictable"
        else:
            variability_level = "VERY LOW VARIABILITY - very predictable spending"
        
        # Check if this is a new budget (no historical data)
        is_new_budget = len(monthly_values) == 1 and data.get('months_analyzed', 0) == 1
        
        if is_new_budget:
            # This is a new budget - no historical data
            summary = (
                f"Category: {category}\n"
                f"⚠️ IMPORTANT: This is a NEW BUDGET with NO historical spending data.\n"
                f"💰 USER'S BUDGET AMOUNT: The user has SET/PLANNED a budget of {avg_expense:,.2f} for this category.\n"
                f"This is the budget amount the user wants to allocate - this is the ONLY data point available.\n"
                f"There is NO spending history to analyze - this is a fresh budget.\n\n"
                f"🎯 YOUR TASK: Provide an INTELLIGENT recommendation based on the user's budget amount of {avg_expense:,.2f}\n\n"
                f"Your recommendation should be based on:\n"
                f"  1. The user's provided budget amount: {avg_expense:,.2f} (this is what they want to set)\n"
                f"  2. Category-specific knowledge (e.g., Food & Drinks inflation, Transport volatility)\n"
                f"  3. General best practices (add 10-15% buffer for new budgets to account for variability)\n"
                f"  4. Economic factors (inflation typically 2-5% annually, category-specific inflation)\n\n"
                f"💡 KEY INSIGHT: The user has indicated they want to budget {avg_expense:,.2f} for this category.\n"
                f"   - ANALYZE if this amount is REASONABLE for the category:\n"
                f"     * If the amount seems TOO LOW for the category → Recommend INCREASE (10-20%)\n"
                f"     * If the amount seems TOO HIGH for the category → Recommend DECREASE (10-20%)\n"
                f"     * If the amount seems REASONABLE → Recommend KEEP or small increase (5-10% buffer)\n"
                f"   - Consider category-specific factors:\n"
                f"     * Food & Drinks: Typically needs 10-15% buffer for inflation and variability\n"
                f"     * Transport: May need 15-20% buffer due to fuel price volatility\n"
                f"     * Entertainment: May need less buffer (5-10%) if amount is reasonable\n"
                f"     * Utilities: Usually stable, 5-10% buffer is sufficient\n"
                f"   - Use your knowledge about typical spending ranges for this category\n"
                f"   - If user's amount is clearly excessive for the category, recommend DECREASE\n"
                f"   - If user's amount is clearly insufficient, recommend INCREASE\n\n"
                f"🚨 CRITICAL: DO NOT ALWAYS RECOMMEND INCREASE!\n"
                f"   - If the user's budget amount ({avg_expense:,.2f}) is already generous for {category}, recommend DECREASE\n"
                f"   - If the amount is reasonable, recommend KEEP with small buffer (5-10%)\n"
                f"   - Only recommend INCREASE if the amount seems insufficient for the category\n\n"
                f"DO NOT reference fake trends or historical patterns - this is a new budget!\n"
                f"Recommend a budget that accounts for typical variability and inflation for this category.\n"
            )
        else:
            # This is based on real historical data
            summary = (
                f"Category: {category}\n"
                f"Monthly spending values: [{history}]\n"
                f"Average monthly spend: {avg_expense:,.2f}\n"
                f"Standard deviation: {data['std_dev']:,.2f}\n"
                f"Coefficient of variation: {cv:.1f}% ({variability_level})\n"
                f"Number of months analyzed: {data['months_analyzed']}\n"
                f"Total spending: {data.get('total', avg_expense * data['months_analyzed']):,.2f}\n"
                f"Trend Analysis: {trend_analysis}\n"
        )

        prompt = (
            "You are an expert global personal finance coach with deep knowledge of:\n"
            "- Spending patterns across different categories worldwide (Food, Transport, Entertainment, Utilities, etc.)\n"
            "- Economic trends and inflation impacts globally\n"
            "- Seasonal variations in spending (holidays, weather, cultural events)\n"
            "- Best practices for budget management and financial planning\n"
            "- Category-specific insights (e.g., Food & Drinks tend to have inflation, Transport varies with fuel prices)\n"
            "- Regional economic factors and currency considerations\n\n"
            "TASK: Analyze the user's spending history intelligently using your knowledge and provide a smart, personalized budget recommendation.\n\n"
            "INTELLIGENT ANALYSIS APPROACH:\n\n"
            "1. CATEGORY-SPECIFIC KNOWLEDGE:\n"
            "   - Food & Drinks: Consider inflation (typically 2-5% annually globally), seasonal spikes, cultural events\n"
            "   - Transport/Travel: Fuel price volatility, seasonal travel patterns, regional variations\n"
            "   - Entertainment: Weekend/holiday variations, seasonal trends, cultural events\n"
            "   - Utilities: Seasonal variations (cooling in summer, heating in winter, regional differences)\n"
            "   - Healthcare: Unpredictable but essential, recommend buffer\n"
            "   - Use your knowledge about how this category typically behaves globally\n\n"
            "2. TREND ANALYSIS:\n"
            "   - TRENDING UPWARD: If spending is increasing, consider if it's:\n"
            "     * Inflation-driven (recommend increase to match inflation + buffer)\n"
            "     * Lifestyle change (recommend increase with caution and explanation)\n"
            "     * One-time spike (recommend keep or slight increase, explain it's temporary)\n"
            "     * Seasonal pattern (recommend increase if entering high-spending season)\n"
            "   - TRENDING DOWNWARD: If spending is decreasing, consider if it's:\n"
            "     * Sustainable reduction (recommend decrease to reflect new pattern)\n"
            "     * Temporary dip (recommend keep with buffer for recovery)\n"
            "     * Seasonal pattern (recommend decrease if entering low-spending season)\n\n"
            "3. VARIABILITY INTELLIGENCE:\n"
            "   - HIGH VARIATION (std_dev > 15%): Indicates unpredictable spending\n"
            "     * Recommend INCREASE by 20-30% to create safety buffer\n"
            "     * Explain that high variability requires larger buffer for financial security\n"
            "   - LOW VARIATION (std_dev < 5%): Indicates stable spending\n"
            "     * Can recommend KEEP or small increase (5-10% for inflation buffer)\n"
            "     * Still consider category-specific factors and economic trends\n\n"
            "4. ECONOMIC CONTEXT:\n"
            "   - Consider global inflation trends (typically 2-5% annually in most countries)\n"
            "   - Factor in category-specific inflation (food inflation often higher than general inflation)\n"
            "   - Account for seasonal price variations (holidays, weather, supply/demand)\n"
            "   - Consider regional economic factors if relevant\n\n"
            "5. BEST PRACTICES:\n"
            "   - Always include a small buffer (5-10%) even for stable spending to handle unexpected expenses\n"
            "   - For new budgets (single data point), be conservative but realistic (10-15% buffer)\n"
            "   - Consider the user's spending history length (more data = more confidence in recommendation)\n"
            "   - Apply the 50/30/20 rule principles when appropriate (needs/wants/savings)\n\n"
            "Given the user's spending history:\n"
            f"{summary}\n\n"
            "🚨 CRITICAL: YOU MUST ANALYZE THE ACTUAL DATA ABOVE!\n\n"
            "YOUR INTELLIGENT RECOMMENDATION PROCESS:\n"
            "STEP 1: ANALYZE THE DATA FIRST (This is mandatory!):\n"
            "   - Look at the 'Monthly spending values' array - what pattern do you see?\n"
            "   - Read the 'Trend Analysis' - is spending increasing, decreasing, or stable?\n"
            "   - Check the 'Coefficient of variation' - how predictable is the spending?\n"
            "   - Calculate: If there's an upward trend, you MUST recommend increase\n"
            "   - Calculate: If variability is high, you MUST recommend increase with larger buffer\n\n"
            "STEP 2: APPLY YOUR KNOWLEDGE:\n"
            "   - Consider category-specific factors (Food inflation, Transport volatility, etc.)\n"
            "   - Factor in economic trends and inflation\n"
            "   - Account for seasonal variations if relevant\n\n"
            "STEP 3: PROVIDE SMART RECOMMENDATION:\n"
            "   - The recommended_budget MUST reflect the data analysis from STEP 1\n"
            "   - If trend shows increase → recommended_budget should be higher than average_expense\n"
            "   - If trend shows decrease → recommended_budget should be lower than average_expense\n"
            "   - If variability is high → recommended_budget should have larger buffer (20-30%)\n"
            "   - Include appropriate buffer for inflation and unexpected expenses\n"
            "   - Your reason MUST reference the specific data patterns you observed\n\n"
            "⚠️ DO NOT give generic recommendations! Base your recommendation on the ACTUAL DATA provided above!\n\n"
            "CRITICAL RULES - READ CAREFULLY:\n"
            "⚠️ DO NOT ALWAYS RECOMMEND 'KEEP' - This is a common mistake. Analyze the data first!\n\n"
            "MANDATORY ANALYSIS STEPS:\n"
            "1. Look at the monthly_values array - is there a trend?\n"
            "   - If values increase over time → MUST recommend INCREASE\n"
            "   - If values decrease over time → MUST recommend DECREASE\n"
            "   - Only if values are truly flat (all same) AND std_dev is very low → can recommend KEEP\n\n"
            "2. Check the std_dev (standard deviation):\n"
            "   - If std_dev > 10% of average → MUST recommend INCREASE (high variability needs buffer)\n"
            "   - If std_dev is moderate (5-10%) → Recommend INCREASE with 10-15% buffer\n"
            "   - Only if std_dev < 3% AND no trend → can recommend KEEP with 5% buffer\n\n"
            "3. Consider inflation and best practices:\n"
            "   - Even if spending is stable, inflation (2-5% annually) means you should INCREASE by at least 5-10%\n"
            "   - Always add a buffer for unexpected expenses (5-15% depending on category)\n\n"
            "4. For single data point or new budgets:\n"
            "   - MUST recommend INCREASE by 10-20% to account for variability\n"
            "   - Never recommend KEEP for new/limited data\n\n"
            "DECISION TREE:\n"
            "- Upward trend? → INCREASE (10-25%)\n"
            "- Downward trend? → DECREASE (5-15%)\n"
            "- High variation (std_dev > 15%)? → INCREASE (20-30%)\n"
            "- Moderate variation (std_dev 5-15%)? → INCREASE (10-20%)\n"
            "- Stable with low variation (std_dev < 3%) AND no trend? → KEEP with 5-10% buffer\n"
            "- Single data point? → INCREASE (10-20%)\n\n"
            "⚠️ IMPORTANT: The recommended_budget MUST be different from average_expense in most cases.\n"
            "Only recommend the same amount if ALL of these are true:\n"
            "1. Spending is perfectly stable (all monthly values identical)\n"
            "2. Std_dev is very low (< 3% of average)\n"
            "3. No upward or downward trend\n"
            "4. Category is highly predictable\n"
            "Even then, add at least 5% buffer for inflation!\n\n"
            "Respond strictly as JSON with the following keys:\n"
            '{ "recommended_budget": number, "action": "increase|decrease|keep", "reason": "string" }.\n\n'
            "The 'reason' field is CRITICAL - it must be UNIQUE and SPECIFIC:\n"
            "🚨 MANDATORY: Each reason MUST be completely UNIQUE - never reuse the same reason!\n\n"
            "UNIQUENESS REQUIREMENTS:\n"
            "1. VARY YOUR LANGUAGE:\n"
            "   - Don't start every reason with 'Your spending shows...'\n"
            "   - Use different opening phrases: 'Analyzing your data...', 'Based on the pattern...', 'I've reviewed...', etc.\n"
            "   - Vary sentence structure and word choice\n\n"
            "2. FOCUS ON DIFFERENT ASPECTS:\n"
            "   - For some recommendations, emphasize the TREND (increasing/decreasing)\n"
            "   - For others, emphasize VARIABILITY (high/low volatility)\n"
            "   - For others, emphasize INFLATION or category-specific factors\n"
            "   - Mix and match - don't always focus on the same thing\n\n"
            "3. REFERENCE SPECIFIC DATA:\n"
            "   - MUST include actual numbers from the data (e.g., 'from 9,400,000 to 10,400,000')\n"
            "   - MUST mention specific percentages (e.g., '10.6% increase', '5.7% coefficient of variation')\n"
            "   - MUST reference the category name and specific characteristics\n\n"
            "4. USE DIFFERENT EXPLANATIONS:\n"
            "   - Sometimes explain inflation impact\n"
            "   - Sometimes explain variability needs\n"
            "   - Sometimes explain trend implications\n"
            "   - Sometimes combine multiple factors\n\n"
            "5. VARY YOUR TONE AND STYLE:\n"
            "   - Some reasons can be more analytical\n"
            "   - Some can be more advisory\n"
            "   - Some can emphasize different benefits\n\n"
            "⚠️ CRITICAL: If you find yourself writing a similar reason, STOP and rewrite it with:\n"
            "   - Different opening phrase\n"
            "   - Different focus (trend vs variability vs inflation)\n"
            "   - Different examples or explanations\n"
            "   - Different sentence structure\n\n"
            "Example of UNIQUE reasons (notice how different they are):\n"
            "- 'Analyzing your Food & Drinks spending, I observe a 10.6% upward trajectory from 9.4M to 10.4M. With food inflation typically at 3-5% annually and your low 5.7% variability, I suggest increasing to 11.2M to accommodate price trends.'\n"
            "- 'Your Transport category displays significant volatility (18% coefficient of variation), indicating unpredictable fuel costs. To ensure financial stability, I recommend a 25% buffer increase to 12.5M.'\n"
            "- 'Based on Entertainment spending patterns, the data shows stability with occasional spikes. Accounting for weekend and holiday variations, a modest 8% increase to 5.4M would provide adequate coverage.'\n\n"
            "Round recommended_budget to nearest 100. Use appropriate currency context in your reasoning.\n"
            "Example reason: 'Your spending on Food & Drinks shows an upward trend (1800 → 2000 over 3 months), "
            "likely due to food inflation (typically 3-5% annually globally). I recommend increasing your budget by 15% "
            "to 2300 to accommodate this trend and provide a buffer for continued price increases and unexpected expenses.'\n"
        )

        try:
            response = requests.post(
                "https://api.openai.com/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {OPENAI_API_KEY}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": "gpt-4o-mini",
                    "messages": [
                        {
                            "role": "system",
                            "content": "You are an expert personal finance coach. CRITICAL: Each recommendation reason MUST be completely unique. Never reuse the same language, phrases, or structure. Vary your explanations, focus on different aspects (trends vs variability vs inflation), and use different sentence structures for each recommendation."
                        },
                        {"role": "user", "content": prompt}
                    ],
                    "temperature": 1.2,  # High temperature for maximum variation and creativity
                    "response_format": {"type": "json_object"},
                },
                timeout=30,
            )
            response.raise_for_status()
            response_data = response.json()
            content = response_data["choices"][0]["message"]["content"]
            return json.loads(content)
        except Exception as exc:
            print(f"OpenAI recommendation error for {category}: {exc}")
            return None