beavergram/backend/src/watcher/registry.py

import asyncio
import contextlib
from dataclasses import dataclass

import asyncpg
from pyrogram import Client, enums
from pyrogram.errors import RPCError
from pyrogram.handlers import (
    ChatMemberUpdatedHandler,
    EditedMessageHandler,
    MessageHandler,
)

from userbot import PyroClient
from userbot.modules.capture import capture_message
from userbot.modules.capture.chat_meta import meta_from_chat
from userbot.modules.jobs import JobConsumer
from utils.events import notify_bg_event
from utils.jobs import enqueue
from utils.logging import logger
from utils.read import accounts, bots
from utils.storage import ContentAddressedStorage
from watcher import tokens
from watcher.context import build_bot_capture_context
from watcher.drain import drain_backlog

DEVICE_MODEL = "Beavergram Bot"
TOKEN_MISSING = "Токен не найден на диске"  # noqa: S105
INVALID_TOKEN_ERRORS = ("ACCESS_TOKEN_INVALID", "ACCESS_TOKEN_EXPIRED")


@dataclass
class RunningBot:
    client: PyroClient
    account_id: int
    tasks: list[asyncio.Task]


async def _cancel(task: asyncio.Task) -> None:
    task.cancel()
    with contextlib.suppress(asyncio.CancelledError):
        await task


def _handlers(pool: asyncpg.Pool) -> tuple:
    async def on_message(client: Client, message) -> None:  # noqa: ANN001
        ctx = getattr(client, "capture", None)
        if ctx is None or message.empty or message.chat is None or message.date is None:
            return
        meta = meta_from_chat(message.chat, ctx.contacts.ids)
        await ctx.watches.on_text(
            meta.chat_id, message.id, message.text or message.caption
        )
        toggles = ctx.resolve(meta)
        if not toggles.messages:
            return
        await capture_message(client, message, ctx, toggles)
        await notify_bg_event(
            ctx.pool,
            "message",
            ctx.account_id,
            chat_id=meta.chat_id,
            message_id=message.id,
        )
        if message.chat.type is not enums.ChatType.PRIVATE:
            await enqueue(
                pool, ctx.account_id, "enrich_chat", {"chat_id": meta.chat_id}
            )

    async def on_member(client: Client, update) -> None:  # noqa: ANN001
        ctx = getattr(client, "capture", None)
        if ctx is None or update.chat is None:
            return
        await enqueue(pool, ctx.account_id, "enrich_chat", {"chat_id": update.chat.id})

    return (
        MessageHandler(on_message),
        EditedMessageHandler(on_message),
        ChatMemberUpdatedHandler(on_member),
    )


def _is_invalid_token(exc: RPCError) -> bool:
    return any(marker in str(exc) for marker in INVALID_TOKEN_ERRORS)


class BotRegistry:
    def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
        self._pool = pool
        self._storage = storage
        self._running: dict[int, RunningBot] = {}
        self._lock = asyncio.Lock()

    @property
    def running(self) -> list[int]:
        return list(self._running)

    async def sync(self) -> None:
        async with self._lock:
            await self._adopt_dropped_in()
            enabled = await bots.enabled_bot_ids(self._pool)
            for bot_id in set(self._running) - enabled:
                await self._stop(bot_id)
            for bot_id in enabled - set(self._running):
                await self._start(bot_id)
            for bot_id in await bots.removed_bot_ids(self._pool):
                tokens.drop(bot_id)

    async def close(self) -> None:
        for bot_id in list(self._running):
            await self._stop(bot_id)

    async def _adopt_dropped_in(self) -> None:
        unknown = tokens.stored_bot_ids() - await bots.known_bot_ids(self._pool)
        for bot_id in unknown:
            await bots.register_bot(self._pool, bot_id)
            logger.info(f"[green]Bot token adopted from disk:[/] {bot_id}")

    def _client(self, bot_id: int, token: str) -> PyroClient:
        return PyroClient(
            str(bot_id),
            workdir=str(tokens.sessions_dir()),
            device_model=DEVICE_MODEL,
            load_handlers=False,
            bot_token=token,
            skip_updates=False,
        )

    async def _start(self, bot_id: int) -> None:
        token = tokens.load(bot_id)
        if token is None:
            await bots.set_status(
                self._pool, bot_id, bots.STATUS_UNAUTHORIZED, detail=TOKEN_MISSING
            )
            return
        client = self._client(bot_id, token)
        for handler in _handlers(self._pool):
            client.add_handler(handler)
        try:
            await client.start()
            me = client.me
            if me is None:
                msg = "bot session is not authorized"
                raise RuntimeError(msg)
        except RPCError as exc:
            status = (
                bots.STATUS_UNAUTHORIZED
                if _is_invalid_token(exc)
                else bots.STATUS_ERROR
            )
            await bots.set_status(self._pool, bot_id, status, detail=str(exc))
            logger.warning(f"[yellow]Bot {bot_id} failed to start:[/] {exc}")
            with contextlib.suppress(Exception):
                await client.stop()
            return
        except Exception as exc:
            await bots.set_status(
                self._pool, bot_id, bots.STATUS_ERROR, detail=str(exc)
            )
            logger.exception(f"[red]Bot {bot_id} failed to start[/]")
            with contextlib.suppress(Exception):
                await client.stop()
            return

        account_id = await accounts.sync_account(
            self._pool, me, str(bot_id), is_bot=True
        )
        client.capture = await build_bot_capture_context(
            client, self._pool, self._storage, account_id
        )
        await bots.link_account(self._pool, bot_id, account_id)
        await bots.set_status(
            self._pool, bot_id, bots.STATUS_OK, username=me.username, title=me.full_name
        )
        await accounts.notify_accounts_changed(self._pool)

        consumer = JobConsumer(client, self._pool, account_id)
        tasks = [asyncio.create_task(consumer.run())]
        if await bots.needs_backfill(self._pool, bot_id):
            tasks.append(asyncio.create_task(self._backlog(client, bot_id, account_id)))
        self._running[bot_id] = RunningBot(client, account_id, tasks)
        logger.info(
            f"[green]Bot account {account_id} watched:[/] @{me.username} ({bot_id})"
        )

    async def _backlog(self, client: PyroClient, bot_id: int, account_id: int) -> None:
        stored = await drain_backlog(client, account_id)
        await bots.mark_backfilled(self._pool, bot_id)
        if stored:
            await accounts.notify_accounts_changed(self._pool)

    async def _stop(self, bot_id: int) -> None:
        running = self._running.pop(bot_id, None)
        if running is None:
            return
        for task in running.tasks:
            await _cancel(task)
        with contextlib.suppress(Exception):
            await running.client.stop()
        logger.info(f"[yellow]Bot unwatched:[/] {bot_id}")