plm-lottery/app/withdrawals/service.py

import logging
from datetime import datetime, timezone

from embit import script
from sqlalchemy import select
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.errors import ApiError
from app.audit.log import write_audit_log
from app.db.base import begin_immediate
from app.db.models import PendingTransaction, Round, User, UtxoEvent, Withdrawal
from app.electrum.client import ElectrumClient
from app.rounds.config import get_round_config
from app.rounds.events import broadcaster
from app.tx.versions import record_version
from app.tx.outcome import BroadcastOutcome, TransactionNetwork, broadcast_with_verdict
from app.wallet.address import is_valid_plm_address
from app.wallet.balance import compute_pending_balance, recompute_balance
from app.wallet.hd import derive_pool_address, derive_user_key
from app.wallet.psbt_builder import (
    BuiltTransaction,
    InsufficientFundsError,
    Utxo,
    build_signed_transaction,
)


logger = logging.getLogger(__name__)


class WithdrawalError(ApiError):
    pass


async def request_withdrawal(
    session: AsyncSession,
    client: ElectrumClient,
    user: User,
    external_address: str,
    amount_sats: int,
    network: TransactionNetwork | None = None,
) -> Withdrawal:
    # Checked before anything else: an address from another chain parses fine as a
    # witness program (see wallet/address.py), so without this the tx would build,
    # broadcast and be irrecoverable rather than fail.
    if not is_valid_plm_address(external_address):
        raise WithdrawalError("invalid_address", "not a valid PLM bech32 address")

    # Withdrawing to your own deposit address is a no-op that costs a network fee,
    # and it breaks two things that assume the recipient and the change are
    # distinguishable by address: the RBF bump would shrink the recipient output
    # instead of the change (tx/broadcast.py:_find_change_output), and
    # compute_pending_balance would count the amount twice (B-17).
    if external_address == user.address:
        raise WithdrawalError(
            "withdrawal_to_own_address",
            "that is your own deposit address — withdraw to an external wallet instead",
        )

    # B-88 (DAT-74, TX-28): the write lock before the first read that decides
    # anything, so the UTXOs selected below are the ones phase 1 then marks spent.
    # Taken only after the checks that need no database at all, so a malformed
    # request never queues for it. No lock within busy_timeout means nothing has been
    # built or broadcast: the caller is asked to retry.
    try:
        await begin_immediate(session)
    except OperationalError as exc:
        raise WithdrawalError(
            "service_busy", "the service is busy right now, nothing was sent — please try again"
        ) from exc

    config = await get_round_config(session)

    # B-80 (WDR-14): never to an address of the platform itself. Coins sent to a round's
    # pool address land among that round's UTXOs, and to the reserve or fee address they
    # leave the user's balance for good; the amount would be spent from the balance and
    # then belong to nobody. Compared lower-cased, since bech32 is case-insensitive.
    platform_addresses = {derive_pool_address().lower(), config.fee_address.strip().lower()}
    platform_addresses.update(
        address.lower() for address in await session.scalars(select(Round.pool_address).distinct()) if address
    )
    platform_addresses.discard("")
    if external_address.strip().lower() in platform_addresses:
        raise WithdrawalError(
            "withdrawal_to_platform_address",
            "that is an address of the platform itself — withdraw to an external wallet instead",
        )

    if amount_sats < config.bet_amount_sats:
        raise WithdrawalError(
            "amount_below_minimum",
            f"amount below the minimum of {config.bet_amount_sats} sats",
            minimum_sats=config.bet_amount_sats,
        )

    unspent = (
        await session.scalars(
            select(UtxoEvent).where(UtxoEvent.user_id == user.id, UtxoEvent.spent_txid.is_(None))
        )
    ).all()
    confirmed_sats = sum(u.amount_sats for u in unspent)
    if confirmed_sats < amount_sats:
        # B-37: cached_balance_sats (== confirmed_sats here) can understate the real
        # balance by a whole unconfirmed change output right after a bet/withdrawal —
        # the UI shows pending_balance_sats instead (compute_pending_balance), which
        # can cover an amount this check would otherwise reject as flatly
        # "insufficient". Distinguish "you don't have the money" from "your money
        # hasn't confirmed yet" so the error doesn't contradict what the user is
        # looking at on screen.
        # B-82 (DEP-13): the displayed balance includes deposits still short of their
        # threshold, so a withdrawal they would cover gets "wait for the confirmation".
        waiting_sats = network.pending_deposit_sats(user.id) if hasattr(network, "pending_deposit_sats") else 0
        pending_inclusive_sats, has_pending = await compute_pending_balance(session, user, waiting_sats)
        if has_pending and pending_inclusive_sats >= amount_sats:
            raise WithdrawalError(
                "balance_pending_confirmation",
                "the requested amount is covered by your pending balance, which has not confirmed yet",
                pending_sats=pending_inclusive_sats - confirmed_sats,
            )
        raise WithdrawalError("insufficient_balance", "insufficient balance", required_sats=amount_sats)

    user_key = derive_user_key(user.derivation_index)
    from_script = script.p2wpkh(user_key.to_public())
    utxos = [Utxo(u.txid, u.vout, u.amount_sats) for u in unspent]

    try:
        built = build_signed_transaction(
            signing_key=user_key,
            from_script=from_script,
            utxos=utxos,
            to_address=external_address,
            amount_sats=amount_sats,
            change_address=user.address,
            fee_rate_sat_vb=config.fee_rate_sat_vb,
            # B-62: "withdraw everything" asks for the whole confirmed balance, which
            # would leave no change output and therefore nothing bump_fee could
            # shrink — the one tx shape RBF cannot rescue, and the UI's default
            # withdrawal path at that. Move a dust limit less instead of producing an
            # unbumpable transaction; amount_sent_sats below records what actually
            # went out, which is already how a fee-deducted withdrawal is reported.
            reduce_amount_to_keep_change=True,
        )
    except InsufficientFundsError as exc:
        raise WithdrawalError(exc.code, str(exc), **exc.params) from exc

    # Persist the intent before broadcasting, and only promote the rows once the
    # network has accepted the tx — same two-phase shape as place_bet (B-08).
    spent_by_key = {(u.txid, u.vout): u for u in unspent}
    for spent in built.spent_utxos:
        spent_by_key[(spent.txid, spent.vout)].spent_txid = built.txid
    await recompute_balance(session, user.id)

    withdrawal = Withdrawal(
        user_id=user.id,
        external_address=external_address,
        amount_requested_sats=amount_sats,
        amount_sent_sats=built.recipient_sats,
        txid=built.txid,
        status="building",
    )
    session.add(withdrawal)
    await session.flush()
    pending = PendingTransaction(
        kind="withdrawal",
        withdrawal_id=withdrawal.id,
        user_id=user.id,
        current_txid=built.txid,
        fee_rate_sat_vb=config.fee_rate_sat_vb,
        raw_tx_hex=built.raw_hex,
        status="building",
    )
    session.add(pending)
    await record_version(session, pending, built.txid, built.raw_hex)  # B-77: the first version, with the row
    await session.commit()

    # B-76 (TX-04, TX-20): only a verified refusal releases the coins. A timeout, a
    # drop after the write, or a refusal nobody could verify may mean the transaction
    # left: the withdrawal would then be paid out while its row said "failed" and its
    # coins were spendable again. Those outcomes leave the rows "building".
    verdict = await broadcast_with_verdict(client, network, built.raw_hex, built.txid)
    if verdict.outcome is BroadcastOutcome.REFUSED:
        try:
            await _release_failed_withdrawal(session, withdrawal, pending, built, user.id, verdict.reason)
        except OperationalError:
            # B-88: the release could not be written (database busy). The rows stay
            # "building", the state committed before the broadcast, and the
            # reconciler resolves them against the chain.
            await session.rollback()
            logger.error(
                "could not release the failed withdrawal %s: database busy, left to the reconciler", built.txid
            )
        raise WithdrawalError("broadcast_failed", f"the network refused the transaction: {verdict.reason}")
    if verdict.outcome is BroadcastOutcome.UNCERTAIN:
        logger.warning(
            "withdrawal %s: broadcast outcome uncertain (%s), left 'building' for the reconciler",
            built.txid,
            verdict.reason,
        )
        await session.refresh(withdrawal)
        broadcaster.publish()
        return withdrawal

    # B-88: after the broadcast nothing is undone. Without the lock the rows stay
    # "building" and the reconciler promotes them; the caller sees the real status.
    try:
        await begin_immediate(session)
    except OperationalError:
        await session.rollback()
        logger.error(
            "withdrawal %s broadcast but the database is busy: left 'building' for the reconciler", built.txid
        )
        await session.refresh(withdrawal)
        return withdrawal
    withdrawal.status = "broadcast"
    pending.status = "pending"
    await write_audit_log(
        session,
        "withdrawal_sent",
        {"txid": built.txid, "amount_sent_sats": built.recipient_sats, "external_address": external_address},
        user_id=user.id,
    )

    await session.commit()
    await session.refresh(withdrawal)
    broadcaster.publish()  # balance just went "pending" — nudge the dashboard to refetch
    return withdrawal


