quickbot/main.py

from contextlib import asynccontextmanager
from typing import Callable, Any
from aiogram import Bot, Dispatcher
from aiogram.client.default import DefaultBotProperties
from aiogram.types import Message, BotCommand as AiogramBotCommand
from aiogram.utils.callback_answer import CallbackAnswerMiddleware
from aiogram.utils.i18n import I18n
from fastapi import FastAPI
from fastapi.applications import Lifespan, AppType
from secrets import token_hex
from logging import getLogger

from .config import Config
from .fsm.db_storage import DbStorage
from .middleware.telegram import AuthMiddleware, I18nMiddleware
from .model.user import UserBase
from .model.entity_metadata import EntityMetadata
from .model.descriptors import BotCommand
from .router import Router


logger = getLogger(__name__)

@asynccontextmanager
async def default_lifespan(app: "QBotApp"):
    logger.debug("starting qbot app")

    if app.lifespan_bot_init:

        if app.config.USE_NGROK:
            app.ngrok_init()

        await app.bot_init()

    logger.info("qbot app started")

    if app.lifespan:
        async with app.lifespan(app):
            yield
    else:
        yield

    logger.info("stopping qbot app")

    if app.lifespan_bot_init:
        await app.bot_close()
        
        if app.config.USE_NGROK:
            app.ngrok_stop()
    
    logger.info("qbot app stopped")


class QBotApp[UserType: UserBase](FastAPI):
    """
    Main class for the QBot application
    """

    def __init__(
        self,
        user_class: UserType = None,
        config: Config | None = None,
        bot_start: Callable[
            [
                Callable[[Message, Any], tuple[UserType, bool]],
                Message,
                Any,
            ],
            None,
        ] = None,
        lifespan: Lifespan[AppType] | None = None,
        lifespan_bot_init: bool = True,
        allowed_updates: list[str] | None = None,
        *args,
        **kwargs,
    ):
        if config is None:
            config = Config()

        if user_class is None:
            from .model.default_user import DefaultUser

            user_class = DefaultUser

        self.allowed_updates = allowed_updates or ["message", "callback_query"]

        self.user_class = user_class
        self.entity_metadata: EntityMetadata = user_class.entity_metadata
        self.config = config
        self.lifespan = lifespan
        self.bot = Bot(
            token=self.config.TELEGRAM_BOT_TOKEN,
            default=DefaultBotProperties(parse_mode="HTML"),
        )

        dp = Dispatcher(storage=DbStorage())

        i18n = I18n(path="locales", default_locale="en", domain="messages")
        i18n_middleware = I18nMiddleware(user_class=user_class, i18n=i18n)
        i18n_middleware.setup(dp)
        dp.callback_query.middleware(CallbackAnswerMiddleware())

        from .bot.handlers.start import router as start_router

        dp.include_router(start_router)

        from .bot.handlers.menu.main import router as main_menu_router

        auth = AuthMiddleware(user_class=user_class)
        main_menu_router.message.middleware.register(auth)
        main_menu_router.callback_query.middleware.register(auth)
        dp.include_router(main_menu_router)

        self.dp = dp

        self.bot_auth_token = token_hex(128)

        self.start_handler = bot_start
        self.bot_commands = dict[str, BotCommand]()

        self.lifespan_bot_init = lifespan_bot_init

        super().__init__(lifespan=default_lifespan, *args, **kwargs)

        from .api_route.telegram import router as telegram_router

        self.include_router(telegram_router, prefix="/api/telegram", tags=["telegram"])
        self.root_router = Router()
        self.root_router._commands = self.bot_commands
        self.command = self.root_router.command

    
    def register_routers(self, *routers: Router):
        for router in routers:
            for command_name, command in router._commands.items():
                self.bot_commands[command_name] = command


    def ngrok_init(self):
        try:
            from pyngrok import ngrok
            from pyngrok.conf import PyngrokConfig

        except ImportError:
            logger.error("pyngrok is not installed")
            raise

        tunnel = ngrok.connect(
            self.config.API_PORT,
            pyngrok_config=PyngrokConfig(auth_token=self.config.NGROK_AUTH_TOKEN),
        )
        self.config.NGROK_URL = tunnel.public_url
    

    def ngrok_stop(self):
        try:
            from pyngrok import ngrok

        except ImportError:
            logger.error("pyngrok is not installed")
            raise

        ngrok.disconnect(self.config.NGROK_URL)
        ngrok.kill()


    
    async def bot_init(self):
        
        commands_captions = dict[str, list[tuple[str, str]]]()

        for command_name, command in self.bot_commands.items():
            if command.show_in_bot_commands:
                if isinstance(command.caption, str) or command.caption is None:
                    if "default" not in commands_captions:
                        commands_captions["default"] = []
                    commands_captions["default"].append(
                        (command_name, command.caption or command_name)
                    )
                else:
                    for locale, description in command.caption.items():
                        if locale not in commands_captions:
                            commands_captions[locale] = []
                        commands_captions[locale].append((command_name, description))

        for locale, commands in commands_captions.items():
            await self.bot.set_my_commands(
                [
                    AiogramBotCommand(command=command[0], description=command[1])
                    for command in commands
                ],
                language_code=None if locale == "default" else locale,
            )

        await self.bot.set_webhook(
            url=f"{self.config.API_URL}/api/telegram/webhook",
            drop_pending_updates=True,
            allowed_updates=self.allowed_updates,
            secret_token=self.bot_auth_token,
        )


    async def bot_close(self):
        await self.bot.delete_webhook()