WealthySmart/backend/app/services/budget_projection.py

import calendar
from datetime import datetime

from sqlmodel import Session, func, select

from app.models.models import (
    RecurringFrequency,
    RecurringItem,
    RecurringItemType,
    Transaction,
    TransactionSource,
    TransactionType,
)


def get_effective_amount(item: RecurringItem, month: int, year: int) -> float | None:
    """Return the effective amount for a recurring item in a given month, or None if inactive."""
    freq = item.frequency

    if freq == RecurringFrequency.MONTHLY:
        if item.override_amounts and str(month) in item.override_amounts:
            return float(item.override_amounts[str(month)])
        return item.amount

    if freq == RecurringFrequency.WEEKLY:
        # Count occurrences of the weekday in this month
        # day_of_month stores day-of-week: 0=Monday
        weekday = item.day_of_month if item.day_of_month is not None else 0
        cal = calendar.monthcalendar(year, month)
        count = sum(1 for week in cal if week[weekday] != 0)
        return item.amount * count

    if freq == RecurringFrequency.QUARTERLY:
        # Active in months 3, 6, 9, 12 by default
        if month % 3 == 0:
            if item.override_amounts and str(month) in item.override_amounts:
                return float(item.override_amounts[str(month)])
            return item.amount
        return None

    if freq == RecurringFrequency.BIANNUAL:
        # Active in month_of_year and 6 months later
        base = item.month_of_year or 1
        second = base + 6 if base <= 6 else base - 6
        if month in (base, second):
            return item.amount
        return None

    if freq == RecurringFrequency.YEARLY:
        if month == (item.month_of_year or 12):
            return item.amount
        return None

    return None


def get_month_range(year: int, month: int) -> tuple[datetime, datetime]:
    """Return (start, end) for a calendar month."""
    start = datetime(year, month, 1)
    if month == 12:
        end = datetime(year + 1, 1, 1)
    else:
        end = datetime(year, month + 1, 1)
    return start, end


def compute_actuals_by_source(
    session: Session, year: int, month: int
) -> dict[str, dict]:
    """Query actual transaction totals for a calendar month, grouped by source."""
    start, end = get_month_range(year, month)

    results = {}
    for source in TransactionSource:
        compra = session.exec(
            select(func.coalesce(func.sum(Transaction.amount), 0)).where(
                Transaction.date >= start,
                Transaction.date < end,
                Transaction.source == source,
                Transaction.transaction_type == TransactionType.COMPRA,
            )
        ).one()
        devolucion = session.exec(
            select(func.coalesce(func.sum(Transaction.amount), 0)).where(
                Transaction.date >= start,
                Transaction.date < end,
                Transaction.source == source,
                Transaction.transaction_type == TransactionType.DEVOLUCION,
            )
        ).one()
        count = session.exec(
            select(func.count()).where(
                Transaction.date >= start,
                Transaction.date < end,
                Transaction.source == source,
                Transaction.transaction_type != TransactionType.DEPOSITO,
            )
        ).one()

        compra_val = float(compra)
        devolucion_val = float(devolucion)
        results[source.value] = {
            "source": source.value,
            "total_compra": compra_val,
            "total_devolucion": devolucion_val,
            "net": compra_val - devolucion_val,
            "count": count,
        }
    return results


def compute_actuals_by_category(
    session: Session, year: int, month: int
) -> dict[int, float]:
    """Return {category_id: net_amount} for actual transactions in a calendar month."""
    start, end = get_month_range(year, month)

    rows = session.exec(
        select(
            Transaction.category_id,
            Transaction.transaction_type,
            func.sum(Transaction.amount),
        )
        .where(
            Transaction.date >= start,
            Transaction.date < end,
            Transaction.category_id.is_not(None),  # type: ignore[union-attr]
            Transaction.transaction_type != TransactionType.DEPOSITO,
        )
        .group_by(Transaction.category_id, Transaction.transaction_type)
    ).all()

    totals: dict[int, float] = {}
    for cat_id, tx_type, amount in rows:
        val = float(amount)
        if tx_type == TransactionType.DEVOLUCION:
            val = -val
        totals[cat_id] = totals.get(cat_id, 0) + val
    return totals


