from __future__ import annotations import re from asyncio import Lock from datetime import UTC, datetime from secrets import token_hex from typing import Any from pymongo import ASCENDING, DESCENDING, ReturnDocument from pymongo.errors import DuplicateKeyError from wbb import BOT_PROFILE_ID, db STATUS_RUNNING = "running" STATUS_DRAWING = "drawing" STATUS_CANCELING = "canceling" STATUS_FINISHED = "finished" STATUS_CANCELED = "canceled" ACTIVE_STATUSES = {STATUS_RUNNING, STATUS_DRAWING, STATUS_CANCELING} ALL_STATUSES = ACTIVE_STATUSES | {STATUS_FINISHED, STATUS_CANCELED} giveawaysdb = db.giveaways participantsdb = db.giveaway_participants giveaway_bansdb = db.giveaway_bans _index_lock = Lock() _indexes_ready = False def utc_now() -> datetime: return datetime.now(UTC) def normalize_giveaway_id(giveaway_id: str) -> str: return giveaway_id.strip().lower().lstrip("#") def as_utc(value: datetime) -> datetime: if value.tzinfo is None: return value.replace(tzinfo=UTC) return value.astimezone(UTC) def _nonnegative_int(value: Any, *, maximum: int = 1_000_000_000) -> int: try: parsed = int(value) except (TypeError, ValueError) as exc: raise ValueError("积分数值必须是整数。") from exc if parsed < 0 or parsed > maximum: raise ValueError(f"积分数值必须在 0 到 {maximum} 之间。") return parsed def normalize_prizes(prizes: list[dict[str, Any]]) -> list[dict[str, Any]]: normalized: list[dict[str, Any]] = [] total_winners = 0 for item in prizes: name = str(item.get("name") or "").strip() count = _nonnegative_int(item.get("count"), maximum=100) points_reward = _nonnegative_int(item.get("points_reward", 0)) if not name or count < 1: raise ValueError("每个奖项都必须填写名称,且中奖人数至少为 1。") total_winners += count normalized.append( {"name": name[:100], "count": count, "points_reward": points_reward} ) if not normalized or total_winners > 100: raise ValueError("抽奖总中奖人数必须在 1 到 100 人之间。") return normalized async def ensure_giveaway_indexes() -> None: global _indexes_ready if _indexes_ready: return async with _index_lock: if _indexes_ready: return await giveawaysdb.create_index([("giveaway_id", ASCENDING)], unique=True) await giveawaysdb.create_index([("status", ASCENDING), ("ends_at", ASCENDING)]) await giveawaysdb.create_index( [("chat_id", ASCENDING), ("created_at", DESCENDING)] ) await participantsdb.create_index( [("giveaway_id", ASCENDING), ("user_id", ASCENDING)], unique=True ) await participantsdb.create_index( [("giveaway_id", ASCENDING), ("active", ASCENDING), ("joined_at", ASCENDING)] ) await giveaway_bansdb.create_index( [("chat_id", ASCENDING), ("user_id", ASCENDING)], unique=True ) await giveaway_bansdb.create_index( [("chat_id", ASCENDING), ("created_at", DESCENDING)] ) _indexes_ready = True async def create_giveaway( *, chat_id: int, creator_id: int, creator_name: str, title: str, description: str, prizes: list[dict[str, Any]], ends_at: datetime, starts_at: datetime | None = None, minimum_points: int = 0, entry_cost: int = 0, participation_reward: int = 0, ) -> dict[str, Any]: await ensure_giveaway_indexes() title = title.strip() if not title: raise ValueError("必须填写抽奖标题。") now = utc_now() normalized_starts_at = as_utc(starts_at) if starts_at else now normalized_ends_at = as_utc(ends_at) if normalized_ends_at <= now: raise ValueError("开奖时间必须晚于当前时间。") if normalized_starts_at >= normalized_ends_at: raise ValueError("报名开始时间必须早于开奖时间。") normalized_prizes = normalize_prizes(prizes) for _ in range(20): giveaway_id = token_hex(3) doc = { "giveaway_id": giveaway_id, "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "creator_id": int(creator_id), "creator_name": creator_name.strip()[:100], "title": title[:200], "description": description.strip()[:4000], "prizes": normalized_prizes, "minimum_points": _nonnegative_int(minimum_points), "entry_cost": _nonnegative_int(entry_cost), "participation_reward": _nonnegative_int(participation_reward), "status": STATUS_RUNNING, "created_at": now, "updated_at": now, "starts_at": normalized_starts_at, "ends_at": normalized_ends_at, "message_id": None, "winners": [], "rerolls": [], } if normalized_starts_at <= now: doc["start_announced_at"] = now try: await giveawaysdb.insert_one(doc) return doc except DuplicateKeyError: continue raise RuntimeError("无法生成唯一的抽奖编号,请重试。") async def attach_giveaway_message( giveaway_id: str, chat_id: int, message_id: int ) -> None: await ensure_giveaway_indexes() await giveawaysdb.update_one( {"giveaway_id": normalize_giveaway_id(giveaway_id)}, { "$set": { "chat_id": int(chat_id), "message_id": int(message_id), "updated_at": utc_now(), } }, ) async def get_giveaway(giveaway_id: str) -> dict[str, Any] | None: await ensure_giveaway_indexes() return await giveawaysdb.find_one( { "giveaway_id": normalize_giveaway_id(giveaway_id), "bot_id": BOT_PROFILE_ID, } ) async def list_running_giveaways( *, chat_id: int | None = None, creator_id: int | None = None, limit: int = 20 ) -> list[dict[str, Any]]: await ensure_giveaway_indexes() query: dict[str, Any] = {"status": STATUS_RUNNING, "bot_id": BOT_PROFILE_ID} if chat_id is not None: query["chat_id"] = int(chat_id) if creator_id is not None: query["creator_id"] = int(creator_id) cursor = giveawaysdb.find(query).sort("created_at", DESCENDING).limit(limit) return [doc async for doc in cursor] async def list_giveaways( *, chat_id: int | None = None, status: str | None = None, limit: int = 50 ) -> list[dict[str, Any]]: items, _ = await list_giveaways_page( chat_id=chat_id, status=status, page=1, page_size=limit ) return items async def list_giveaways_page( *, chat_id: int | None = None, status: str | None = None, query: str = "", page: int = 1, page_size: int = 20, all_bots: bool = False, ) -> tuple[list[dict[str, Any]], int]: await ensure_giveaway_indexes() filters: dict[str, Any] = {} if all_bots else {"bot_id": BOT_PROFILE_ID} if chat_id is not None: filters["chat_id"] = int(chat_id) if status: if status not in ALL_STATUSES: return [], 0 filters["status"] = status query = query.strip() if query: giveaway_id_pattern = re.escape(query.lstrip("#")) title_pattern = re.escape(query) filters["$or"] = [ {"giveaway_id": {"$regex": giveaway_id_pattern, "$options": "i"}}, {"title": {"$regex": title_pattern, "$options": "i"}}, ] page = max(1, int(page)) page_size = max(1, min(int(page_size), 100)) total = await giveawaysdb.count_documents(filters) cursor = ( giveawaysdb.find(filters) .sort("created_at", DESCENDING) .skip((page - 1) * page_size) .limit(page_size) ) return [doc async for doc in cursor], total async def list_due_giveaways(limit: int = 20) -> list[dict[str, Any]]: await ensure_giveaway_indexes() cursor = ( giveawaysdb.find( { "bot_id": BOT_PROFILE_ID, "$or": [ {"status": STATUS_RUNNING, "ends_at": {"$lte": utc_now()}}, {"status": STATUS_DRAWING}, ] } ) .sort("ends_at", ASCENDING) .limit(limit) ) return [doc async for doc in cursor] async def list_unannounced_started_giveaways(limit: int = 20) -> list[dict[str, Any]]: await ensure_giveaway_indexes() now = utc_now() cursor = ( giveawaysdb.find( { "bot_id": BOT_PROFILE_ID, "status": STATUS_RUNNING, "starts_at": {"$lte": now}, "ends_at": {"$gt": now}, "start_announced_at": {"$exists": False}, } ) .sort("starts_at", ASCENDING) .limit(limit) ) return [doc async for doc in cursor] async def mark_giveaway_start_announced(giveaway_id: str) -> None: await ensure_giveaway_indexes() now = utc_now() await giveawaysdb.update_one( { "giveaway_id": normalize_giveaway_id(giveaway_id), "status": STATUS_RUNNING, "start_announced_at": {"$exists": False}, }, {"$set": {"start_announced_at": now, "updated_at": now}}, ) async def list_pending_cancellations(limit: int = 20) -> list[dict[str, Any]]: await ensure_giveaway_indexes() cursor = ( giveawaysdb.find({"status": STATUS_CANCELING, "bot_id": BOT_PROFILE_ID}) .sort("updated_at", ASCENDING) .limit(limit) ) return [doc async for doc in cursor] async def claim_giveaway_draw(giveaway_id: str) -> dict[str, Any] | None: await ensure_giveaway_indexes() normalized_id = normalize_giveaway_id(giveaway_id) claimed = await giveawaysdb.find_one_and_update( {"giveaway_id": normalized_id, "status": STATUS_RUNNING}, { "$set": { "status": STATUS_DRAWING, "draw_started_at": utc_now(), "updated_at": utc_now(), } }, return_document=ReturnDocument.AFTER, ) if claimed: return claimed return await giveawaysdb.find_one( {"giveaway_id": normalized_id, "status": STATUS_DRAWING} ) async def save_pending_winners( giveaway_id: str, winners: list[dict[str, Any]] ) -> dict[str, Any] | None: await ensure_giveaway_indexes() normalized_id = normalize_giveaway_id(giveaway_id) updated = await giveawaysdb.find_one_and_update( { "giveaway_id": normalized_id, "status": STATUS_DRAWING, "pending_winners": {"$exists": False}, }, { "$set": { "pending_winners": winners, "draw_result_created_at": utc_now(), "updated_at": utc_now(), } }, return_document=ReturnDocument.AFTER, ) return updated or await giveawaysdb.find_one( {"giveaway_id": normalized_id, "status": STATUS_DRAWING} ) async def finish_giveaway( giveaway_id: str, winners: list[dict[str, Any]] | None = None ) -> bool: await ensure_giveaway_indexes() values: dict[str, Any] = { "status": STATUS_FINISHED, "finished_at": utc_now(), "updated_at": utc_now(), } if winners is not None: values["winners"] = winners result = await giveawaysdb.update_one( { "giveaway_id": normalize_giveaway_id(giveaway_id), "status": STATUS_DRAWING, }, {"$set": values, "$unset": {"pending_winners": ""}}, ) return result.modified_count == 1 async def claim_giveaway_cancel( giveaway_id: str, chat_id: int ) -> dict[str, Any] | None: await ensure_giveaway_indexes() normalized_id = normalize_giveaway_id(giveaway_id) claimed = await giveawaysdb.find_one_and_update( { "giveaway_id": normalized_id, "chat_id": int(chat_id), "status": STATUS_RUNNING, }, { "$set": { "status": STATUS_CANCELING, "cancel_started_at": utc_now(), "updated_at": utc_now(), } }, return_document=ReturnDocument.AFTER, ) if claimed: return claimed return await giveawaysdb.find_one( { "giveaway_id": normalized_id, "chat_id": int(chat_id), "status": STATUS_CANCELING, } ) async def cancel_giveaway(giveaway_id: str, chat_id: int) -> bool: claimed = await claim_giveaway_cancel(giveaway_id, chat_id) if not claimed: return False return await finalize_giveaway_cancel(giveaway_id) async def finalize_giveaway_cancel(giveaway_id: str) -> bool: await ensure_giveaway_indexes() result = await giveawaysdb.update_one( { "giveaway_id": normalize_giveaway_id(giveaway_id), "status": STATUS_CANCELING, }, { "$set": { "status": STATUS_CANCELED, "canceled_at": utc_now(), "updated_at": utc_now(), } }, ) return result.modified_count == 1 async def record_reroll( giveaway_id: str, winners: list[dict[str, Any]], tier_name: str | None, moderator_id: int, reroll_id: str | None = None, ) -> str: await ensure_giveaway_indexes() reroll_id = reroll_id or token_hex(8) await giveawaysdb.update_one( { "giveaway_id": normalize_giveaway_id(giveaway_id), "status": STATUS_FINISHED, "rerolls.reroll_id": {"$ne": reroll_id}, }, { "$push": { "rerolls": { "reroll_id": reroll_id, "tier_name": tier_name, "moderator_id": int(moderator_id), "winners": winners, "created_at": utc_now(), } }, "$set": {"updated_at": utc_now()}, }, ) return reroll_id async def is_giveaway_banned(chat_id: int, user_id: int) -> bool: await ensure_giveaway_indexes() return ( await giveaway_bansdb.find_one( {"chat_id": int(chat_id), "user_id": int(user_id)} ) ) is not None async def add_giveaway_ban( *, chat_id: int, user_id: int, moderator_id: int, reason: str | None ) -> None: await ensure_giveaway_indexes() now = utc_now() await giveaway_bansdb.update_one( {"chat_id": int(chat_id), "user_id": int(user_id)}, { "$set": { "moderator_id": int(moderator_id), "reason": (reason or "").strip()[:500], "updated_at": now, }, "$setOnInsert": {"created_at": now}, }, upsert=True, ) async def remove_giveaway_ban(chat_id: int, user_id: int) -> bool: await ensure_giveaway_indexes() result = await giveaway_bansdb.delete_one( {"chat_id": int(chat_id), "user_id": int(user_id)} ) return result.deleted_count == 1 async def list_giveaway_bans( *, chat_id: int, page: int = 1, page_size: int = 20 ) -> tuple[list[dict[str, Any]], int]: await ensure_giveaway_indexes() filters = {"chat_id": int(chat_id)} page = max(1, int(page)) page_size = max(1, min(int(page_size), 100)) total = await giveaway_bansdb.count_documents(filters) cursor = ( giveaway_bansdb.find(filters) .sort("created_at", DESCENDING) .skip((page - 1) * page_size) .limit(page_size) ) return [doc async for doc in cursor], total async def add_participant( *, giveaway_id: str, user_id: int, username: str | None, first_name: str | None, display_name: str | None = None, ) -> str: await ensure_giveaway_indexes() giveaway = await get_giveaway(giveaway_id) if not giveaway: return "missing" if giveaway["status"] != STATUS_RUNNING: return "closed" starts_at = giveaway.get("starts_at") or giveaway.get("created_at") if starts_at and as_utc(starts_at) > utc_now(): return "not_started" if as_utc(giveaway["ends_at"]) <= utc_now(): return "ended" if await is_giveaway_banned(giveaway["chat_id"], user_id): return "banned" try: await participantsdb.insert_one( { "giveaway_id": giveaway["giveaway_id"], "chat_id": int(giveaway["chat_id"]), "user_id": int(user_id), "username": username, "first_name": first_name, "display_name": display_name or first_name, "entry_cost": int(giveaway.get("entry_cost", 0)), "active": True, "joined_at": utc_now(), } ) except DuplicateKeyError: return "duplicate" return "ok" async def get_participant(giveaway_id: str, user_id: int) -> dict[str, Any] | None: await ensure_giveaway_indexes() return await participantsdb.find_one( { "giveaway_id": normalize_giveaway_id(giveaway_id), "user_id": int(user_id), } ) async def remove_participant( *, giveaway_id: str, user_id: int, moderator_id: int, reason: str, refund: bool, ) -> dict[str, Any] | None: await ensure_giveaway_indexes() return await participantsdb.find_one_and_update( { "giveaway_id": normalize_giveaway_id(giveaway_id), "user_id": int(user_id), "active": {"$ne": False}, }, { "$set": { "active": False, "removed_at": utc_now(), "removed_by": int(moderator_id), "removal_reason": reason.strip()[:500], "refund_requested": bool(refund), } }, return_document=ReturnDocument.BEFORE, ) async def mark_participant_refunded( giveaway_id: str, user_id: int, refund_reason: str ) -> None: await ensure_giveaway_indexes() await participantsdb.update_one( { "giveaway_id": normalize_giveaway_id(giveaway_id), "user_id": int(user_id), }, { "$set": { "refunded_at": utc_now(), "refund_reason": refund_reason[:100], } }, ) async def count_participants(giveaway_id: str, *, active_only: bool = True) -> int: await ensure_giveaway_indexes() filters: dict[str, Any] = { "giveaway_id": normalize_giveaway_id(giveaway_id) } if active_only: filters["active"] = {"$ne": False} return await participantsdb.count_documents(filters) async def list_participants( giveaway_id: str, limit: int | None = None, *, active_only: bool = True ) -> list[dict[str, Any]]: await ensure_giveaway_indexes() filters: dict[str, Any] = { "giveaway_id": normalize_giveaway_id(giveaway_id) } if active_only: filters["active"] = {"$ne": False} cursor = participantsdb.find(filters).sort("joined_at", ASCENDING) if limit is not None: cursor = cursor.limit(max(1, int(limit))) return [doc async for doc in cursor] async def list_participants_page( *, giveaway_id: str, active_only: bool = False, page: int = 1, page_size: int = 20, ) -> tuple[list[dict[str, Any]], int]: await ensure_giveaway_indexes() filters: dict[str, Any] = { "giveaway_id": normalize_giveaway_id(giveaway_id) } if active_only: filters["active"] = {"$ne": False} page = max(1, int(page)) page_size = max(1, min(int(page_size), 100)) total = await participantsdb.count_documents(filters) cursor = ( participantsdb.find(filters) .sort("joined_at", ASCENDING) .skip((page - 1) * page_size) .limit(page_size) ) return [doc async for doc in cursor], total