beavergram/backend/src/userbot/modules/media/downloader.py

from typing import Any

from pyrogram import Client
from pyrogram.errors import (
    FileIdInvalid,
    FileReferenceExpired,
    FileReferenceInvalid,
    FloodPremiumWait,
    FloodWait,
)
from pyrogram.types import Message

from userbot.modules.capture import repository
from userbot.modules.capture.context import CaptureContext
from userbot.modules.download import Downloadable, download_bytes
from utils.logging import logger
from utils.policy.models import CaptureToggles

STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)

_MEDIA_ATTRS = (
    "photo",
    "video",
    "voice",
    "video_note",
    "document",
    "audio",
    "animation",
    "sticker",
)


_WEB_PAGE_ATTRS = ("photo", "video", "animation", "document", "audio")


def media_object(message: Message) -> tuple[str | None, Any]:
    for attr in _MEDIA_ATTRS:
        obj = getattr(message, attr, None)
        if obj is not None:
            return attr, obj
    web_page = getattr(message, "web_page", None)
    if web_page is not None:
        for attr in _WEB_PAGE_ATTRS:
            obj = getattr(web_page, attr, None)
            if obj is not None:
                return attr, obj
    return None, None


def self_destruct_ttl(message: Message) -> int | None:
    _, obj = media_object(message)
    return getattr(obj, "ttl_seconds", None) if obj is not None else None


def media_unique_id(message: Message) -> str | None:
    _, obj = media_object(message)
    return getattr(obj, "file_unique_id", None) if obj is not None else None


def _download_target(message: Message) -> Downloadable | None:
    kind, obj = media_object(message)
    if obj is None:
        return None
    return message if getattr(message, kind or "", None) is obj else obj


async def _fresh_target(
    client: Client, chat_id: int, message_id: int
) -> Downloadable | None:
    message = await client.get_messages(chat_id, message_id)
    if not isinstance(message, Message) or message.empty:
        return None
    return _download_target(message)


async def _download(
    client: Client, target: Downloadable | None, chat_id: int, message_id: int
) -> bytes | None:
    if target is None:
        return None
    try:
        return await download_bytes(client, target)
    except TimeoutError:
        logger.warning(
            f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]"
        )
        return None
    except STALE_FILE_ID:
        fresh = await _fresh_target(client, chat_id, message_id)
        if fresh is None:
            return None
        try:
            return await download_bytes(client, fresh)
        except (TimeoutError, *STALE_FILE_ID):
            logger.warning(
                f"[yellow]Media download failed after refetch "
                f"for {chat_id}/{message_id}.[/]"
            )
            return None


async def capture_media(  # noqa: PLR0913
    client: Client,
    message: Message,
    ctx: CaptureContext,
    chat_id: int,
    message_id: int,
    toggles: CaptureToggles,
    *,
    force: bool = False,
) -> bool:
    kind, obj = media_object(message)
    if obj is None:
        return True
    unique_id = getattr(obj, "file_unique_id", None)
    ttl = getattr(obj, "ttl_seconds", None)
    want = toggles.self_destruct_media if ttl else toggles.media
    file_size = getattr(obj, "file_size", None)
    mime = getattr(obj, "mime_type", None)
    storage_key: str | None = None
    downloaded = False
    evicted = False
    flood: FloodWait | FloodPremiumWait | None = None
    if want:
        existing = await repository.current_media(
            ctx.pool, ctx.account_id, chat_id, message_id
        )
        if (
            existing is not None
            and existing["downloaded"]
            and existing["unique_id"] == unique_id
            and existing["storage_key"] is not None
        ):
            storage_key = existing["storage_key"]
            file_size = existing["file_size"]
            downloaded = True
        elif force or not await repository.is_evicted(ctx.pool, unique_id):
            target = _download_target(message)
            try:
                data = await _download(client, target, chat_id, message_id)
            except (FloodWait, FloodPremiumWait) as exc:
                flood = exc
                data = None
            if data is not None:
                storage_key = ctx.storage.put(data)
                file_size = len(data)
                downloaded = True
        else:
            evicted = True
    await repository.insert_media(
        ctx.pool,
        ctx.account_id,
        chat_id,
        message_id,
        kind or "unknown",
        storage_key,
        file_size,
        mime,
        ttl,
        unique_id,
        downloaded=downloaded,
        evicted=evicted,
    )
    if (
        force
        and storage_key is not None
        and await repository.unevict(ctx.pool, unique_id, storage_key, file_size)
    ):
        ctx.storage.purge(storage_key)
    if flood is not None:
        raise flood
    return downloaded or not want