def compute_monthly_projection(
    session: Session, year: int, month: int
) -> dict:
    """Compute full monthly projection with no-double-count logic."""
    items = session.exec(
        select(RecurringItem).where(RecurringItem.is_active == True)  # noqa: E712
    ).all()

    actuals_by_source = compute_actuals_by_source(session, year, month)
    actuals_by_category = compute_actuals_by_category(session, year, month)

    income_items = []
    expense_items = []
    savings_items = []

    total_income = 0.0
    total_fixed_expenses = 0.0
    total_savings = 0.0

    for item in items:
        effective = get_effective_amount(item, month, year)
        if effective is None:
            continue

        detail = {
            "id": item.id,
            "name": item.name,
            "amount": effective,
            "item_type": item.item_type.value,
            "frequency": item.frequency.value,
            "category_name": item.category.name if item.category else None,
            "category_id": item.category_id,
            "used_actual": False,
        }

        if item.item_type == RecurringItemType.INCOME:
            income_items.append(detail)
            total_income += effective

        elif item.item_type == RecurringItemType.EXPENSE:
            # No-double-count: if category has actuals, use actual instead
            if item.category_id and item.category_id in actuals_by_category:
                actual_amount = actuals_by_category[item.category_id]
                detail["amount"] = actual_amount
                detail["projected_amount"] = effective
                detail["used_actual"] = True
                total_fixed_expenses += actual_amount
            else:
                total_fixed_expenses += effective
            expense_items.append(detail)

        elif item.item_type == RecurringItemType.SAVINGS:
            savings_items.append(detail)
            total_savings += effective

    # Sum actuals from sources for categories NOT covered by recurring items
    covered_category_ids = {
        item.category_id
        for item in items
        if item.item_type == RecurringItemType.EXPENSE
        and item.category_id is not None
        and get_effective_amount(item, month, year) is not None
    }

    uncovered_actual = 0.0
    for cat_id, amount in actuals_by_category.items():
        if cat_id not in covered_category_ids:
            uncovered_actual += amount

    # Also add transactions with no category
    start, end = get_month_range(year, month)
    uncategorized = session.exec(
        select(
            Transaction.transaction_type,
            func.sum(Transaction.amount),
        )
        .where(
            Transaction.date >= start,
            Transaction.date < end,
            Transaction.category_id.is_(None),  # type: ignore[union-attr]
            Transaction.transaction_type != TransactionType.DEPOSITO,
        )
        .group_by(Transaction.transaction_type)
    ).all()
    for tx_type, amount in uncategorized:
        val = float(amount)
        if tx_type == TransactionType.DEVOLUCION:
            val = -val
        uncovered_actual += val

    actual_credit_card = actuals_by_source.get("CREDIT_CARD", {}).get("net", 0)
    actual_cash = actuals_by_source.get("CASH", {}).get("net", 0)
    actual_transfers = actuals_by_source.get("TRANSFER", {}).get("net", 0)

    gran_total = total_fixed_expenses + uncovered_actual
    # Savings are NOT deducted — they are already deducted from gross salary
    # (the income amounts are net, post-savings)
    net_balance = total_income - gran_total

    return {
        "year": year,
        "month": month,
        "projected_income": total_income,
        "projected_fixed_expenses": total_fixed_expenses,
        "projected_savings": total_savings,
        "actual_credit_card": actual_credit_card,
        "actual_cash": actual_cash,
        "actual_transfers": actual_transfers,
        "uncovered_actual": uncovered_actual,
        "gran_total_egresos": gran_total,
        "net_balance": net_balance,
        "income_items": income_items,
        "expense_items": expense_items,
        "savings_items": savings_items,
        "actuals_by_source": list(actuals_by_source.values()),
    }