TelegramBotAdmin/wbb/modules/giveaway.py

from __future__ import annotations

import asyncio
import re
from datetime import timedelta
from html import escape
from typing import Any

from pyrogram import filters
from pyrogram.enums import ParseMode
from pyrogram.types import CallbackQuery, Message

from wbb import app, log
from wbb.core.decorators.permissions import adminsOnly
from wbb.services.giveaways import (
    GiveawayServiceError,
    cancel_and_refund_giveaway,
    create_and_publish_giveaway,
    finish_and_publish_giveaway,
    join_giveaway,
    remove_and_optionally_refund_participant,
    render_winners,
    reroll_giveaway,
    resume_pending_giveaway,
)
from wbb.services.member_identity import display_name
from wbb.utils.dbgiveaway import (
    add_giveaway_ban,
    count_participants,
    get_giveaway,
    list_due_giveaways,
    list_participants,
    list_pending_cancellations,
    list_running_giveaways,
    normalize_giveaway_id,
    remove_giveaway_ban,
    utc_now,
)
from wbb.utils.functions import extract_user_and_reason

__MODULE__ = "抽奖管理"
__HELP__ = """/giveaway <时长> | <奖项> | <标题> | [说明] | [最低积分] | [报名消耗] | [参与奖励]
在当前群创建抽奖。

奖项格式：`1`、`一等奖:1, 二等奖:3`，或 `一等奖:1:50`；
最后的 50 表示该奖项每位中奖者获得 50 积分。
时长单位：s 秒、m 分钟、h 小时、d 天、w 周。

/gjoin <编号> - 报名参加抽奖。
/glist - 查看进行中的抽奖。
/gparticipants <编号> - 查看参与者。
/gend <编号> - 立即开奖。
/gcancel <编号> - 取消抽奖并退还报名积分。
/greroll <编号> [奖项] - 排除历次中奖者后重抽。
/gremove <编号> <用户 ID> [no-refund] - 移除参与者，可选择不退款。
/gban <成员> [原因] - 禁止成员参加本群抽奖。
/gunban <成员> - 解除抽奖参与限制。"""

MAX_TOTAL_WINNERS = 100
PARTICIPANT_PREVIEW_LIMIT = 30
SWEEP_INTERVAL_SECONDS = 30
_DURATION_RE = re.compile(r"^\s*(\d+)\s*([smhdwSMHDW])\s*$")
_DURATION_UNITS = {
    "s": "seconds",
    "m": "minutes",
    "h": "hours",
    "d": "days",
    "w": "weeks",
}
_sweeper_started = False


def _parse_duration(raw: str):
    match = _DURATION_RE.match(raw)
    if not match:
        return None
    amount = int(match.group(1))
    if amount <= 0:
        return None
    return utc_now() + timedelta(**{_DURATION_UNITS[match.group(2).lower()]: amount})


def _parse_prizes(raw: str) -> list[dict[str, Any]] | None:
    raw = raw.strip()
    if raw.isdigit():
        count = int(raw)
        if 0 < count <= MAX_TOTAL_WINNERS:
            return [{"name": "中奖者", "count": count, "points_reward": 0}]
        return None

    prizes: list[dict[str, Any]] = []
    for piece in re.split(r"[,;\n]+", raw):
        parts = [part.strip() for part in re.split(r"[:：]", piece.strip())]
        if len(parts) not in {2, 3} or not parts[0] or not parts[1].isdigit():
            return None
        count = int(parts[1])
        points_reward = int(parts[2]) if len(parts) == 3 and parts[2].isdigit() else 0
        if count <= 0 or points_reward < 0:
            return None
        prizes.append(
            {"name": parts[0], "count": count, "points_reward": points_reward}
        )
    if not prizes or sum(item["count"] for item in prizes) > MAX_TOTAL_WINNERS:
        return None
    return prizes


def _parse_optional_points(pieces: list[str], index: int) -> int | None:
    if len(pieces) <= index or not pieces[index]:
        return 0
    if not pieces[index].isdigit():
        return None
    return int(pieces[index])


def _create_usage() -> str:
    return (
        "<b>用法</b>\n"
        "/giveaway 1h | 1 | 月度会员抽奖\n"
        "/giveaway 2d | 一等奖:1:50, 二等奖:3:10 | 活动抽奖 | "
        "抽奖说明 | 20 | 5 | 2"
    )


def _actor_id(message: Message) -> int:
    if message.from_user:
        return message.from_user.id
    if message.sender_chat:
        return message.sender_chat.id
    return 0


def _actor_name(message: Message) -> str:
    if message.from_user:
        return message.from_user.first_name or str(message.from_user.id)
    if message.sender_chat:
        return message.sender_chat.title or str(message.sender_chat.id)
    return "匿名管理员"


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'<a href="tg://user?id={user_id}">{escape(str(label))}</a>'


