428th-exchange-bot/bot/services/relay.py

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"<b>{escape(client_display_name(user))}:</b>"


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
    )