TelegramBotAdmin/wbb/services/giveaways.py

from __future__ import annotations

import asyncio
import secrets
from collections import defaultdict
from contextlib import suppress
from html import escape
from typing import Any

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,
)
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)


class GiveawayServiceError(RuntimeError):
    def __init__(self, code: str, message: str, *, status: int = 400):
        super().__init__(message)
        self.code = code
        self.status = status


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>'


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"])
    text = [
        f"<b>抽奖 #{escape(giveaway['giveaway_id'])}</b>",
        f"<b>状态：</b>{escape(giveaway_status_label(giveaway['status']))}",
        f"<b>标题：</b>{escape(giveaway['title'])}",
    ]
    if giveaway.get("description"):
        text.append(f"<b>说明：</b>{escape(giveaway['description'])}")
    text.extend(
        [
            "<b>奖项：</b>",
            _format_prizes(giveaway["prizes"]),
            f"<b>截止时间：</b>{as_utc(giveaway['ends_at']).strftime('%Y-%m-%d %H:%M UTC')}",
            f"<b>参与人数：</b>{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"<b>最低积分：</b>{minimum}")
    if cost:
        text.append(f"<b>报名消耗：</b>{cost}")
    if reward:
        text.append(f"<b>参与奖励：</b>+{reward}")
    return "\n".join(text)


def join_markup(giveaway_id: str) -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup(
        [
            [
                InlineKeyboardButton(
                    "立即报名", callback_data=f"giveaway_join:{giveaway_id}"
                )
            ]
        ]
    )


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"<b>{title} #{escape(giveaway['giveaway_id'])}</b>",
        f"<b>标题：</b>{escape(giveaway['title'])}",
    ]
    if tier_name:
        lines.append(f"<b>奖项：</b>{escape(tier_name)}")
    if not winners:
        lines.append("没有符合条件的参与者。")
        return "\n".join(lines)
    lines.append("<b>中奖者：</b>")
    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
        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,
            )
    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"恭喜你在 <b>{escape(giveaway['title'])}</b> 中获得"
                    f"<b>{escape(str(winner['tier']))}</b>。"
                ),
                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",
        )
    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}"
        )