TelegramBotAdmin/wbb/services/message_cleanup.py

from __future__ import annotations

import asyncio
import logging
from functools import wraps
from typing import Any
from weakref import WeakKeyDictionary

GROUP_INTERACTION_TTL_SECONDS = 60

_LOGGER = logging.getLogger(__name__)
_scheduled_messages: set[tuple[int, int]] = set()
_cleanup_tasks: set[asyncio.Task[Any]] = set()
_active_interactions: WeakKeyDictionary[asyncio.Task[Any], int] = WeakKeyDictionary()
_TRACKED_SEND_METHODS = (
    "send_message",
    "send_photo",
    "send_document",
    "send_video",
    "send_animation",
    "send_audio",
    "send_voice",
    "send_video_note",
    "send_sticker",
    "send_contact",
    "send_location",
    "send_venue",
    "send_poll",
    "send_dice",
    "send_media_group",
    "send_cached_media",
    "copy_message",
    "copy_media_group",
    "forward_messages",
)


def _message_key(message: Any) -> tuple[int, int] | None:
    chat = getattr(message, "chat", None)
    message_id = getattr(message, "id", None)
    if not chat or not message_id:
        return None
    try:
        return int(chat.id), int(message_id)
    except (TypeError, ValueError):
        return None


def _is_group_message(message: Any) -> bool:
    chat = getattr(message, "chat", None)
    if not chat:
        return False
    chat_type = getattr(chat, "type", None)
    chat_type_value = str(getattr(chat_type, "value", chat_type) or "").lower()
    if chat_type_value:
        return chat_type_value in {"group", "supergroup"}
    try:
        return int(chat.id) < 0
    except (TypeError, ValueError):
        return False


def _current_task() -> asyncio.Task[Any] | None:
    try:
        return asyncio.current_task()
    except RuntimeError:
        return None


async def _delete_later(message: Any, delay_seconds: int) -> None:
    key = _message_key(message)
    if key is None:
        return
    try:
        await asyncio.sleep(max(1, int(delay_seconds)))
        try:
            await message.delete(revoke=True)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            _LOGGER.warning(
                "群交互消息自动删除失败 chat_id=%s message_id=%s: %s",
                key[0],
                key[1],
                exc,
            )
    finally:
        _scheduled_messages.discard(key)


def schedule_group_cleanup(
    *messages: Any,
    delay_seconds: int = GROUP_INTERACTION_TTL_SECONDS,
) -> None:
    for message in messages:
        key = _message_key(message)
        if key is None or not _is_group_message(message):
            continue
        if key in _scheduled_messages:
            continue
        _scheduled_messages.add(key)
        task = asyncio.create_task(
            _delete_later(message, delay_seconds),
            name=f"group-interaction-cleanup:{key[0]}:{key[1]}",
        )
        _cleanup_tasks.add(task)
        task.add_done_callback(_cleanup_tasks.discard)


def is_group_cleanup_scheduled(message: Any) -> bool:
    key = _message_key(message)
    return bool(key and key in _scheduled_messages)


def clear_current_group_interaction() -> None:
    task = _current_task()
    if task is not None:
        _active_interactions.pop(task, None)


def begin_group_message_interaction(
    message: Any,
    *,
    menu_texts: set[str] | frozenset[str] = frozenset(),
) -> bool:
    clear_current_group_interaction()
    if not _is_group_message(message):
        return False
    text = str(
        getattr(message, "text", None) or getattr(message, "caption", None) or ""
    ).strip()
    folded = text.casefold()
    is_command = text.startswith("/") or folded in {"@admin", "@admins"}
    is_command = is_command or folded.startswith(("@admin ", "@admins "))
    if not is_command and text not in menu_texts:
        return False
    task = _current_task()
    if task is None:
        return False
    _active_interactions[task] = int(message.chat.id)
    schedule_group_cleanup(message)
    return True


def begin_group_callback_interaction(query: Any) -> bool:
    clear_current_group_interaction()
    message = getattr(query, "message", None)
    if (
        not message
        or not _is_group_message(message)
        or not is_group_cleanup_scheduled(message)
    ):
        return False
    task = _current_task()
    if task is None:
        return False
    _active_interactions[task] = int(message.chat.id)
    return True


def track_current_group_interaction(
    result: Any,
    *,
    reply_to_message_id: Any = None,
) -> None:
    if reply_to_message_id is None:
        return
    task = _current_task()
    if task is None:
        return
    chat_id = _active_interactions.get(task)
    if chat_id is None:
        return
    messages = result if isinstance(result, (list, tuple)) else (result,)
    for message in messages:
        key = _message_key(message)
        if key and key[0] == chat_id and _is_group_message(message):
            schedule_group_cleanup(message)


def install_group_interaction_tracking(client: Any) -> None:
    if getattr(client, "_wbb_group_interaction_tracking", False):
        return
    for method_name in _TRACKED_SEND_METHODS:
        original = getattr(client, method_name, None)
        if not callable(original):
            continue

        @wraps(original)
        async def tracked_send(*args: Any, _original=original, **kwargs: Any):
            result = await _original(*args, **kwargs)
            track_current_group_interaction(
                result,
                reply_to_message_id=kwargs.get("reply_to_message_id"),
            )
            return result

        setattr(client, method_name, tracked_send)
    client._wbb_group_interaction_tracking = True