budget-tracker/backend/app/services/budget_service.py

import uuid
from datetime import date

from fastapi import HTTPException, status
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.models.budget import Budget
from app.models.transaction import Transaction, TransactionType
from app.schemas.budget import BudgetCreate, BudgetResponse, BudgetUpdate
from app.utils import utcnow


def _month_bounds(month: str) -> tuple[date, date]:
    year, mon = map(int, month.split("-"))
    start = date(year, mon, 1)
    if mon == 12:
        end = date(year + 1, 1, 1)
    else:
        end = date(year, mon + 1, 1)
    return start, end


def _prev_month(month: str) -> str:
    year, mon = map(int, month.split("-"))
    if mon == 1:
        return f"{year - 1}-12"
    return f"{year}-{mon - 1:02d}"


async def _spent_cents(
    session: AsyncSession,
    user_id: uuid.UUID,
    category_id: uuid.UUID,
    month: str,
) -> int:
    start, end = _month_bounds(month)
    result = await session.execute(
        select(func.coalesce(func.sum(Transaction.amount_cents), 0)).where(
            Transaction.user_id == user_id,
            Transaction.deleted_at.is_(None),
            Transaction.type == TransactionType.expense,
            Transaction.category_id == category_id,
            Transaction.transaction_date >= start,
            Transaction.transaction_date < end,
        )
    )
    return result.scalar_one()


async def _to_response(
    session: AsyncSession,
    user_id: uuid.UUID,
    budget: Budget,
) -> BudgetResponse:
    spent = await _spent_cents(session, user_id, budget.category_id, budget.month)
    return BudgetResponse(
        id=budget.id,
        category_id=budget.category_id,
        category_name=budget.category.name,
        category_color=budget.category.color,
        month=budget.month,
        limit_cents=budget.limit_cents,
        spent_cents=spent,
        remaining_cents=max(0, budget.limit_cents - spent),
        created_at=budget.created_at,
        updated_at=budget.updated_at,
    )


async def list_budgets(
    session: AsyncSession,
    user_id: uuid.UUID,
    month: str,
) -> list[BudgetResponse]:
    result = await session.execute(
        select(Budget)
        .options(selectinload(Budget.category))
        .where(Budget.user_id == user_id, Budget.month == month)
        .order_by(Budget.created_at)
    )
    budgets = result.scalars().all()
    return [await _to_response(session, user_id, b) for b in budgets]


async def create_budget(
    session: AsyncSession,
    user_id: uuid.UUID,
    data: BudgetCreate,
) -> BudgetResponse:
    # Check for duplicate
    existing = await session.execute(
        select(Budget).where(
            Budget.user_id == user_id,
            Budget.category_id == data.category_id,
            Budget.month == data.month,
        )
    )
    if existing.scalar_one_or_none() is not None:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail="Budget already exists for this category and month",
        )

    budget = Budget(
        user_id=user_id,
        category_id=data.category_id,
        month=data.month,
        limit_cents=data.limit_cents,
    )
    session.add(budget)
    await session.commit()

    result = await session.execute(
        select(Budget)
        .options(selectinload(Budget.category))
        .where(Budget.id == budget.id)
    )
    budget = result.scalar_one()
    return await _to_response(session, user_id, budget)


async def update_budget(
    session: AsyncSession,
    user_id: uuid.UUID,
    budget_id: uuid.UUID,
    data: BudgetUpdate,
) -> BudgetResponse:
    result = await session.execute(
        select(Budget)
        .options(selectinload(Budget.category))
        .where(Budget.id == budget_id, Budget.user_id == user_id)
    )
    budget = result.scalar_one_or_none()
    if budget is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Budget not found")

    budget.limit_cents = data.limit_cents
    budget.updated_at = utcnow()
    await session.commit()

    result = await session.execute(
        select(Budget)
        .options(selectinload(Budget.category))
        .where(Budget.id == budget_id)
    )
    budget = result.scalar_one()
    return await _to_response(session, user_id, budget)


async def delete_budget(
    session: AsyncSession,
    user_id: uuid.UUID,
    budget_id: uuid.UUID,
) -> None:
    result = await session.execute(
        select(Budget).where(Budget.id == budget_id, Budget.user_id == user_id)
    )
    budget = result.scalar_one_or_none()
    if budget is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Budget not found")
    await session.delete(budget)
    await session.commit()


async def rollover_budgets(
    session: AsyncSession,
    user_id: uuid.UUID,
    target_month: str,
) -> list[BudgetResponse]:
    """Copy budgets from previous month to target_month (skip existing)."""
    source_month = _prev_month(target_month)
    result = await session.execute(
        select(Budget)
        .options(selectinload(Budget.category))
        .where(Budget.user_id == user_id, Budget.month == source_month)
    )
    source_budgets = result.scalars().all()

    created = []
    for src in source_budgets:
        existing = await session.execute(
            select(Budget).where(
                Budget.user_id == user_id,
                Budget.category_id == src.category_id,
                Budget.month == target_month,
            )
        )
        if existing.scalar_one_or_none() is not None:
            continue
        new_budget = Budget(
            user_id=user_id,
            category_id=src.category_id,
            month=target_month,
            limit_cents=src.limit_cents,
        )
        session.add(new_budget)
        created.append(new_budget)

    await session.commit()

    responses = []
    for b in created:
        result = await session.execute(
            select(Budget)
            .options(selectinload(Budget.category))
            .where(Budget.id == b.id)
        )
        b = result.scalar_one()
        responses.append(await _to_response(session, user_id, b))

    return responses