"""Budget validation helpers for finance-aware print dispatch.""" import calendar from datetime import datetime, timezone from zoneinfo import ZoneInfo, ZoneInfoNotFoundError from fastapi import HTTPException from sqlalchemy import case, func, select from sqlalchemy.ext.asyncio import AsyncSession from backend.app.models.finance import BudgetReservation, CostCenter, CostCenterMember, WalletTransaction from backend.app.models.print_queue import PrintQueueItem from backend.app.models.settings import Settings from backend.app.models.user import User async def is_billing_enabled(db: AsyncSession) -> bool: # Consider any 'billing_enabled' setting with a true-ish value as enabling billing. result = await db.execute( select(func.count()) .select_from(Settings) .where(Settings.key == "billing_enabled", func.lower(func.coalesce(Settings.value, "")) == "true") ) count = int(result.scalar_one() or 0) return count > 0 async def is_printer_kill_switch_enabled(db: AsyncSession) -> bool: """Return True when billing and the printer kill-switch are both enabled.""" result = await db.execute( select(Settings.key, Settings.value).where(Settings.key.in_(("billing_enabled", "printer_kill_switch_enabled"))) ) values = {key: (value or "").strip().lower() for key, value in result.all()} return values.get("billing_enabled") == "true" and values.get("printer_kill_switch_enabled") == "true" async def _get_budget_window_start_utc(db: AsyncSession) -> datetime: result = await db.execute( select(Settings).where(Settings.key.in_(["finance_budget_reset_day", "finance_budget_reset_timezone"])) ) values = {setting.key: setting.value for setting in result.scalars().all()} desired_day = 1 try: parsed = int(values.get("finance_budget_reset_day") or 1) if 1 <= parsed <= 31: desired_day = parsed except (TypeError, ValueError): pass timezone_name = values.get("finance_budget_reset_timezone") or "UTC" try: tz = ZoneInfo(timezone_name) except ZoneInfoNotFoundError: tz = ZoneInfo("UTC") now = datetime.now(tz) current_month_reset_day = min(desired_day, calendar.monthrange(now.year, now.month)[1]) if now.day < current_month_reset_day: month = now.month - 1 year = now.year if month == 0: month = 12 year -= 1 else: month = now.month year = now.year reset_day = min(desired_day, calendar.monthrange(year, month)[1]) return datetime(year, month, reset_day, tzinfo=tz).astimezone(timezone.utc) async def _cost_center_spend(db: AsyncSession, cost_center_id: int, *, monthly: bool) -> float: spend_expr = case((WalletTransaction.amount < 0, -WalletTransaction.amount), else_=0.0) conditions = [ WalletTransaction.cost_center_id == cost_center_id, WalletTransaction.cost_center_id.is_not(None), ] if monthly: conditions.append(WalletTransaction.created_at >= await _get_budget_window_start_utc(db)) result = await db.execute(select(func.coalesce(func.sum(spend_expr), 0.0)).where(*conditions)) return float(result.scalar() or 0.0) async def _cost_center_open_queue_reservations( db: AsyncSession, cost_center_id: int, *, exclude_queue_item_id: int | None = None, ) -> float: active_queue_reservation = ( select(BudgetReservation.id) .where( BudgetReservation.status == "active", BudgetReservation.source_type == "print_queue", BudgetReservation.source_id == PrintQueueItem.id, ) .exists() ) conditions = [ PrintQueueItem.cost_center_id == cost_center_id, PrintQueueItem.status.in_(("pending", "printing")), ~active_queue_reservation, ] if exclude_queue_item_id is not None: conditions.append(PrintQueueItem.id != exclude_queue_item_id) result = await db.execute(select(func.coalesce(func.sum(PrintQueueItem.estimated_cost), 0.0)).where(*conditions)) return float(result.scalar() or 0.0) async def _cost_center_active_budget_reservations( db: AsyncSession, cost_center_id: int, *, exclude_source_type: str | None = None, exclude_source_id: int | None = None, ) -> float: conditions = [ BudgetReservation.cost_center_id == cost_center_id, BudgetReservation.status == "active", ] if exclude_source_type is not None and exclude_source_id is not None: conditions.append( ~( (BudgetReservation.source_type == exclude_source_type) & (BudgetReservation.source_id == exclude_source_id) ) ) result = await db.execute(select(func.coalesce(func.sum(BudgetReservation.amount), 0.0)).where(*conditions)) return float(result.scalar() or 0.0) async def validate_print_budget( db: AsyncSession, *, cost_center_id: int | None, estimated_cost: float | None, current_user: User | None, quantity: int = 1, exclude_queue_item_id: int | None = None, exclude_reservation_source_type: str | None = None, exclude_reservation_source_id: int | None = None, ) -> None: """Validate that a print can be assigned to a cost center budget.""" if not await is_billing_enabled(db): return if cost_center_id is None: raise HTTPException(status_code=400, detail="Cost center is required when billing is enabled") if estimated_cost is None or estimated_cost <= 0: raise HTTPException(status_code=400, detail="Estimated cost is required for cost center prints") center = await db.scalar(select(CostCenter).where(CostCenter.id == cost_center_id).with_for_update()) if not center: raise HTTPException(status_code=404, detail="Cost center not found") if not center.is_active: raise HTTPException(status_code=400, detail="Cost center is inactive") if current_user is not None and not current_user.is_admin: if center.is_private: if center.owner_user_id != current_user.id: raise HTTPException(status_code=403, detail="You cannot print with this private cost center") else: member = await db.scalar( select(CostCenterMember).where( CostCenterMember.cost_center_id == cost_center_id, CostCenterMember.user_id == current_user.id, ) ) if not member or not member.can_print: raise HTTPException(status_code=403, detail="You cannot print with this cost center") budget_limit = center.monthly_budget if center.monthly_budget is not None else center.total_budget if budget_limit is None: return used = await _cost_center_spend(db, cost_center_id, monthly=center.monthly_budget is not None) reserved = await _cost_center_open_queue_reservations( db, cost_center_id, exclude_queue_item_id=exclude_queue_item_id, ) reserved += await _cost_center_active_budget_reservations( db, cost_center_id, exclude_source_type=exclude_reservation_source_type, exclude_source_id=exclude_reservation_source_id, ) requested = estimated_cost * max(1, quantity) available = float(budget_limit) - used - reserved if requested > available: raise HTTPException( status_code=400, detail=f"Estimated print cost exceeds available cost center budget ({requested:.2f} > {available:.2f})", ) async def create_budget_reservation( db: AsyncSession, *, cost_center_id: int | None, estimated_cost: float | None, current_user: User | None, source_type: str, source_id: int | None, print_archive_id: int | None = None, exclude_queue_item_id: int | None = None, ) -> BudgetReservation | None: if not await is_billing_enabled(db): return None if cost_center_id is None: raise HTTPException(status_code=400, detail="Cost center is required when billing is enabled") await validate_print_budget( db, cost_center_id=cost_center_id, estimated_cost=estimated_cost, current_user=current_user, exclude_queue_item_id=exclude_queue_item_id, exclude_reservation_source_type=source_type, exclude_reservation_source_id=source_id, ) existing = None if source_id is not None: existing = await db.scalar( select(BudgetReservation).where( BudgetReservation.status == "active", BudgetReservation.source_type == source_type, BudgetReservation.source_id == source_id, ) ) if existing is not None: existing.cost_center_id = cost_center_id existing.amount = float(estimated_cost or 0.0) if print_archive_id is not None: existing.print_archive_id = print_archive_id await db.flush() return existing reservation = BudgetReservation( cost_center_id=cost_center_id, amount=float(estimated_cost or 0.0), status="active", source_type=source_type, source_id=source_id, print_archive_id=print_archive_id, ) db.add(reservation) await db.flush() return reservation async def release_budget_reservation( db: AsyncSession, *, source_type: str | None = None, source_id: int | None = None, print_archive_id: int | None = None, status: str = "released", ) -> int: conditions = [BudgetReservation.status == "active"] if print_archive_id is not None: conditions.append(BudgetReservation.print_archive_id == print_archive_id) else: conditions.extend( [ BudgetReservation.source_type == source_type, BudgetReservation.source_id == source_id, ] ) result = await db.execute(select(BudgetReservation).where(*conditions)) reservations = result.scalars().all() for reservation in reservations: reservation.status = status reservation.released_at = datetime.now(timezone.utc) if reservations: await db.flush() return len(reservations)