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,
update_running_giveaway,
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 _mongo_datetime_millis(value: datetime) -> int:
return int(as_utc(value).timestamp() * 1000)
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
async def update_and_refresh_giveaway(
giveaway_id: str,
*,
title: str,
description: str,
starts_at: datetime,
ends_at: datetime,
expected_updated_at: datetime,
) -> dict[str, Any]:
async with _giveaway_locks[giveaway_id]:
giveaway = await get_giveaway(giveaway_id)
if not giveaway:
raise GiveawayServiceError(
"giveaway_not_found", "未找到该抽奖。", status=404
)
if giveaway["status"] != STATUS_RUNNING:
raise GiveawayServiceError(
"giveaway_not_running", "只有进行中的抽奖可以编辑。", status=409
)
if _mongo_datetime_millis(giveaway["updated_at"]) != _mongo_datetime_millis(
expected_updated_at
):
raise GiveawayServiceError(
"giveaway_stale",
"抽奖已被其他操作更新,请刷新页面后重试。",
status=409,
)
current_starts_at = giveaway_starts_at(giveaway)
if current_starts_at <= utc_now() and _mongo_datetime_millis(
starts_at
) != _mongo_datetime_millis(current_starts_at):
raise GiveawayServiceError(
"giveaway_registration_started",
"报名开始后不能修改报名开始时间。",
status=409,
)
try:
updated = await update_running_giveaway(
giveaway_id,
title=title,
description=description,
starts_at=starts_at,
ends_at=ends_at,
expected_updated_at=giveaway["updated_at"],
)
except ValueError as exc:
raise GiveawayServiceError("invalid_giveaway", str(exc)) from exc
if not updated:
raise GiveawayServiceError(
"giveaway_stale",
"抽奖状态或内容已变化,请刷新页面后重试。",
status=409,
)
if not await refresh_giveaway_message(updated):
raise GiveawayServiceError(
"giveaway_message_refresh_failed",
"抽奖内容已保存,但群消息刷新失败,请刷新页面后重试保存。",
status=502,
)
return updated
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}"
)