from __future__ import annotations import asyncio import logging from html import escape from typing import Any, Dict, Optional from aiogram import Bot from aiogram.enums import ChatType, ParseMode from aiogram.exceptions import ( TelegramAPIError, TelegramBadRequest, TelegramForbiddenError, TelegramRetryAfter, ) from aiogram.types import Message, MessageReactionUpdated, ReactionTypeCustomEmoji, ReactionTypeEmoji, User from aiogram.utils.text_decorations import html_decoration from bot.db.repository import Client, ClientRepository, MessageLink from bot.i18n import ( DEFAULT_LANGUAGE, format_labeled_html, format_manager_html, format_prefix, t, ) logger = logging.getLogger(__name__) # General topic in forum groups uses thread id 1 (or None in some clients). GENERAL_TOPIC_ID = 1 TOPIC_NAME_MAX_LEN = 128 CAPTION_MAX_LEN = 1024 # How many times to wait out Telegram flood control before giving up. FLOOD_RETRY_ATTEMPTS = 5 async def _with_flood_retry(api_call, *args, **kwargs): """Call Telegram API; on 429 wait retry_after and try again.""" for attempt in range(1, FLOOD_RETRY_ATTEMPTS + 1): try: return await api_call(*args, **kwargs) except TelegramRetryAfter as exc: wait = max(int(exc.retry_after), 1) method = getattr(api_call, "__name__", type(api_call).__name__) if attempt >= FLOOD_RETRY_ATTEMPTS: logger.error( "Rate limited on %s after %s attempt(s); giving up (retry_after=%ss)", method, attempt, wait, ) raise logger.warning( "Rate limited on %s; waiting %ss (attempt %s/%s)", method, wait, attempt, FLOOD_RETRY_ATTEMPTS, ) await asyncio.sleep(wait) def build_topic_name(user: User) -> str: name = (user.full_name or "Client").strip() or "Client" if user.username: label = f"{name} (@{user.username}) · {user.id}" else: label = f"{name} · {user.id}" if len(label) <= TOPIC_NAME_MAX_LEN: return label # Keep trailing id; trim the name part. suffix = f" · {user.id}" if user.username: suffix = f" (@{user.username}){suffix}" keep = TOPIC_NAME_MAX_LEN - len(suffix) return f"{name[: max(1, keep)].rstrip()}{suffix}" def _is_captionable(message: Message) -> bool: return bool( message.photo or message.document or message.video or message.audio or message.animation or message.voice ) def _html_body(text: Optional[str], entities) -> str: if not text: return "" return html_decoration.unparse(text, entities or []) def _safe_prefixed_caption( prefix_html: str, caption: Optional[str], entities, ) -> str: """Build prefix + caption HTML without chopping mid-entity / mid-escape.""" if len(prefix_html) > CAPTION_MAX_LEN: return prefix_html[: CAPTION_MAX_LEN - 1] + "…" if not caption: return prefix_html rich = f"{prefix_html}\n\n{_html_body(caption, entities)}" if len(rich) <= CAPTION_MAX_LEN: return rich # Overflow: fall back to escaped plain text and shrink until it fits. plain = caption while plain: candidate = f"{prefix_html}\n\n{escape(plain)}" if len(candidate) <= CAPTION_MAX_LEN: if plain != caption: with_ellipsis = f"{prefix_html}\n\n{escape(plain + '…')}" if len(with_ellipsis) <= CAPTION_MAX_LEN: return with_ellipsis return candidate plain = plain[:-1] return prefix_html def client_display_name(user: User) -> str: if user.username: return f"@{user.username}" name = (user.full_name or "").strip() return name or "Client" def client_prefix_html(user: User) -> str: return f"{escape(client_display_name(user))}:" def bot_reactions_from_update(event: MessageReactionUpdated) -> list: """Bots may set at most one reaction; skip paid types.""" for reaction in event.new_reaction: if isinstance(reaction, ReactionTypeEmoji): return [ReactionTypeEmoji(emoji=reaction.emoji)] if isinstance(reaction, ReactionTypeCustomEmoji): return [ReactionTypeCustomEmoji(custom_emoji_id=reaction.custom_emoji_id)] return [] def _is_message_gone_error(exc: BaseException) -> bool: text = str(exc).lower() return any( marker in text for marker in ( "message to edit not found", "message to delete not found", "message can't be edited", "message not found", "message_id_invalid", ) ) def _is_reply_error(exc: BaseException) -> bool: text = str(exc).lower() # Telegram uses "replied" (y→i), so "reply" is not a substring of "replied". return any( marker in text for marker in ( "message to be replied", "replied message", "message to reply not found", ) ) or ("reply" in text and "not found" in text) def _is_topic_missing_error(exc: BaseException) -> bool: text = str(exc).lower() upper = str(exc).upper() return any( marker in text for marker in ( "message thread not found", "thread not found", "topic not found", "topic_deleted", ) ) or "TOPIC_DELETED" in upper def _thread_matches(expected: int, actual: Optional[int]) -> bool: """False when Telegram dropped the thread (e.g. deleted topic → General).""" if actual is None or actual == GENERAL_TOPIC_ID: return False return actual == expected class RelayService: def __init__(self, repo: ClientRepository, manager_group_id: int) -> None: self._repo = repo self._manager_group_id = manager_group_id self._bot_id: Optional[int] = None self._topic_locks: Dict[int, asyncio.Lock] = {} self._topic_locks_guard = asyncio.Lock() def set_bot_id(self, bot_id: int) -> None: self._bot_id = bot_id def is_self(self, user_id: Optional[int]) -> bool: return self._bot_id is not None and user_id == self._bot_id async def _lock_for_user(self, user_id: int) -> asyncio.Lock: async with self._topic_locks_guard: lock = self._topic_locks.get(user_id) if lock is None: lock = asyncio.Lock() self._topic_locks[user_id] = lock return lock async def ensure_client(self, user: User) -> Client: return await self._repo.upsert_client( telegram_id=user.id, username=user.username, first_name=user.first_name, last_name=user.last_name, ) async def _topic_exists(self, bot: Bot, topic_id: int) -> bool: """True if the forum topic still exists. Uses send_chat_action (no visible side effects). Closed topics still accept chat actions; deleted ones raise a topic-missing error. """ try: await bot.send_chat_action( chat_id=self._manager_group_id, action="typing", message_thread_id=topic_id, ) return True except TelegramBadRequest as exc: if _is_topic_missing_error(exc): return False logger.warning("Unexpected error verifying topic %s: %s", topic_id, exc) return True async def ensure_topic(self, bot: Bot, user: User, client: Client) -> int: if client.topic_id is not None: if await self._topic_exists(bot, client.topic_id): return client.topic_id logger.warning( "Stored topic %s for client %s is gone; will recreate", client.topic_id, user.id, ) await self._repo.clear_topic_id(user.id) lock = await self._lock_for_user(user.id) async with lock: # Re-read under lock — another coroutine may have created the topic. fresh = await self._repo.get_by_telegram_id(user.id) if fresh is not None and fresh.topic_id is not None: if await self._topic_exists(bot, fresh.topic_id): return fresh.topic_id await self._repo.clear_topic_id(user.id) topic = await _with_flood_retry( bot.create_forum_topic, chat_id=self._manager_group_id, name=build_topic_name(user), ) await self._repo.set_topic_id(user.id, topic.message_thread_id) logger.info( "Created topic %s for client %s", topic.message_thread_id, user.id, ) return topic.message_thread_id async def _link_messages( self, client_telegram_id: int, client_message_id: int, manager_message_id: int, ) -> None: try: await self._repo.save_message_link( client_telegram_id=client_telegram_id, client_message_id=client_message_id, manager_message_id=manager_message_id, ) except Exception: logger.exception( "Failed to save message link client=%s/%s manager=%s", client_telegram_id, client_message_id, manager_message_id, ) async def _reply_to_manager_id( self, client_telegram_id: int, message: Message, ) -> Optional[int]: reply = message.reply_to_message if reply is None: return None link = await self._repo.get_link_by_client_message(client_telegram_id, reply.message_id) return link.manager_message_id if link else None async def _reply_to_client_id(self, message: Message) -> Optional[int]: reply = message.reply_to_message if reply is None: return None link = await self._repo.get_link_by_manager_message(reply.message_id) return link.client_message_id if link else None async def _call_with_optional_reply(self, api_call, kwargs: Dict[str, Any]): """Call Telegram API; retry once without reply_to if the reply target is gone.""" try: return await _with_flood_retry(api_call, **kwargs) except TelegramBadRequest as exc: if kwargs.get("reply_to_message_id") is not None and _is_reply_error(exc): retry = dict(kwargs) retry.pop("reply_to_message_id", None) return await _with_flood_retry(api_call, **retry) raise async def _copy_client_message( self, bot: Bot, message: Message, topic_id: int, user: User, ) -> tuple[Optional[int], Optional[int]]: """Copy/send client content into the topic. Returns (manager_message_id, actual_message_thread_id). actual_message_thread_id is set for send_message; None for copy_message (Telegram only returns MessageId there). """ prefix = client_prefix_html(user) reply_to = await self._reply_to_manager_id(user.id, message) if message.text is not None and not _is_captionable(message): body = _html_body(message.text, message.entities) kwargs: Dict[str, Any] = { "chat_id": self._manager_group_id, "message_thread_id": topic_id, "text": format_labeled_html(f"{client_display_name(user)}:", body), "parse_mode": ParseMode.HTML, } if reply_to is not None: kwargs["reply_to_message_id"] = reply_to sent = await self._call_with_optional_reply(bot.send_message, kwargs) return sent.message_id, sent.message_thread_id if _is_captionable(message): caption = _safe_prefixed_caption( prefix, message.caption, message.caption_entities, ) kwargs = { "chat_id": self._manager_group_id, "from_chat_id": message.chat.id, "message_id": message.message_id, "message_thread_id": topic_id, "caption": caption, "parse_mode": ParseMode.HTML, } if reply_to is not None: kwargs["reply_to_message_id"] = reply_to result = await self._call_with_optional_reply(bot.copy_message, kwargs) return result.message_id, None # Stickers, video notes, etc. have no caption — topic already identifies the client. kwargs = { "chat_id": self._manager_group_id, "from_chat_id": message.chat.id, "message_id": message.message_id, "message_thread_id": topic_id, } if reply_to is not None: kwargs["reply_to_message_id"] = reply_to result = await self._call_with_optional_reply(bot.copy_message, kwargs) return result.message_id, None async def _recreate_topic_for_user(self, bot: Bot, user: User, stale_topic_id: int) -> int: logger.warning("Topic %s missing for client %s; recreating", stale_topic_id, user.id) await self._repo.clear_topic_id(user.id) client = await self._repo.get_by_telegram_id(user.id) if client is None: client = await self.ensure_client(user) else: # Ensure ensure_topic does not trust the cleared id from a stale object. client.topic_id = None return await self.ensure_topic(bot, user, client) async def relay_client_to_manager(self, bot: Bot, message: Message) -> None: if message.from_user is None: return user = message.from_user client = await self.ensure_client(user) topic_id = await self.ensure_topic(bot, user, client) try: manager_message_id, actual_thread = await self._copy_client_message( bot, message, topic_id, user ) except TelegramBadRequest as exc: # Topic may have been deleted manually — recreate once. if not _is_topic_missing_error(exc): logger.exception("Failed to copy client message to topic %s", topic_id) raise topic_id = await self._recreate_topic_for_user(bot, user, topic_id) manager_message_id, actual_thread = await self._copy_client_message( bot, message, topic_id, user ) # Text sends can silently land in General when the topic is gone. if ( manager_message_id is not None and actual_thread is not None and not _thread_matches(topic_id, actual_thread) ): logger.warning( "Message for client %s landed in thread %s instead of %s; recreating topic", user.id, actual_thread, topic_id, ) await self._safe_delete(bot, self._manager_group_id, manager_message_id) topic_id = await self._recreate_topic_for_user(bot, user, topic_id) manager_message_id, actual_thread = await self._copy_client_message( bot, message, topic_id, user ) if ( manager_message_id is not None and actual_thread is not None and not _thread_matches(topic_id, actual_thread) ): await self._safe_delete(bot, self._manager_group_id, manager_message_id) raise RuntimeError( f"Failed to deliver client {user.id} message into topic {topic_id}" ) if manager_message_id is not None: await self._link_messages(user.id, message.message_id, manager_message_id) async def _deliver_manager_message( self, bot: Bot, client: Client, message: Message, ) -> Optional[int]: prefix = format_prefix(client.language, "prefix_manager") reply_to = await self._reply_to_client_id(message) if message.text is not None and not _is_captionable(message): body = _html_body(message.text, message.entities) kwargs: Dict[str, Any] = { "chat_id": client.telegram_id, "text": format_manager_html(client.language, body), "parse_mode": ParseMode.HTML, } if reply_to is not None: kwargs["reply_to_message_id"] = reply_to sent = await self._call_with_optional_reply(bot.send_message, kwargs) return sent.message_id if _is_captionable(message): caption = _safe_prefixed_caption( prefix, message.caption, message.caption_entities, ) kwargs = { "chat_id": client.telegram_id, "from_chat_id": message.chat.id, "message_id": message.message_id, "caption": caption, "parse_mode": ParseMode.HTML, } if reply_to is not None: kwargs["reply_to_message_id"] = reply_to result = await self._call_with_optional_reply(bot.copy_message, kwargs) return result.message_id kwargs = { "chat_id": client.telegram_id, "from_chat_id": message.chat.id, "message_id": message.message_id, } if reply_to is not None: kwargs["reply_to_message_id"] = reply_to result = await self._call_with_optional_reply(bot.copy_message, kwargs) return result.message_id async def _notify_managers(self, bot: Bot, topic_id: int, key: str) -> None: try: await _with_flood_retry( bot.send_message, chat_id=self._manager_group_id, message_thread_id=topic_id, text=t(DEFAULT_LANGUAGE, key), ) except TelegramAPIError: logger.exception("Failed to notify managers (%s) in topic %s", key, topic_id) async def relay_manager_to_client(self, bot: Bot, message: Message) -> None: topic_id = message.message_thread_id if topic_id is None or topic_id == GENERAL_TOPIC_ID: return client = await self._repo.get_by_topic_id(topic_id) if client is None: logger.warning("No client mapped to topic %s", topic_id) return try: client_message_id = await self._deliver_manager_message(bot, client, message) if client_message_id is not None: await self._link_messages( client.telegram_id, client_message_id, message.message_id, ) except TelegramForbiddenError: logger.warning("Client %s blocked the bot", client.telegram_id) await self._notify_managers(bot, topic_id, "client_blocked") except TelegramBadRequest as exc: logger.exception( "Failed to deliver manager message to client %s: %s", client.telegram_id, exc, ) await self._notify_managers(bot, topic_id, "deliver_failed") async def relay_client_edit(self, bot: Bot, message: Message) -> None: """Mirror a client message edit into the manager topic.""" if message.from_user is None: return link = await self._repo.get_link_by_client_message( message.from_user.id, message.message_id, ) if link is None: return user = message.from_user try: if message.text is not None and not _is_captionable(message): body = _html_body(message.text, message.entities) await _with_flood_retry( bot.edit_message_text, chat_id=self._manager_group_id, message_id=link.manager_message_id, text=format_labeled_html(f"{client_display_name(user)}:", body), parse_mode=ParseMode.HTML, ) return if _is_captionable(message): caption = _safe_prefixed_caption( client_prefix_html(user), message.caption, message.caption_entities, ) await _with_flood_retry( bot.edit_message_caption, chat_id=self._manager_group_id, message_id=link.manager_message_id, caption=caption, parse_mode=ParseMode.HTML, ) except TelegramBadRequest as exc: if _is_message_gone_error(exc): await self._delete_paired_after_gone( bot, client_telegram_id=link.client_telegram_id, client_message_id=link.client_message_id, manager_message_id=link.manager_message_id, delete_manager=False, delete_client=True, ) else: logger.exception( "Failed to sync client edit to manager message %s", link.manager_message_id, ) async def relay_manager_edit(self, bot: Bot, message: Message) -> None: """Mirror a manager message edit into the client DM.""" topic_id = message.message_thread_id if topic_id is None or topic_id == GENERAL_TOPIC_ID: return link = await self._repo.get_link_by_manager_message(message.message_id) if link is None: return client = await self._repo.get_by_telegram_id(link.client_telegram_id) if client is None: return try: if message.text is not None and not _is_captionable(message): body = _html_body(message.text, message.entities) await _with_flood_retry( bot.edit_message_text, chat_id=client.telegram_id, message_id=link.client_message_id, text=format_manager_html(client.language, body), parse_mode=ParseMode.HTML, ) return if _is_captionable(message): caption = _safe_prefixed_caption( format_prefix(client.language, "prefix_manager"), message.caption, message.caption_entities, ) await _with_flood_retry( bot.edit_message_caption, chat_id=client.telegram_id, message_id=link.client_message_id, caption=caption, parse_mode=ParseMode.HTML, ) except TelegramForbiddenError: await self._notify_managers(bot, topic_id, "client_blocked") except TelegramBadRequest as exc: if _is_message_gone_error(exc): await self._delete_paired_after_gone( bot, client_telegram_id=link.client_telegram_id, client_message_id=link.client_message_id, manager_message_id=link.manager_message_id, delete_manager=True, delete_client=False, ) else: logger.exception( "Failed to sync manager edit to client message %s", link.client_message_id, ) await self._notify_managers(bot, topic_id, "edit_failed") async def delete_paired_from_client(self, bot: Bot, message: Message) -> bool: """Delete the manager-side copy for a client message (reply+/del).""" if message.from_user is None or message.reply_to_message is None: return False link = await self._repo.get_link_by_client_message( message.from_user.id, message.reply_to_message.message_id, ) if link is None: return False await self._delete_both_sides(bot, link.client_telegram_id, link) return True async def delete_paired_from_manager(self, bot: Bot, message: Message) -> bool: """Delete the client-side copy for a manager message (reply+/del).""" if message.reply_to_message is None: return False link = await self._repo.get_link_by_manager_message( message.reply_to_message.message_id, ) if link is None: return False await self._delete_both_sides(bot, link.client_telegram_id, link) return True async def _delete_both_sides( self, bot: Bot, client_telegram_id: int, link: MessageLink, ) -> None: await self._safe_delete(bot, client_telegram_id, link.client_message_id) await self._safe_delete(bot, self._manager_group_id, link.manager_message_id) await self._repo.delete_message_link( client_telegram_id=link.client_telegram_id, client_message_id=link.client_message_id, manager_message_id=link.manager_message_id, ) async def _delete_paired_after_gone( self, bot: Bot, *, client_telegram_id: int, client_message_id: int, manager_message_id: int, delete_manager: bool, delete_client: bool, ) -> None: if delete_client: await self._safe_delete(bot, client_telegram_id, client_message_id) if delete_manager: await self._safe_delete(bot, self._manager_group_id, manager_message_id) await self._repo.delete_message_link( client_telegram_id=client_telegram_id, client_message_id=client_message_id, manager_message_id=manager_message_id, ) async def _safe_delete(self, bot: Bot, chat_id: int, message_id: int) -> None: try: await _with_flood_retry( bot.delete_message, chat_id=chat_id, message_id=message_id, ) except TelegramAPIError: logger.info("Could not delete message %s in chat %s", message_id, chat_id) async def relay_reaction(self, bot: Bot, event: MessageReactionUpdated) -> None: if event.user is not None and self.is_self(event.user.id): return reactions = bot_reactions_from_update(event) if event.chat.type == ChatType.PRIVATE: link = await self._repo.get_link_by_client_message( event.chat.id, event.message_id, ) if link is None: return try: await _with_flood_retry( bot.set_message_reaction, chat_id=self._manager_group_id, message_id=link.manager_message_id, reaction=reactions, ) except TelegramAPIError: logger.exception( "Failed to sync client reaction to manager message %s", link.manager_message_id, ) return if event.chat.id != self._manager_group_id: return link = await self._repo.get_link_by_manager_message(event.message_id) if link is None: return try: await _with_flood_retry( bot.set_message_reaction, chat_id=link.client_telegram_id, message_id=link.client_message_id, reaction=reactions, ) except TelegramAPIError: logger.exception( "Failed to sync manager reaction to client message %s", link.client_message_id, ) def is_relayable_content(message: Message) -> bool: """True for user content we can copy (not service messages).""" return bool( message.text or message.caption or message.photo or message.document or message.video or message.audio or message.voice or message.video_note or message.sticker or message.animation or message.contact or message.location or message.venue or message.poll or message.dice )