async def _join_user(giveaway_id: str, user) -> str:
    status, _ = await join_giveaway(
        giveaway_id,
        user_id=user.id,
        username=user.username,
        first_name=user.first_name,
        display_name=display_name(user.first_name, user.last_name),
    )
    return {
        "ok": "报名成功。",
        "duplicate": "你已经报名过该抽奖。",
        "missing": "未找到该抽奖。",
        "closed": "该抽奖已经关闭。",
        "ended": "该抽奖已经结束。",
        "banned": "你已被禁止参加本群抽奖。",
        "not_member": "只有当前群成员可以参加该抽奖。",
        "insufficient_points": "你的积分不足，无法报名。",
    }.get(status, "报名失败，请稍后重试。")


async def _sweep_due_giveaways() -> None:
    await asyncio.sleep(5)
    while True:
        try:
            pending = await list_due_giveaways()
            pending.extend(await list_pending_cancellations())
            for giveaway in pending:
                await resume_pending_giveaway(giveaway)
        except Exception as exc:
            log.error(f"抽奖定时任务失败：{exc}")
        await asyncio.sleep(SWEEP_INTERVAL_SECONDS)


def _start_sweeper() -> None:
    global _sweeper_started
    if _sweeper_started:
        return
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        return
    _sweeper_started = True
    loop.create_task(_sweep_due_giveaways(), name="giveaway-sweeper")


@app.on_message(filters.command(["giveaway", "gcreate"]) & ~filters.private)
@adminsOnly("can_change_info")
async def create_giveaway_command(_, message: Message):
    payload = message.text.split(None, 1)[1].strip() if message.text and " " in message.text else ""
    pieces = [piece.strip() for piece in payload.split("|")]
    if len(pieces) < 3:
        return await message.reply_text(_create_usage(), parse_mode=ParseMode.HTML)
    ends_at = _parse_duration(pieces[0])
    prizes = _parse_prizes(pieces[1])
    minimum_points = _parse_optional_points(pieces, 4)
    entry_cost = _parse_optional_points(pieces, 5)
    participation_reward = _parse_optional_points(pieces, 6)
    if (
        not ends_at
        or not prizes
        or not pieces[2]
        or minimum_points is None
        or entry_cost is None
        or participation_reward is None
    ):
        return await message.reply_text(_create_usage(), parse_mode=ParseMode.HTML)
    giveaway = await create_and_publish_giveaway(
        chat_id=message.chat.id,
        creator_id=_actor_id(message),
        creator_name=_actor_name(message),
        title=pieces[2],
        description=pieces[3] if len(pieces) > 3 else "",
        prizes=prizes,
        ends_at=ends_at,
        minimum_points=minimum_points,
        entry_cost=entry_cost,
        participation_reward=participation_reward,
    )
    return await message.reply_text(f"抽奖 #{giveaway['giveaway_id']} 已创建。")


@app.on_callback_query(filters.regex(r"^giveaway_join:"))
async def giveaway_join_callback(_, query: CallbackQuery):
    giveaway_id = normalize_giveaway_id(query.data.split(":", 1)[1])
    await query.answer(await _join_user(giveaway_id, query.from_user), show_alert=True)


@app.on_message(filters.command("gjoin") & ~filters.private)
async def giveaway_join_command(_, message: Message):
    if not message.from_user:
        return await message.reply_text("只有普通用户身份可以参加抽奖。")
    if len(message.command) < 2:
        return await message.reply_text("用法：/gjoin <抽奖编号>")
    return await message.reply_text(await _join_user(message.command[1], message.from_user))


@app.on_message(filters.command("glist") & ~filters.private)
@adminsOnly("can_change_info")
async def list_giveaways_command(_, message: Message):
    giveaways = await list_running_giveaways(chat_id=message.chat.id, limit=20)
    if not giveaways:
        return await message.reply_text("本群当前没有进行中的抽奖。")
    lines = ["<b>进行中的抽奖</b>"]
    for giveaway in giveaways:
        lines.append(
            "#{id} - {title} - {count} 人参与 - 截止于 {ends_at}".format(
                id=escape(giveaway["giveaway_id"]),
                title=escape(giveaway["title"]),
                count=await count_participants(giveaway["giveaway_id"]),
                ends_at=giveaway["ends_at"].strftime("%Y-%m-%d %H:%M UTC"),
            )
        )
    return await message.reply_text("\n".join(lines), parse_mode=ParseMode.HTML)


@app.on_message(filters.command("gparticipants") & ~filters.private)
@adminsOnly("can_change_info")
async def participants_command(_, message: Message):
    if len(message.command) < 2:
        return await message.reply_text("用法：/gparticipants <抽奖编号>")
    giveaway_id = normalize_giveaway_id(message.command[1])
    giveaway = await get_giveaway(giveaway_id)
    if not giveaway or giveaway["chat_id"] != message.chat.id:
        return await message.reply_text("本群中未找到该抽奖。")
    total = await count_participants(giveaway_id)
    participants = await list_participants(giveaway_id, limit=PARTICIPANT_PREVIEW_LIMIT)
    lines = [f"<b>抽奖 #{escape(giveaway_id)} 的参与者</b>", f"共 {total} 人"]
    lines.extend("- " + _participant_mention(item) for item in participants)
    if total > len(participants):
        lines.append(f"另有 {total - len(participants)} 人未显示。")
    return await message.reply_text(
        "\n".join(lines), parse_mode=ParseMode.HTML, disable_web_page_preview=True
    )


