plm-lottery/app/wallet/balance.py

from embit.transaction import Transaction
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import PendingTransaction, User, UtxoEvent
from app.wallet.plm_network import PLM_MAINNET


async def recompute_balance(session: AsyncSession, user_id: int) -> int:
    """Source of truth: sum of this user's confirmed, unspent UTXOs. Updates and
    returns the read-cache column (User.cached_balance_sats). Must be called
    within the same transaction as whatever inserted/updated utxo_events rows."""
    balance = await session.scalar(
        select(func.sum(UtxoEvent.amount_sats)).where(UtxoEvent.user_id == user_id, UtxoEvent.spent_txid.is_(None))
    )
    user = await session.get(User, user_id)
    user.cached_balance_sats = balance or 0
    return user.cached_balance_sats


async def compute_pending_balance(
    session: AsyncSession, user: User, pending_deposit_sats: int = 0
) -> tuple[int, bool]:
    """Balance including the user's own change still in flight.

    Placing a bet or a withdrawal spends whatever UTXOs cover the amount — often
    much larger than the amount actually moving, since select_utxos() picks
    whole UTXOs — and recompute_balance() drops that entire input total from
    cached_balance_sats the moment the tx is broadcast (spent_txid is set right
    away, well before the tx has any confirmations). The change output that
    returns the difference only becomes a UtxoEvent (and so counts toward the
    confirmed balance again) once it gets its own 1st confirmation. In between,
    User.cached_balance_sats alone understates the user's real balance by the
    full unconfirmed change amount, which can look like a much bigger loss than
    the tx actually represents.

    This walks every in-flight (status="pending") bet/withdrawal PendingTransaction
    of this user, decodes its current raw tx (kept up to date across RBF bumps —
    see tx/broadcast.py:bump_fee), and sums whichever outputs pay back to the
    user's own address. Adding that to cached_balance_sats gives the balance the
    user will end up with once everything currently in flight confirms.

    The change output's own confirmation is credited by two independent, unordered
    paths: the Electrum listener (event-driven, near-instant — app/deposits/service.py
    turns it into a UtxoEvent and folds it into cached_balance_sats via
    recompute_balance) and this module's PendingTransaction.status flip
    (app/tx/confirmation.py, polled every 10s). The listener usually wins that race,
    so for the gap until the poller catches up the row is still "pending" here while
    the same sats are already inside cached_balance_sats — double-counting the
    change unless excluded below.

    B-82 (DEP-12 (b), DEP-14): `pending_deposit_sats` is the total of the user's deposits
    that are confirmed but still short of their threshold (the listener keeps them, in
    memory); they are added to what is shown, and make `has_pending` true, exactly like
    change in flight — without ever being spendable.

    Returns (pending_inclusive_balance_sats, has_pending) — has_pending tells the
    caller whether this differs from the confirmed-only balance at all.
    """
    pending = (
        await session.scalars(
            select(PendingTransaction).where(
                PendingTransaction.user_id == user.id,
                PendingTransaction.kind.in_(("bet", "withdrawal")),
                # "building" as well as "pending": a building row's UTXOs are already
                # marked spent (see place_bet's two phases), so leaving it out would
                # make the displayed balance dip for the duration of the broadcast.
                PendingTransaction.status.in_(("building", "pending")),
            )
        )
    ).all()

    already_credited = {
        (txid, vout)
        for txid, vout in (
            await session.execute(select(UtxoEvent.txid, UtxoEvent.vout).where(UtxoEvent.user_id == user.id))
        ).all()
    }

    pending_change_sats = 0
    for row in pending:
        tx = Transaction.parse(bytes.fromhex(row.raw_tx_hex))
        for vout, out in enumerate(tx.vout):
            if (row.current_txid, vout) in already_credited:
                continue
            if out.script_pubkey.address(network=PLM_MAINNET) == user.address:
                pending_change_sats += out.value

    return user.cached_balance_sats + pending_change_sats + pending_deposit_sats, bool(pending) or pending_deposit_sats > 0