hyperboria/nexus/bot/handlers/search.py

import asyncio
import logging
import re
import time

from grpc import StatusCode
from grpc.experimental.aio import AioRpcError
from library.telegram.base import RequestContext
from nexus.bot.widgets.search_widget import SearchWidget
from nexus.translations import t
from nexus.views.telegram.common import close_button
from nexus.views.telegram.registry import parse_typed_document_to_view
from telethon import (
    events,
    functions,
)

from .base import (
    BaseCallbackQueryHandler,
    BaseHandler,
)


class BaseSearchHandler(BaseHandler):
    should_reset_last_widget = False

    async def do_search(
        self,
        event: events.ChatAction,
        request_context: RequestContext,
        prefetch_message,
        query: str,
        is_group_mode: bool = False,
        is_shortpath_enabled: bool = False,
    ):
        session_id = self.generate_session_id()
        message_id = prefetch_message.id

        request_context.add_default_fields(is_group_mode=is_group_mode, mode='search', session_id=session_id)

        start_time = time.time()
        try:
            search_widget = await SearchWidget.create(
                application=self.application,
                chat=request_context.chat,
                session_id=session_id,
                message_id=message_id,
                request_id=request_context.request_id,
                query=query,
                is_group_mode=is_group_mode,
            )
        except AioRpcError as e:
            actions = [
                self.application.telegram_client.delete_messages(
                    request_context.chat.chat_id,
                    [message_id],
                )
            ]
            if e.code() == StatusCode.INVALID_ARGUMENT:
                too_difficult_picture_url = self.application.config['application'].get('too_difficult_picture_url', '')
                if e.details() == 'url_query_error':
                    actions.append(
                        event.reply(
                            t('INVALID_QUERY_ERROR', language=request_context.chat.language).format(
                                too_difficult_picture_url=too_difficult_picture_url,
                            ),
                            buttons=[close_button()],
                        )
                    )
                elif e.details() == 'invalid_query_error':
                    actions.append(
                        event.reply(
                            t('INVALID_SYNTAX_ERROR', language=request_context.chat.language).format(
                                too_difficult_picture_url=too_difficult_picture_url,
                            ),
                            buttons=[close_button()],
                        )
                    )
                return await asyncio.gather(*actions)
            elif e.code() == StatusCode.CANCELLED:
                maintenance_picture_url = self.application.config['application'].get('maintenance_picture_url', '')
                request_context.error_log(e)
                actions.append(event.reply(
                    t('MAINTENANCE', language=request_context.chat.language).format(
                        maintenance_picture_url=maintenance_picture_url,
                    ),
                    buttons=[close_button()],
                ))
                return await asyncio.gather(*actions)
            raise e

        action = 'documents_found'
        if len(search_widget.scored_documents) == 0:
            action = 'documents_not_found'

        request_context.statbox(
            action=action,
            duration=time.time() - start_time,
            query=f'page:0 query:{query}',
        )

        if len(search_widget.scored_documents) == 1 and is_shortpath_enabled:
            scored_document = search_widget.scored_documents[0]
            document_view = parse_typed_document_to_view(scored_document.typed_document)
            # Second (re-)fetching is required to retrieve duplicates
            document_view = await self.resolve_document(
                schema=scored_document.typed_document.WhichOneof('document'),
                document_id=document_view.id,
                position=0,
                session_id=session_id,
                request_context=request_context,
            )
            view, buttons = document_view.get_view(
                language=request_context.chat.language,
                session_id=session_id,
                bot_external_name=self.application.config['telegram']['bot_external_name'],
                with_buttons=not is_group_mode,
            )
            return await asyncio.gather(
                self.application.telegram_client.edit_message(
                    request_context.chat.chat_id,
                    message_id,
                    view,
                    buttons=buttons,
                ),
            )

        serp, buttons = await search_widget.render()
        return await self.application.telegram_client.edit_message(
            request_context.chat.chat_id,
            message_id,
            serp,
            buttons=buttons,
            link_preview=False,
        )


