plm-lottery/app/rounds/scheduler.py

import asyncio
import logging
from datetime import datetime, timedelta, timezone

from embit import script
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import async_sessionmaker

from app.audit.log import write_audit_log
from app.db.models import PendingTransaction, Round, RoundParticipant, User
from app.electrum.listener import ElectrumListener
from app.electrum.scripthash import address_to_scripthash
from app.rounds.config import get_round_config
from app.rounds.draw import draw_winner, header_hex_to_block_hash
from app.rounds.events import broadcaster
from app.rounds.service import open_new_round_if_needed
from app.wallet.hd import derive_pool_key
from app.wallet.plm_network import PLM_MAINNET
from app.wallet.psbt_builder import InsufficientFundsError, Utxo, build_payout_transaction

logger = logging.getLogger(__name__)

_TICK_INTERVAL_SECONDS = 5


class RoundScheduler:
    """Background task implementing flowchart.mmd's DRAW subgraph: closes the
    round on its timer (once any in-flight bets have confirmed), draws a winner
    from the next confirmed block, and broadcasts the payout. The next round only
    opens once this one is fully closed (rounds/service.get_active_round)."""

    def __init__(self, session_factory: async_sessionmaker, listener: ElectrumListener):
        self._session_factory = session_factory
        self._listener = listener

    async def run(self) -> None:
        while True:
            try:
                await self._tick()
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("round scheduler tick failed")
            await asyncio.sleep(_TICK_INTERVAL_SECONDS)

    async def _tick(self) -> None:
        if self._listener.client is None:
            return

        async with self._session_factory() as session:
            round_ = await open_new_round_if_needed(session)
            await session.commit()
            if round_ is None:
                return  # still in the cooldown window after the last round closed
            round_id, status, opened_at = round_.id, round_.status, round_.opened_at
            round_duration_seconds = (await get_round_config(session)).round_duration_seconds

        if status not in ("open", "closing"):
            return  # already drawing/paying_out; progress happens elsewhere

        if status == "open":
            opened_at = opened_at.replace(tzinfo=timezone.utc)
            if datetime.now(timezone.utc) < opened_at + timedelta(seconds=round_duration_seconds):
                return

            async with self._session_factory() as session:
                round_ = await session.get(Round, round_id)
                # No new bets from here on, regardless of how long the pending-bet
                # wait below takes — flip to "closing" immediately so it's observable
                # via /rounds/current (e.g. "round closed, waiting for jackpot
                # confirmation") instead of silently staying "open" past the deadline.
                round_.status = "closing"
                round_.closed_at = datetime.now(timezone.utc)
                await session.commit()
                broadcaster.publish()

        async with self._session_factory() as session:
            # "building" counts as in-flight too: it's a bet mid-broadcast (see
            # bets/service.py's two-phase write). A bet that never confirms is
            # eventually removed by app/tx/reconcile.py, which is what stops this
            # wait from being unbounded.
            pending_count = await session.scalar(
                select(func.count())
                .select_from(RoundParticipant)
                .where(
                    RoundParticipant.round_id == round_id,
                    RoundParticipant.status.in_(("building", "broadcast")),
                )
            )
        if pending_count:
            return  # wait for in-flight bets to confirm before closing; stays "closing"

        await self._close_and_draw(round_id)

    async def _close_and_draw(self, round_id: int) -> None:
        async with self._session_factory() as session:
            round_ = await session.get(Round, round_id)

            participants = (
                await session.scalars(
                    select(RoundParticipant)
                    .where(RoundParticipant.round_id == round_id, RoundParticipant.status == "confirmed")
                    .order_by(RoundParticipant.broadcast_at)
                )
            ).all()

            if not participants:
                round_.status = "closed"
                await write_audit_log(session, "round_closed", {"participants": 0}, round_id=round_id)
                await session.commit()
                broadcaster.publish()
                logger.info("round %s closed with no participants", round_id)
                return

            pool_amount = sum(p.bet_amount_sats for p in participants)
            addresses: list[str] = []
            user_by_address: dict[str, int] = {}
            for p in participants:
                user = await session.get(User, p.user_id)
                addresses.append(user.address)
                user_by_address[user.address] = user.id

            round_.status = "drawing"
            await session.commit()
        broadcaster.publish()

        tip_at_close = self._listener.tip_height
        block_height, block_hash = await self._wait_for_next_block(tip_at_close)
        winner_address = draw_winner(addresses, block_hash)

        async with self._session_factory() as session:
            round_ = await session.get(Round, round_id)
            round_.draw_block_height = block_height
            round_.draw_block_hash = block_hash
            round_.seed_int = str(int(block_hash, 16))
            round_.winner_user_id = user_by_address[winner_address]
            round_.pool_amount_sats = pool_amount
            round_.status = "paying_out"
            await write_audit_log(
                session,
                "winner_drawn",
                {
                    "winner_address": winner_address,
                    "pool_amount_sats": pool_amount,
                    "block_height": block_height,
                    "block_hash": block_hash,
                    "participants": len(addresses),
                },
                user_id=user_by_address[winner_address],
                round_id=round_id,
            )
            await session.commit()
        broadcaster.publish()

        logger.info("round %s: winner=%s pool=%s", round_id, winner_address, pool_amount)
        await self._trigger_payout(round_id)

    async def _wait_for_next_block(self, tip_at_close: int) -> tuple[int, str]:
        while self._listener.tip_height <= tip_at_close or not self._listener.tip_header_hex:
            await asyncio.sleep(_TICK_INTERVAL_SECONDS)
        return self._listener.tip_height, header_hex_to_block_hash(self._listener.tip_header_hex)

    async def _trigger_payout(self, round_id: int) -> None:
        """Three phases, so no DB session is held across a network call (B-18): read
        what's needed, do the chain work, then persist the outcome."""
        client = self._listener.client
        if client is None:
            logger.error("round %s payout deferred: not connected", round_id)
            return

        # --- Phase 1: read (session closed before any network I/O) ---------------
        async with self._session_factory() as session:
            round_ = await session.get(Round, round_id)
            config = await get_round_config(session)
            fee_address = config.fee_address
            fee_rate = config.fee_rate_sat_vb
            pool_amount_sats = round_.pool_amount_sats
            winner_user_id = round_.winner_user_id
            winner = await session.get(User, winner_user_id)
            winner_address = winner.address if winner is not None else None
            await session.commit()  # get_round_config may have created the row

        if not fee_address:
            logger.error(
                "round %s payout blocked: no fee_address configured (set it via the admin endpoint)", round_id
            )
            return
        if winner_address is None:
            logger.error("round %s payout blocked: winner user %s not found", round_id, winner_user_id)
            return

        winner_share = pool_amount_sats * 70 // 100
        commission_share = pool_amount_sats - winner_share  # remainder from rounding goes to fees

        # --- Phase 2: build and broadcast ----------------------------------------
        try:
            pool_key = derive_pool_key()
            pool_script_obj = script.p2wpkh(pool_key.to_public())
            pool_address = pool_script_obj.address(network=PLM_MAINNET)
            entries = await client.listunspent(address_to_scripthash(pool_address))
            utxos = [Utxo(e["tx_hash"], e["tx_pos"], e["value"]) for e in entries if e["height"] > 0]

            built = build_payout_transaction(
                signing_key=pool_key,
                from_script=pool_script_obj,
                utxos=utxos,
                winner_address=winner_address,
                winner_share_sats=winner_share,
                fee_address=fee_address,
                commission_sats=commission_share,
                change_address=pool_address,
                fee_rate_sat_vb=fee_rate,
            )
            await client.broadcast(built.raw_hex)
        except InsufficientFundsError:
            logger.exception("round %s payout failed: insufficient pool UTXOs", round_id)
            return
        except Exception:
            # Anything else — a malformed fee_address (EmbitError), a rejected
            # broadcast, a dead connection. This used to escape all the way to
            # run()'s catch-all, which logged it without recording anything, leaving
            # no trace of *why* the round was stuck (B-05). The round stays in
            # "paying_out" either way: automatic payout retry is still an open gap.
            logger.exception("round %s payout failed", round_id)
            await self._log_payout_failure(round_id, winner_user_id)
            return

        # --- Phase 3: persist -----------------------------------------------------
        async with self._session_factory() as session:
            round_ = await session.get(Round, round_id)
            round_.winner_amount_sats = built.winner_sats
            round_.fee_amount_sats = built.commission_sats
            round_.payout_txid = built.txid
            session.add(
                PendingTransaction(
                    kind="payout",
                    round_id=round_id,
                    current_txid=built.txid,
                    fee_rate_sat_vb=fee_rate,
                    raw_tx_hex=built.raw_hex,
                    status="pending",
                )
            )
            await write_audit_log(
                session,
                "payout_sent",
                {"txid": built.txid, "winner_sats": built.winner_sats, "commission_sats": built.commission_sats},
                user_id=winner_user_id,
                round_id=round_id,
            )
            await session.commit()

        logger.info("round %s payout broadcast: txid=%s", round_id, built.txid)

    async def _log_payout_failure(self, round_id: int, winner_user_id: int | None) -> None:
        """Leaves an operator-visible trace in the audit log for a round stuck in
        "paying_out" — the logs alone don't show up in /admin."""
        try:
            async with self._session_factory() as session:
                await write_audit_log(
                    session,
                    "payout_failed",
                    {"round_id": round_id},
                    user_id=winner_user_id,
                    round_id=round_id,
                )
                await session.commit()
        except Exception:
            logger.exception("could not record the payout failure of round %s", round_id)