bambuddy/backend/app/services/finance_budget.py

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