import json
import datetime
from pathlib import Path
from typing import Optional, Tuple, Dict, Any
from utils.date_utils import parse_month, days_remaining, parse_flexible_month
from utils.expenses_core import get_expenses_by_month
from utils.status import get_status
# -----------------------------------------
# STORAGE LOCATION
# -----------------------------------------
BUDGET_FILE = Path("data/budgets.json")
# -----------------------------------------
# BASIC FILE I/O
# -----------------------------------------
def load_budget() -> Dict[str, Any]:
if not BUDGET_FILE.exists():
return {}
return json.loads(BUDGET_FILE.read_text())
def save_budget(data: Dict[str, Any]):
BUDGET_FILE.parent.mkdir(exist_ok=True)
BUDGET_FILE.write_text(json.dumps(data, indent=2))
def list_all_budgets() -> Dict[str, Any]:
"""Return all budgets as { '2025-01': {...}, ... }."""
return load_budget()
# -----------------------------------------
# HELPERS REQUESTED BY MODULES
# -----------------------------------------
def get_current_month() -> str:
"""Return the current month in YYYY-MM format."""
today = datetime.date.today()
return today.strftime("%Y-%m")
def budget_exists_for_month(month: str) -> bool:
"""Check if a budget exists for a given month."""
month_key = parse_month(month)
return month_key in load_budget()
# -----------------------------------------
# BUDGET MANAGEMENT
# -----------------------------------------
def set_budget(total: float, month: str, categories: Optional[Dict[str, float]] = None) -> Dict[str, Any]:
"""
Save a budget for a month.
categories: dict of category->amount OR None for total-only mode.
"""
month_key = parse_month(month)
data = load_budget()
data[month_key] = {
"total": total,
"categories": categories or {},
"updated_at": datetime.datetime.now().isoformat(),
}
save_budget(data)
return data[month_key]
def get_budget_for_month(month: str) -> Optional[Dict[str, Any]]:
return load_budget().get(parse_month(month))
# -----------------------------------------
# CALCULATE REMAINING BUDGET
# -----------------------------------------
def calculate_remaining_budget(month: str, category: Optional[str] = None) -> Tuple[bool, Dict[str, Any]]:
"""
Used by multiple MCP modules.
MUST return: (success, result_dict)
NEVER raise errors for missing budgets — return (False, {...error...})
"""
month_key = parse_month(month)
budget = get_budget_for_month(month_key)
if not budget:
return False, {"error": f"Budget not set for {month_key}"}
# Fetch expenses grouped by category
expenses = get_expenses_by_month(month_key)
category_totals = {}
total_spent = 0
for exp in expenses:
cat = exp["category"]
amt = exp["amount"]
category_totals[cat] = category_totals.get(cat, 0) + amt
total_spent += amt
categories = budget["categories"]
# =====================================================================
# CASE 1: TOTAL-ONLY MODE
# =====================================================================
if not categories:
remaining_total = budget["total"] - total_spent
pct_used = (total_spent / budget["total"] * 100) if budget["total"] else 0
result = {
"budget_mode": "total_only",
"month": month_key,
"total_budget": budget["total"],
"total_spent": total_spent,
"remaining_total": remaining_total,
"utilization_percentage": round(pct_used, 2),
"status": get_status(pct_used),
"days_remaining": days_remaining(),
# Category-specific placeholders (for compatibility with modules)
"remaining": remaining_total,
"budget": budget["total"],
}
return True, result
# =====================================================================
# CASE 2: CATEGORY MODE
# =====================================================================
breakdown = {}
for cat, limit in categories.items():
spent = category_totals.get(cat, 0)
remaining = limit - spent
pct = (spent / limit * 100) if limit else 0
breakdown[cat] = {
"budget": limit,
"spent": spent,
"remaining": remaining,
"percentage_used": round(pct, 2),
"status": get_status(pct),
}
remaining_total = budget["total"] - total_spent
result = {
"budget_mode": "category_mode",
"month": month_key,
"total_budget": budget["total"],
"total_spent": total_spent,
"remaining_total": remaining_total,
"days_remaining": days_remaining(),
"category_breakdown": breakdown,
}
# If module asked for category-specific remaining
if category:
cat_info = breakdown.get(category)
if not cat_info:
# Category not in budget — module can interpret gracefully
return False, {
"error": f"'{category}' is not part of the budget for {month_key}",
"available_categories": list(breakdown.keys()),
}
result["remaining"] = cat_info["remaining"]
result["budget"] = cat_info["budget"]
return True, result
