from __future__ import annotations import asyncio import hashlib import logging import time from html import escape from typing import Any, Dict, Optional, Tuple 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 # Collapse identical manager warnings so a tag-change flood cannot spam topics. NOTIFY_COOLDOWN_SECONDS = 120.0 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 _entities_key(entities) -> str: if not entities: return "" chunks = [] for entity in entities: kind = getattr(entity.type, "value", entity.type) url = getattr(entity, "url", None) or "" custom_emoji_id = getattr(entity, "custom_emoji_id", None) or "" chunks.append(f"{kind}:{entity.offset}:{entity.length}:{url}:{custom_emoji_id}") return "|".join(chunks) def _media_key(message: Message) -> str: if message.photo: return f"photo:{message.photo[-1].file_unique_id}" for attr in ("document", "video", "audio", "animation", "voice", "video_note", "sticker"): obj = getattr(message, attr, None) if obj is not None: uid = getattr(obj, "file_unique_id", "") or "" return f"{attr}:{uid}" return "" def message_content_fingerprint(message: Message) -> str: """Hash of fields we relay. Ignores sender_tag, edit_date, and similar.""" payload = "\n".join( ( message.text or "", message.caption or "", _entities_key(message.entities), _entities_key(message.caption_entities), _media_key(message), ) ) return hashlib.sha256(payload.encode("utf-8")).hexdigest() 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 not found", "message_id_invalid", ) ) def _is_not_modified_error(exc: BaseException) -> bool: return "message is not modified" in str(exc).lower() def _is_uneditable_error(exc: BaseException) -> bool: text = str(exc).lower() return any( marker in text for marker in ( "message can't be edited", "message can not be edited", "message_too_old", "message is too old", ) ) def _is_benign_edit_error(exc: BaseException) -> bool: """True for errors that do not mean the paired message is gone.""" return _is_not_modified_error(exc) or _is_uneditable_error(exc) 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() self._notify_last: Dict[Tuple[int, str], float] = {} 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, source: Message, ) -> 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, content_fingerprint=message_content_fingerprint(source), ) 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 _remember_fingerprint(self, link: MessageLink, fingerprint: str) -> None: try: await self._repo.set_link_fingerprint( client_telegram_id=link.client_telegram_id, client_message_id=link.client_message_id, manager_message_id=link.manager_message_id, content_fingerprint=fingerprint, ) except Exception: logger.exception( "Failed to store content fingerprint for manager message %s", link.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, message ) 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: now = time.monotonic() stamp_key = (topic_id, key) last = self._notify_last.get(stamp_key) if last is not None and now - last < NOTIFY_COOLDOWN_SECONDS: logger.info( "Skipping duplicate manager notify %s in topic %s", key, topic_id, ) return self._notify_last[stamp_key] = now 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, message, ) 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 fingerprint = message_content_fingerprint(message) if link.content_fingerprint == fingerprint: 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, ) elif _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, ) await self._remember_fingerprint(link, fingerprint) except TelegramBadRequest as exc: if _is_benign_edit_error(exc): await self._remember_fingerprint(link, fingerprint) return 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 fingerprint = message_content_fingerprint(message) if link.content_fingerprint == fingerprint: 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, ) elif _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, ) await self._remember_fingerprint(link, fingerprint) except TelegramForbiddenError: await self._notify_managers(bot, topic_id, "client_blocked") except TelegramBadRequest as exc: if _is_benign_edit_error(exc): await self._remember_fingerprint(link, fingerprint) return 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 )