beavergram/backend/src/userbot/runner.py

import asyncio
import contextlib
import signal
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import asyncpg
import uvloop

from dependencies.container import container
from userbot.modules.capture import CaptureContext, build_capture_context
from userbot.modules.client import PyroClient
from userbot.modules.jobs import JobConsumer
from utils.env import env
from utils.jobs import enqueue
from utils.listen import Listener
from utils.logging import logger, setup_logging
from utils.read.accounts import (
    ACCOUNTS_CHANGED_CHANNEL,
    inactive_session_names,
    session_device_models,
    sync_account,
)
from utils.read.watches import WATCHES_CHANGED_CHANNEL
from utils.storage import ContentAddressedStorage

setup_logging()

SESSION_FLUSH_INTERVAL = 60


@dataclass
class RunningAccount:
    client: PyroClient
    consumer_task: asyncio.Task
    device_model: str | None


def _sessions_dir() -> Path:
    sessions_dir = Path(env.tg.sessions_dir)
    sessions_dir.mkdir(parents=True, exist_ok=True)
    return sessions_dir


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


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

    @property
    def clients(self) -> list[PyroClient]:
        return [account.client for account in self._running.values()]

    async def sync(self) -> None:
        async with self._lock:
            inactive = await inactive_session_names(self._pool)
            models = await session_device_models(self._pool)
            present: set[str] = set()
            for path in sorted(_sessions_dir().glob("*.session")):
                if path.stem in inactive:
                    await self._log_out(path)
                    continue
                present.add(path.stem)
                device_model = models.get(path.stem)
                running = self._running.get(path.stem)
                if running is not None and running.device_model != device_model:
                    await self._stop(path.stem)
                    running = None
                if running is None:
                    await self._start(path, device_model)
            for session_name in set(self._running) - present:
                await self._stop(session_name)

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

    async def _start(self, path: Path, device_model: str | None) -> None:
        session_name = path.stem
        client = PyroClient(
            session_name, workdir=str(path.parent), device_model=device_model
        )
        try:
            await client.start()
            me = client.me
            if me is None:
                msg = f"session {session_name} is not authorized"
                raise RuntimeError(msg)
            account_id = await sync_account(self._pool, me, session_name)
            client.capture = await build_capture_context(
                client, self._pool, self._storage, account_id
            )
        except Exception:
            logger.exception(f"[red]Failed to start session:[/] {session_name}")
            with contextlib.suppress(Exception):
                await client.stop()
            return
        consumer = JobConsumer(client, self._pool, account_id)
        self._running[session_name] = RunningAccount(
            client, asyncio.create_task(consumer.run()), device_model
        )
        logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
        await enqueue(self._pool, account_id, "sync_dialogs", {})
        await enqueue(self._pool, account_id, "sync_contacts", {})

    async def _stop(self, session_name: str) -> None:
        account = self._running.pop(session_name, None)
        if account is None:
            return
        await _cancel(account.consumer_task)
        with contextlib.suppress(Exception):
            await account.client.stop()
        logger.info(f"[yellow]Client stopped:[/] {session_name}")

    async def _log_out(self, path: Path) -> None:
        account = self._running.pop(path.stem, None)
        if account is not None:
            await _cancel(account.consumer_task)
            client = account.client
        else:
            client = PyroClient(
                path.stem, workdir=str(path.parent), load_handlers=False
            )
        with contextlib.suppress(Exception):
            if account is None:
                await client.start()
            await client.log_out()
        with contextlib.suppress(Exception):
            await client.stop()
        path.unlink(missing_ok=True)  # noqa: ASYNC240
        logger.info(f"[yellow]Account logged out:[/] {path.stem}")


async def _flush_sessions(registry: AccountRegistry) -> None:
    while True:
        await asyncio.sleep(SESSION_FLUSH_INTERVAL)
        for client in registry.clients:
            with contextlib.suppress(Exception):
                await client.storage.save()


def _change_listener(registry: AccountRegistry, tasks: set[asyncio.Task]) -> Listener:
    def spawn(coro: Coroutine[Any, Any, None]) -> None:
        task = asyncio.create_task(coro)
        tasks.add(task)
        task.add_done_callback(tasks.discard)

    def reload(
        make_coro: Callable[[CaptureContext], Coroutine[Any, Any, None]],
    ) -> None:
        for client in registry.clients:
            if client.capture is not None:
                spawn(make_coro(client.capture))

    def on_policy(_payload: str) -> None:
        reload(lambda capture: capture.reload_policies())

    def on_watch(_payload: str) -> None:
        reload(lambda capture: capture.watches.refresh())

    def on_accounts(_payload: str) -> None:
        spawn(registry.sync())

    def on_reconnect() -> None:
        on_policy("")
        on_watch("")
        on_accounts("")

    return Listener(
        "userbot",
        {
            "policy_changed": on_policy,
            WATCHES_CHANGED_CHANNEL: on_watch,
            ACCOUNTS_CHANGED_CHANNEL: on_accounts,
        },
        on_reconnect,
    )


async def runner() -> None:
    pool = await container.get(asyncpg.Pool)
    storage = await container.get(ContentAddressedStorage)

    registry = AccountRegistry(pool, storage)
    tasks: set[asyncio.Task] = set()
    listener = _change_listener(registry, tasks)
    stop = asyncio.Event()
    asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
    flush = asyncio.create_task(_flush_sessions(registry))
    try:
        await registry.sync()
        if not registry.clients:
            logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]")
        await listener.start()
        logger.info("[green]Userbot running.[/]")
        await stop.wait()
    finally:
        await _cancel(flush)
        await listener.stop()
        await registry.close()
        await container.close()


def main() -> None:
    uvloop.install()
    logger.info("Starting userbot...")
    with contextlib.suppress(KeyboardInterrupt):
        asyncio.run(runner())
    logger.info("[red]Userbot stopped.[/]")