@app.on_message(filters.command("gend") & ~filters.private)
@adminsOnly("can_change_info")
async def end_giveaway_command(_, message: Message):
    if len(message.command) < 2:
        return await message.reply_text("用法：/gend <抽奖编号>")
    giveaway_id = normalize_giveaway_id(message.command[1])
    giveaway = await get_giveaway(giveaway_id)
    if not giveaway or giveaway["chat_id"] != message.chat.id:
        return await message.reply_text("本群中未找到该抽奖。")
    _, text, _ = await finish_and_publish_giveaway(giveaway_id)
    return await message.reply_text(text)


@app.on_message(filters.command("gcancel") & ~filters.private)
@adminsOnly("can_change_info")
async def cancel_giveaway_command(_, message: Message):
    if len(message.command) < 2:
        return await message.reply_text("用法：/gcancel <抽奖编号>")
    giveaway_id = normalize_giveaway_id(message.command[1])
    giveaway = await get_giveaway(giveaway_id)
    if not giveaway or giveaway["chat_id"] != message.chat.id:
        return await message.reply_text("本群中未找到该抽奖。")
    _, text, _ = await cancel_and_refund_giveaway(
        giveaway_id, chat_id=message.chat.id
    )
    return await message.reply_text(text)


@app.on_message(filters.command("greroll") & ~filters.private)
@adminsOnly("can_change_info")
async def reroll_giveaway_command(_, message: Message):
    parts = message.text.split(None, 2) if message.text else []
    if len(parts) < 2:
        return await message.reply_text("用法：/greroll <抽奖编号> [奖项名称]")
    giveaway_id = normalize_giveaway_id(parts[1])
    giveaway = await get_giveaway(giveaway_id)
    if not giveaway or giveaway["chat_id"] != message.chat.id:
        return await message.reply_text("本群中未找到该抽奖。")
    try:
        winners, _ = await reroll_giveaway(
            giveaway_id=giveaway_id,
            moderator_id=_actor_id(message),
            tier_name=parts[2].strip() if len(parts) > 2 else None,
            publish=False,
        )
    except GiveawayServiceError as exc:
        return await message.reply_text(str(exc))
    return await message.reply_text(
        render_winners(
            giveaway,
            winners,
            reroll=True,
            tier_name=parts[2].strip() if len(parts) > 2 else None,
        ),
        parse_mode=ParseMode.HTML,
        disable_web_page_preview=True,
    )


@app.on_message(filters.command("gremove") & ~filters.private)
@adminsOnly("can_change_info")
async def remove_participant_command(_, message: Message):
    if len(message.command) < 3 or not message.command[2].lstrip("-").isdigit():
        return await message.reply_text(
            "用法：/gremove <抽奖编号> <用户 ID> [no-refund]"
        )
    giveaway_id = normalize_giveaway_id(message.command[1])
    giveaway = await get_giveaway(giveaway_id)
    if not giveaway or giveaway["chat_id"] != message.chat.id:
        return await message.reply_text("本群中未找到该抽奖。")
    refund = not (len(message.command) > 3 and message.command[3].lower() == "no-refund")
    removed = await remove_and_optionally_refund_participant(
        giveaway_id=giveaway_id,
        user_id=int(message.command[2]),
        moderator_id=_actor_id(message),
        reason="由群管理员移除",
        refund=refund,
    )
    return await message.reply_text(
        "参与者已移除。"
        if removed
        else "未找到该参与者，或该参与者此前已被移除。"
    )


@app.on_message(filters.command("gban") & ~filters.private)
@adminsOnly("can_change_info")
async def giveaway_ban_command(_, message: Message):
    user_id, reason = await extract_user_and_reason(message)
    if not user_id:
        return await message.reply_text("用法：/gban <成员> [原因]")
    await add_giveaway_ban(
        chat_id=message.chat.id,
        user_id=int(user_id),
        moderator_id=_actor_id(message),
        reason=reason,
    )
    return await message.reply_text(f"已禁止用户 {user_id} 参加本群抽奖。")


@app.on_message(filters.command("gunban") & ~filters.private)
@adminsOnly("can_change_info")
async def giveaway_unban_command(_, message: Message):
    user_id, _ = await extract_user_and_reason(message)
    if not user_id:
        return await message.reply_text("用法：/gunban <成员>")
    removed = await remove_giveaway_ban(message.chat.id, int(user_id))
    return await message.reply_text(
        f"已允许用户 {user_id} 参加本群抽奖。"
        if removed
        else "该用户没有被限制参加抽奖。"
    )


_start_sweeper()