"""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), 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 | 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)