SurfSense/surfsense_backend/app/services/wallet_credit.py

"""Shared credit-wallet primitives for the flat-rate per-unit billers.

Both :class:`app.services.web_crawl_credit_service.WebCrawlCreditService` and
:class:`app.services.platform_scrape_credit_service.PlatformScrapeCreditService`
follow the same gate -> pre-check -> post-charge model against the unified
``User.credit_micros_*`` wallet. The wallet math lives here once instead of
being copied per service:

- :func:`spendable_micros` — ``balance - reserved`` (ungated by any flag)
- :func:`check_balance` — raise :class:`InsufficientCreditsError` if short
- :func:`apply_debit` — debit + commit + best-effort auto-reload

``InsufficientCreditsError`` is re-exported from ``etl_credit_service`` so every
per-unit biller (ETL, crawl, platform scrape) shares one "out of credit" type —
the capability doors already catch exactly that one.
"""

from uuid import UUID

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.services.etl_credit_service import InsufficientCreditsError

__all__ = [
    "InsufficientCreditsError",
    "apply_debit",
    "check_balance",
    "spendable_micros",
]


async def spendable_micros(session: AsyncSession, user_id: str | UUID) -> int:
    """Raw ``balance - reserved`` read, **ungated** by any billing flag."""
    from app.db import User

    result = await session.execute(
        select(User.credit_micros_balance, User.credit_micros_reserved).where(
            User.id == user_id
        )
    )
    row = result.first()
    if not row:
        raise ValueError(f"User with ID {user_id} not found")

    balance, reserved = row
    return balance - reserved


async def check_balance(
    session: AsyncSession, user_id: str | UUID, required_micros: int
) -> None:
    """Raise :class:`InsufficientCreditsError` if the wallet can't cover
    ``required_micros``. Generic and **ungated** — the caller decides when at
    least one relevant biller is enabled. No-op for a non-positive requirement.
    """
    if required_micros <= 0:
        return
    available = await spendable_micros(session, user_id)
    if required_micros > available:
        raise InsufficientCreditsError(
            message=(
                "This run would exceed your available credit. "
                f"Available: ${available / 1_000_000:.2f}, "
                f"estimated need: ${required_micros / 1_000_000:.2f}. "
                "Add more credits to continue."
            ),
            balance_micros=available,
            required_micros=required_micros,
        )


async def apply_debit(
    session: AsyncSession, user_id: str | UUID, cost_micros: int
) -> int | None:
    """Debit ``cost_micros`` from the wallet and commit.

    Flushes any audit row the caller staged before this, then fires a
    best-effort auto-reload check. No-op for a non-positive cost; returns the
    new balance in micros, or ``None`` when nothing was charged.
    """
    if cost_micros <= 0:
        return None

    from app.db import User

    result = await session.execute(select(User).where(User.id == user_id))
    user = result.unique().scalar_one_or_none()
    if not user:
        raise ValueError(f"User with ID {user_id} not found")

    user.credit_micros_balance -= cost_micros
    await session.commit()
    await session.refresh(user)

    # Best-effort: fire an auto-reload check if the balance dropped low.
    try:
        from app.services.auto_reload_service import maybe_trigger_auto_reload

        await maybe_trigger_auto_reload(user_id)
    except Exception:
        pass

    return user.credit_micros_balance