TelegramBotAdmin/wbb/utils/dbgiveaway.py

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,
    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("必须填写抽奖标题。")
    normalized_ends_at = as_utc(ends_at)
    if normalized_ends_at <= utc_now():
        raise ValueError("抽奖截止时间必须晚于当前时间。")
    normalized_prizes = normalize_prizes(prizes)
    now = utc_now()
    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,
            "ends_at": normalized_ends_at,
            "message_id": None,
            "winners": [],
            "rerolls": [],
        }
        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_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"
    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