class SearchHandler(BaseSearchHandler):
    filter = events.NewMessage(incoming=True, pattern=re.compile('^(/search\\s+)?(.*)', flags=re.DOTALL))
    is_group_handler = True
    should_reset_last_widget = False
    is_subscription_required_for_handler = True

    async def ban_handler(self, event: events.ChatAction, request_context: RequestContext, ban_timeout: float):
        logging.getLogger('statbox').info({
            'bot_name': self.application.config['telegram']['bot_name'],
            'action': 'user_flood_ban',
            'mode': 'search',
            'ban_timeout_seconds': ban_timeout,
            'chat_id': request_context.chat.chat_id,
        })
        ban_reason = t(
            'BAN_MESSAGE_TOO_MANY_REQUESTS',
            language=request_context.chat.language
        )
        return await event.reply(t(
            'BANNED_FOR_SECONDS',
            language=request_context.chat.language
        ).format(
            seconds=str(ban_timeout),
            reason=ban_reason,
        ))

    async def handler(self, event: events.ChatAction, request_context: RequestContext):
        ban_timeout = self.application.user_manager.check_search_ban_timeout(user_id=request_context.chat.chat_id)
        if ban_timeout:
            return await self.ban_handler(event, request_context, ban_timeout)
        self.application.user_manager.add_search_time(user_id=request_context.chat.chat_id, search_time=time.time())

        search_prefix = event.pattern_match.group(1)
        query = event.pattern_match.group(2)
        is_group_mode = event.is_group or event.is_channel

        if is_group_mode and not search_prefix:
            return
        if not is_group_mode and search_prefix:
            query = event.raw_text
        prefetch_message = await event.reply(
            t("SEARCHING", language=request_context.chat.language),
        )
        self.application.user_manager.last_widget[request_context.chat.chat_id] = prefetch_message.id
        try:
            await self.do_search(
                event, request_context, prefetch_message,
                query=query,
                is_group_mode=is_group_mode,
                is_shortpath_enabled=True,
            )
        except (AioRpcError, asyncio.CancelledError) as e:
            await asyncio.gather(
                event.delete(),
                prefetch_message.delete(),
            )
            raise e


class SearchEditHandler(BaseSearchHandler):
    filter = events.MessageEdited(incoming=True, pattern=re.compile('^(/search\\s+)?(.*)', flags=re.DOTALL))
    is_group_handler = True
    should_reset_last_widget = False

    async def handler(self, event: events.ChatAction, request_context: RequestContext):
        request_context.add_default_fields(mode='search_edit')
        search_prefix = event.pattern_match.group(1)
        query = event.pattern_match.group(2)
        is_group_mode = event.is_group or event.is_channel
        if is_group_mode and not search_prefix:
            return
        if not is_group_mode and search_prefix:
            query = event.raw_text
        result = await self.application.telegram_client(functions.messages.GetMessagesRequest(
            id=list(range(event.id + 1, event.id + 10)))
        )
        if not result:
            request_context.statbox(action='failed')
            return await event.reply(
                t('REPLY_MESSAGE_HAS_BEEN_DELETED', language=request_context.chat.language),
            )
        for next_message in result.messages:
            if next_message.is_reply and event.id == next_message.reply_to_msg_id:
                request_context.statbox(action='resolved')
                await self.do_search(
                    event,
                    request_context,
                    prefetch_message=next_message,
                    query=query,
                    is_group_mode=is_group_mode,
                )
                return
        request_context.statbox(action='failed')
        return await event.reply(
            t('REPLY_MESSAGE_HAS_BEEN_DELETED', language=request_context.chat.language),
        )


class SearchPagingHandler(BaseCallbackQueryHandler):
    filter = events.CallbackQuery(pattern='^/search_([A-Za-z0-9]+)_([0-9]+)_([0-9]+)$')
    should_reset_last_widget = False

    async def handler(self, event: events.ChatAction, request_context: RequestContext):
        session_id = event.pattern_match.group(1).decode()
        message_id = int(event.pattern_match.group(2).decode())
        page = int(event.pattern_match.group(3).decode())

        request_context.add_default_fields(mode='search_paging', session_id=session_id)

        message = await event.get_message()
        if not message:
            return await event.answer()
        reply_message = await message.get_reply_message()
        if not reply_message:
            return await event.respond(
                t('REPLY_MESSAGE_HAS_BEEN_DELETED', language=request_context.chat.language),
            )

        start_time = time.time()
        query = reply_message.raw_text
        try:
            search_widget = await SearchWidget.create(
                application=self.application,
                chat=request_context.chat,
                session_id=session_id,
                message_id=message_id,
                request_id=request_context.request_id,
                query=query,
                page=page,
            )
        except AioRpcError as e:
            if e.code() == StatusCode.INVALID_ARGUMENT or e.code() == StatusCode.CANCELLED:
                request_context.error_log(e)
                return await event.answer(
                    t('MAINTENANCE_WO_PIC', language=request_context.chat.language),
                )
            raise e

        action = 'documents_found'
        if len(search_widget.scored_documents) == 0:
            action = 'documents_not_found'

        request_context.statbox(
            action=action,
            duration=time.time() - start_time,
            query=f'page:{page} query:{query}',
        )
        serp, buttons = await search_widget.render()
        return await asyncio.gather(
            event.answer(),
            message.edit(serp, buttons=buttons, link_preview=False)
        )