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