async def _release_failed_withdrawal(
    session: AsyncSession,
    withdrawal: Withdrawal,
    pending: PendingTransaction,
    built: BuiltTransaction,
    user_id: int,
    reason: str,
) -> None:
    """Nothing reached the chain, so free the reserved UTXOs and restore the balance.
    The Withdrawal row is kept (marked "failed") rather than deleted: unlike a bet, a
    withdrawal is an instruction the user gave, and they should be able to see that it
    didn't go through."""
    await begin_immediate(session)  # B-88
    for spent in built.spent_utxos:
        row = await session.scalar(
            select(UtxoEvent).where(UtxoEvent.txid == spent.txid, UtxoEvent.vout == spent.vout)
        )
        if row is not None:
            row.spent_txid = None
    withdrawal.status = "failed"
    withdrawal.txid = None
    pending.status = "failed"
    pending.failure_reason = reason[:128]
    pending.failed_at = datetime.now(timezone.utc)  # B-78 (DAT-36): starts the 14-day watch
    await recompute_balance(session, user_id)
    await write_audit_log(
        session,
        "withdrawal_broadcast_failed",
        {"txid": built.txid, "reason": reason[:200]},
        user_id=user_id,
    )
    await session.commit()
    broadcaster.publish()  # the reserved UTXOs are spendable again — refetch the balance (B-49)