"""Budget validation helpers for finance-aware print dispatch.""" import calendar from datetime import datetime, timezone from typing import TYPE_CHECKING 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 if TYPE_CHECKING: from backend.app.core.auth import ApiKeyActor 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), WalletTransaction.is_voided.is_(False), ] 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 get_cost_center_reserved_map( db: AsyncSession, cost_center_ids: list[int], *, exclude_queue_item_id: int | None = None, exclude_reservation_source_type: str | None = None, exclude_reservation_source_id: int | None = None, ) -> dict[int, float]: """Return active holds plus unreserved open queue estimates per cost center. Queue items that already have an active ``print_queue`` reservation are excluded from the queue sum because the reservation is their replacement, not an additional hold. """ if not cost_center_ids: return {} active_queue_reservation = ( select(BudgetReservation.id) .where( BudgetReservation.status == "active", BudgetReservation.source_type == "print_queue", BudgetReservation.source_id == PrintQueueItem.id, ) .exists() ) queue_conditions = [ PrintQueueItem.cost_center_id.in_(cost_center_ids), PrintQueueItem.status.in_(("pending", "printing")), ~active_queue_reservation, ] if exclude_queue_item_id is not None: queue_conditions.append(PrintQueueItem.id != exclude_queue_item_id) queue_rows = await db.execute( select(PrintQueueItem.cost_center_id, func.coalesce(func.sum(PrintQueueItem.estimated_cost), 0.0)) .where(*queue_conditions) .group_by(PrintQueueItem.cost_center_id) ) reserved_map = {int(center_id): float(value) for center_id, value in queue_rows.all() if center_id is not None} reservation_conditions = [ BudgetReservation.cost_center_id.in_(cost_center_ids), BudgetReservation.status == "active", ] if exclude_reservation_source_type is not None and exclude_reservation_source_id is not None: reservation_conditions.append( ~( (BudgetReservation.source_type == exclude_reservation_source_type) & (BudgetReservation.source_id == exclude_reservation_source_id) ) ) reservation_rows = await db.execute( select(BudgetReservation.cost_center_id, func.coalesce(func.sum(BudgetReservation.amount), 0.0)) .where(*reservation_conditions) .group_by(BudgetReservation.cost_center_id) ) for center_id, value in reservation_rows.all(): if center_id is not None: reserved_map[int(center_id)] = reserved_map.get(int(center_id), 0.0) + float(value or 0.0) return reserved_map async def validate_print_budget( db: AsyncSession, *, cost_center_id: int | None, estimated_cost: float | None, current_user: "User | ApiKeyActor | 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_map = await get_cost_center_reserved_map( db, [cost_center_id], exclude_queue_item_id=exclude_queue_item_id, exclude_reservation_source_type=exclude_reservation_source_type, exclude_reservation_source_id=exclude_reservation_source_id, ) reserved = reserved_map.get(cost_center_id, 0.0) 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)