from __future__ import annotations import asyncio import secrets from collections import defaultdict from contextlib import suppress from datetime import UTC, datetime from html import escape from typing import Any from zoneinfo import ZoneInfo from pyrogram.enums import ChatMemberStatus, ParseMode from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup from wbb import app, log from wbb.utils.dbgiveaway import ( STATUS_CANCELED, STATUS_CANCELING, STATUS_DRAWING, STATUS_FINISHED, STATUS_RUNNING, add_participant, as_utc, attach_giveaway_message, claim_giveaway_cancel, claim_giveaway_draw, count_participants, create_giveaway, finalize_giveaway_cancel, finish_giveaway, get_giveaway, get_participant, list_participants, mark_participant_refunded, record_reroll, remove_participant, save_pending_winners, utc_now, ) from wbb.utils.dbpoints import ( SOURCE_GIVEAWAY_ENTRY, SOURCE_GIVEAWAY_PARTICIPATION, SOURCE_GIVEAWAY_REFUND, SOURCE_GIVEAWAY_WINNER, InsufficientPoints, adjust_points, get_point_account, get_point_transaction_by_key, ) from wbb.utils.i18n import giveaway_status_label _RANDOM = secrets.SystemRandom() _giveaway_locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock) GIVEAWAY_TIMEZONE = ZoneInfo("Asia/Shanghai") class GiveawayServiceError(RuntimeError): def __init__(self, code: str, message: str, *, status: int = 400): super().__init__(message) self.code = code self.status = status def parse_giveaway_time(raw: str) -> datetime | None: try: local_time = datetime.strptime(raw.strip(), "%Y-%m-%d %H:%M") except ValueError: return None return local_time.replace(tzinfo=GIVEAWAY_TIMEZONE).astimezone(UTC) def format_giveaway_time(value: datetime) -> str: return as_utc(value).astimezone(GIVEAWAY_TIMEZONE).strftime("%Y-%m-%d %H:%M") def giveaway_starts_at(giveaway: dict[str, Any]) -> datetime: return as_utc(giveaway.get("starts_at") or giveaway["created_at"]) def _participant_mention(participant: dict[str, Any]) -> str: user_id = int(participant["user_id"]) label = ( participant.get("display_name") or ("@" + str(participant["username"]) if participant.get("username") else None) or participant.get("first_name") or str(user_id) ) return f'{escape(str(label))}' def _format_prizes(prizes: list[dict[str, Any]]) -> str: lines: list[str] = [] for prize in prizes: points = int(prize.get("points_reward", 0)) suffix = f"(每人 +{points} 积分)" if points else "" lines.append( f"- {escape(str(prize['name']))}:{int(prize['count'])} 人{suffix}" ) return "\n".join(lines) async def render_giveaway(giveaway: dict[str, Any]) -> str: participants = await count_participants(giveaway["giveaway_id"]) status_label = giveaway_status_label(giveaway["status"]) if giveaway["status"] == STATUS_RUNNING and giveaway_starts_at(giveaway) > utc_now(): status_label = "待报名" text = [ f"抽奖 #{escape(giveaway['giveaway_id'])}", f"状态:{escape(status_label)}", f"标题:{escape(giveaway['title'])}", ] if giveaway.get("description"): text.append(f"说明:{escape(giveaway['description'])}") text.extend( [ "奖项:", _format_prizes(giveaway["prizes"]), f"报名开始:{format_giveaway_time(giveaway_starts_at(giveaway))}", f"开奖时间:{format_giveaway_time(giveaway['ends_at'])}", "时区:北京时间", f"参与人数:{participants}", ] ) minimum = int(giveaway.get("minimum_points", 0)) cost = int(giveaway.get("entry_cost", 0)) reward = int(giveaway.get("participation_reward", 0)) if minimum: text.append(f"最低积分:{minimum}") if cost: text.append(f"报名消耗:{cost}") if reward: text.append(f"参与奖励:+{reward}") return "\n".join(text) def join_markup(giveaway_id: str) -> InlineKeyboardMarkup: return InlineKeyboardMarkup( [ [ InlineKeyboardButton( "立即报名", callback_data=f"giveaway_join:{giveaway_id}" ) ] ] ) async def refresh_giveaway_message(giveaway: dict[str, Any]) -> bool: message_id = giveaway.get("message_id") if not message_id: return True try: await app.edit_message_text( int(giveaway["chat_id"]), int(message_id), await render_giveaway(giveaway), parse_mode=ParseMode.HTML, reply_markup=join_markup(giveaway["giveaway_id"]), ) except Exception: return False return True def pick_winners( participants: list[dict[str, Any]], prizes: list[dict[str, Any]] ) -> list[dict[str, Any]]: winners: list[dict[str, Any]] = [] picked_user_ids: set[int] = set() for tier_index, prize in enumerate(prizes): remaining = [ participant for participant in participants if int(participant["user_id"]) not in picked_user_ids ] picked = _RANDOM.sample( remaining, min(int(prize["count"]), len(remaining)) ) for participant in picked: user_id = int(participant["user_id"]) picked_user_ids.add(user_id) winners.append( { "user_id": user_id, "username": participant.get("username"), "first_name": participant.get("first_name"), "display_name": participant.get("display_name"), "tier": str(prize["name"]), "tier_index": tier_index, "points_reward": int(prize.get("points_reward", 0)), } ) return winners def render_winners( giveaway: dict[str, Any], winners: list[dict[str, Any]], *, reroll: bool = False, tier_name: str | None = None, ) -> str: title = "抽奖重抽结果" if reroll else "抽奖结果" lines = [ f"{title} #{escape(giveaway['giveaway_id'])}", f"标题:{escape(giveaway['title'])}", ] if tier_name: lines.append(f"奖项:{escape(tier_name)}") if not winners: lines.append("没有符合条件的参与者。") return "\n".join(lines) lines.append("中奖者:") for winner in winners: points = int(winner.get("points_reward", 0)) suffix = f"(+{points} 积分)" if points else "" lines.append( f"- {escape(str(winner['tier']))}:{_participant_mention(winner)}{suffix}" ) return "\n".join(lines) async def create_and_publish_giveaway(**values: Any) -> dict[str, Any]: giveaway = await create_giveaway(**values) try: sent = await app.send_message( giveaway["chat_id"], await render_giveaway(giveaway), parse_mode=ParseMode.HTML, reply_markup=join_markup(giveaway["giveaway_id"]), disable_web_page_preview=True, ) except Exception: await claim_giveaway_cancel(giveaway["giveaway_id"], giveaway["chat_id"]) await finalize_giveaway_cancel(giveaway["giveaway_id"]) raise await attach_giveaway_message(giveaway["giveaway_id"], sent.chat.id, sent.id) giveaway["message_id"] = sent.id return giveaway async def _is_current_member(chat_id: int, user_id: int) -> bool: try: member = await app.get_chat_member(chat_id, user_id) except Exception: return False return member.status not in {ChatMemberStatus.LEFT, ChatMemberStatus.BANNED} async def join_giveaway( giveaway_id: str, *, user_id: int, username: str | None, first_name: str | None, display_name: str | None = None, verify_membership: bool = True, ) -> tuple[str, dict[str, Any] | None]: async with _giveaway_locks[giveaway_id]: giveaway = await get_giveaway(giveaway_id) if not giveaway: return "missing", None if giveaway["status"] != STATUS_RUNNING: return "closed", giveaway now = utc_now() if giveaway_starts_at(giveaway) > now: return "not_started", giveaway if as_utc(giveaway["ends_at"]) <= now: return "ended", giveaway existing = await get_participant(giveaway_id, user_id) if existing: return "duplicate", giveaway if verify_membership and not await _is_current_member( int(giveaway["chat_id"]), user_id ): return "not_member", giveaway minimum = int(giveaway.get("minimum_points", 0)) entry_cost = int(giveaway.get("entry_cost", 0)) account = await get_point_account(int(giveaway["chat_id"]), user_id) debit_key: str | None = None refund_key: str | None = None pending_debit = False if entry_cost: debit_key, refund_key, pending_debit = await _entry_attempt_keys( giveaway_id, user_id ) available_balance = int(account.get("balance", 0)) + ( entry_cost if pending_debit else 0 ) if available_balance < max(minimum, entry_cost): return "insufficient_points", giveaway if entry_cost and debit_key: try: await adjust_points( chat_id=int(giveaway["chat_id"]), user_id=user_id, delta=-entry_cost, source=SOURCE_GIVEAWAY_ENTRY, idempotency_key=debit_key, reference_id=giveaway_id, reason=f"抽奖 #{giveaway_id} 报名消耗", username=username, first_name=first_name, display_name=display_name, ) except InsufficientPoints: return "insufficient_points", giveaway try: result = await add_participant( giveaway_id=giveaway_id, user_id=user_id, username=username, first_name=first_name, display_name=display_name, ) except Exception: if entry_cost: await _refund_entry( giveaway, user_id=user_id, username=username, first_name=first_name, display_name=display_name, key_suffix="join-failed", idempotency_key=refund_key, ) raise if result != "ok" and entry_cost: await _refund_entry( giveaway, user_id=user_id, username=username, first_name=first_name, display_name=display_name, key_suffix=f"join-{result}", idempotency_key=refund_key, ) if result == "ok": await refresh_giveaway_message(giveaway) return result, giveaway async def _entry_attempt_keys( giveaway_id: str, user_id: int ) -> tuple[str, str, bool]: base = f"giveaway-entry:{giveaway_id}:{user_id}" for attempt in range(100): suffix = "" if attempt == 0 else f":{attempt}" debit_key = f"{base}{suffix}" refund_key = f"giveaway-refund:{giveaway_id}:{user_id}:entry-attempt:{attempt}" debit = await get_point_transaction_by_key(debit_key) if not debit: return debit_key, refund_key, False if not await get_point_transaction_by_key(refund_key): return debit_key, refund_key, True raise GiveawayServiceError( "entry_retry_limit", "报名重试次数过多,请稍后再试。", status=409 ) async def _refund_entry( giveaway: dict[str, Any], *, user_id: int, username: str | None, first_name: str | None, display_name: str | None, key_suffix: str, idempotency_key: str | None = None, ) -> bool: entry_cost = int(giveaway.get("entry_cost", 0)) if entry_cost <= 0: return False _, created = await adjust_points( chat_id=int(giveaway["chat_id"]), user_id=user_id, delta=entry_cost, source=SOURCE_GIVEAWAY_REFUND, idempotency_key=idempotency_key or f"giveaway-refund:{giveaway['giveaway_id']}:{user_id}:{key_suffix}", reference_id=giveaway["giveaway_id"], reason=f"抽奖 #{giveaway['giveaway_id']} 报名退款", username=username, first_name=first_name, display_name=display_name, ) if created: await mark_participant_refunded( giveaway["giveaway_id"], user_id, key_suffix ) return created async def _award_draw_points( giveaway: dict[str, Any], participants: list[dict[str, Any]], winners: list[dict[str, Any]], ) -> None: participation_reward = int(giveaway.get("participation_reward", 0)) if participation_reward: for participant in participants: await adjust_points( chat_id=int(giveaway["chat_id"]), user_id=int(participant["user_id"]), delta=participation_reward, source=SOURCE_GIVEAWAY_PARTICIPATION, idempotency_key=( f"giveaway-participation:{giveaway['giveaway_id']}:" f"{participant['user_id']}" ), reference_id=giveaway["giveaway_id"], reason=f"抽奖 #{giveaway['giveaway_id']} 参与奖励", username=participant.get("username"), first_name=participant.get("first_name"), display_name=participant.get("display_name"), ) for winner in winners: reward = int(winner.get("points_reward", 0)) if not reward: continue await adjust_points( chat_id=int(giveaway["chat_id"]), user_id=int(winner["user_id"]), delta=reward, source=SOURCE_GIVEAWAY_WINNER, idempotency_key=( f"giveaway-winner:{giveaway['giveaway_id']}:" f"{winner.get('tier_index', 0)}:{winner['user_id']}" ), reference_id=giveaway["giveaway_id"], reason=( f"抽奖 #{giveaway['giveaway_id']} 中奖奖励" f"({winner['tier']})" ), username=winner.get("username"), first_name=winner.get("first_name"), display_name=winner.get("display_name"), ) async def finish_and_publish_giveaway( giveaway_id: str, *, publish: bool = True ) -> tuple[bool, str, dict[str, Any] | None]: async with _giveaway_locks[giveaway_id]: giveaway = await claim_giveaway_draw(giveaway_id) if not giveaway: existing = await get_giveaway(giveaway_id) if not existing: return False, "未找到该抽奖。", None if existing["status"] == STATUS_FINISHED: return False, "该抽奖已经完成开奖。", existing return False, "该抽奖当前不在进行中。", existing participants = await list_participants(giveaway_id) if giveaway.get("pending_winners") is None: pending = pick_winners(participants, giveaway["prizes"]) giveaway = await save_pending_winners(giveaway_id, pending) or giveaway winners = list(giveaway.get("pending_winners", [])) await _award_draw_points(giveaway, participants, winners) marked = await finish_giveaway(giveaway_id, winners) if not marked: current = await get_giveaway(giveaway_id) if not current or current["status"] != STATUS_FINISHED: return False, "抽奖无法完成,请稍后重试。", current giveaway = current winners = list(current.get("winners", [])) else: giveaway = {**giveaway, "status": STATUS_FINISHED, "winners": winners} if publish: await publish_finished(giveaway, winners) return True, f"抽奖 #{giveaway_id} 已完成开奖。", giveaway async def publish_finished( giveaway: dict[str, Any], winners: list[dict[str, Any]] ) -> None: chat_id = int(giveaway["chat_id"]) message_id = giveaway.get("message_id") text = render_winners(giveaway, winners) result_posted = False if message_id: with suppress(Exception): await app.edit_message_text( chat_id, int(message_id), await render_giveaway({**giveaway, "status": STATUS_FINISHED}), parse_mode=ParseMode.HTML, reply_markup=None, ) with suppress(Exception): await app.send_message( chat_id, text, parse_mode=ParseMode.HTML, reply_to_message_id=int(message_id), ) result_posted = True if not result_posted: with suppress(Exception): await app.send_message(chat_id, text, parse_mode=ParseMode.HTML) for winner in winners: with suppress(Exception): await app.send_message( int(winner["user_id"]), ( f"恭喜你在 {escape(giveaway['title'])} 中获得" f"{escape(str(winner['tier']))}。" ), parse_mode=ParseMode.HTML, ) async def cancel_and_refund_giveaway( giveaway_id: str, *, chat_id: int, publish: bool = True ) -> tuple[bool, str, dict[str, Any] | None]: async with _giveaway_locks[giveaway_id]: giveaway = await claim_giveaway_cancel(giveaway_id, chat_id) if not giveaway: existing = await get_giveaway(giveaway_id) if not existing: return False, "未找到该抽奖。", None if existing["status"] == STATUS_CANCELED: return False, "该抽奖已经取消。", existing return False, "该抽奖当前不在进行中。", existing participants = await list_participants(giveaway_id) for participant in participants: await _refund_entry( giveaway, user_id=int(participant["user_id"]), username=participant.get("username"), first_name=participant.get("first_name"), display_name=participant.get("display_name"), key_suffix="cancel", ) await finalize_giveaway_cancel(giveaway_id) giveaway = {**giveaway, "status": STATUS_CANCELED} if publish and giveaway.get("message_id"): with suppress(Exception): await app.edit_message_text( int(giveaway["chat_id"]), int(giveaway["message_id"]), await render_giveaway(giveaway), parse_mode=ParseMode.HTML, reply_markup=None, ) return True, f"抽奖 #{giveaway_id} 已取消,报名积分已退还。", giveaway async def remove_and_optionally_refund_participant( *, giveaway_id: str, user_id: int, moderator_id: int, reason: str, refund: bool = True, ) -> bool: giveaway = await get_giveaway(giveaway_id) if not giveaway or giveaway["status"] != STATUS_RUNNING: raise GiveawayServiceError( "giveaway_not_running", "该抽奖当前不在进行中。" ) participant = await remove_participant( giveaway_id=giveaway_id, user_id=user_id, moderator_id=moderator_id, reason=reason, refund=refund, ) if not participant: return False if refund: await _refund_entry( giveaway, user_id=user_id, username=participant.get("username"), first_name=participant.get("first_name"), display_name=participant.get("display_name"), key_suffix="participant-removed", ) await refresh_giveaway_message(giveaway) return True async def reroll_giveaway( *, giveaway_id: str, moderator_id: int, tier_name: str | None = None, reroll_id: str | None = None, publish: bool = True, ) -> tuple[list[dict[str, Any]], str]: reroll_id = reroll_id or secrets.token_hex(8) async with _giveaway_locks[giveaway_id]: giveaway = await get_giveaway(giveaway_id) if not giveaway or giveaway["status"] != STATUS_FINISHED: raise GiveawayServiceError( "giveaway_not_finished", "只有已开奖的抽奖可以重抽。" ) for previous in giveaway.get("rerolls", []): if previous.get("reroll_id") == reroll_id: return list(previous.get("winners", [])), reroll_id prizes = list(giveaway["prizes"]) if tier_name: prizes = [ prize for prize in prizes if str(prize["name"]).strip().lower() == tier_name.strip().lower() ] if not prizes: raise GiveawayServiceError("tier_not_found", "未找到该奖项。") excluded = {int(winner["user_id"]) for winner in giveaway.get("winners", [])} for previous in giveaway.get("rerolls", []): excluded.update( int(winner["user_id"]) for winner in previous.get("winners", []) ) participants = [ participant for participant in await list_participants(giveaway_id) if int(participant["user_id"]) not in excluded ] winners = pick_winners(participants, prizes) await record_reroll( giveaway_id, winners, tier_name, moderator_id, reroll_id=reroll_id, ) for winner in winners: reward = int(winner.get("points_reward", 0)) if reward: await adjust_points( chat_id=int(giveaway["chat_id"]), user_id=int(winner["user_id"]), delta=reward, source=SOURCE_GIVEAWAY_WINNER, idempotency_key=( f"giveaway-reroll-winner:{giveaway_id}:{reroll_id}:" f"{winner.get('tier_index', 0)}:{winner['user_id']}" ), reference_id=giveaway_id, reason=f"抽奖 #{giveaway_id} 重抽中奖奖励", username=winner.get("username"), first_name=winner.get("first_name"), display_name=winner.get("display_name"), ) if publish: with suppress(Exception): await app.send_message( int(giveaway["chat_id"]), render_winners(giveaway, winners, reroll=True, tier_name=tier_name), parse_mode=ParseMode.HTML, disable_web_page_preview=True, ) return winners, reroll_id async def resume_pending_giveaway(giveaway: dict[str, Any]) -> None: try: if giveaway["status"] in {STATUS_RUNNING, STATUS_DRAWING}: await finish_and_publish_giveaway(giveaway["giveaway_id"]) elif giveaway["status"] == STATUS_CANCELING: await cancel_and_refund_giveaway( giveaway["giveaway_id"], chat_id=int(giveaway["chat_id"]) ) except Exception as exc: log.error( f"抽奖 {giveaway.get('giveaway_id')} 恢复任务失败:{exc}" )