ForcePilot/backend/package/yuxi/channels/adapters/telegram/send.py

from __future__ import annotations

import asyncio
import random
from typing import Any

from telegram import Bot, ReplyParameters
from telegram.error import BadRequest, NetworkError, RetryAfter, TelegramError, TimedOut
from telegram.ext import Application

from yuxi.channels.sdk.format import markdown_to_html, strip_html_tags as strip_all_tags
from yuxi.utils.logging_config import logger


async def send_with_retry(
    bot: Bot,
    chat_id: str,
    payload: dict[str, Any],
    config: dict[str, Any] | None = None,
) -> Any:
    cfg = config or {}
    retry_cfg = cfg.get("retry", {})
    max_retries = retry_cfg.get("attempts", 3) if isinstance(retry_cfg, dict) else 3
    min_delay_ms = retry_cfg.get("min_delay_ms", 400) if isinstance(retry_cfg, dict) else 400
    max_delay_ms = retry_cfg.get("max_delay_ms", 30000) if isinstance(retry_cfg, dict) else 30000
    jitter = retry_cfg.get("jitter", 0.1) if isinstance(retry_cfg, dict) else 0.1

    last_exception = None
    for attempt in range(max_retries):
        try:
            return await bot.send_message(**payload)
        except RetryAfter as e:
            await asyncio.sleep(e.retry_after)
            last_exception = e
        except BadRequest as e:
            if "can't parse entities" in str(e).lower() and payload.get("parse_mode") == "HTML":
                payload = dict(payload)
                payload.pop("parse_mode", None)
                if "text" in payload:
                    payload["text"] = strip_all_tags(payload["text"])
                try:
                    return await bot.send_message(**payload)
                except Exception:
                    raise
            if (
                "reply message not found" in str(e).lower() or "replied message not found" in str(e).lower()
            ) and payload.get("reply_to_message_id"):
                payload = dict(payload)
                payload.pop("reply_to_message_id", None)
                payload.pop("reply_parameters", None)
                logger.debug(f"[Telegram] Reply target not found, retrying without reply_to for chat {chat_id}")
                return await bot.send_message(**payload)
            raise
        except (TimedOut, NetworkError, OSError) as e:
            last_exception = e
            delay = min(max_delay_ms / 1000, (min_delay_ms / 1000) * (2**attempt))
            delay += random.uniform(0, delay * jitter)
            logger.warning(f"[Telegram] Send retry {attempt + 1}/{max_retries} after {delay:.1f}s: {e}")
            await asyncio.sleep(delay)

    if last_exception:
        raise last_exception
    raise TelegramError("Send failed after retries")


async def send_stream_edit(
    application: Application,
    chat_id: str,
    msg_id: str,
    chunk: str,
    config: dict[str, Any] | None = None,
) -> Any:
    html = markdown_to_html(chunk)
    try:
        await application.bot.edit_message_text(
            chat_id=chat_id,
            message_id=int(msg_id),
            text=html,
            parse_mode="HTML",
        )
    except BadRequest as e:
        if "can't parse entities" in str(e).lower():
            clean_text = strip_all_tags(html)
            try:
                await application.bot.edit_message_text(
                    chat_id=chat_id,
                    message_id=int(msg_id),
                    text=clean_text,
                )
            except Exception:
                pass
        else:
            logger.debug(f"[Telegram] Stream edit skipped (BadRequest): {e}")
    except RetryAfter as e:
        await asyncio.sleep(e.retry_after)
    except Exception as e:
        logger.debug(f"[Telegram] Stream edit error: {e}")


async def send_with_format_fallback(
    bot: Bot,
    chat_id: str,
    text: str,
    **kwargs,
) -> Any:
    try:
        return await bot.send_message(
            chat_id=chat_id,
            text=text,
            parse_mode="HTML",
            **kwargs,
        )
    except BadRequest as e:
        if "can't parse entities" in str(e).lower():
            clean_text = strip_all_tags(text)
            logger.warning(f"[Telegram] HTML parse failed, falling back to plain text for chat {chat_id}")
            return await bot.send_message(
                chat_id=chat_id,
                text=clean_text,
                **kwargs,
            )
        raise


async def send_media_with_retry(
    bot: Bot,
    chat_id: str,
    media_type: str,
    media_data: Any,
    caption: str | None = None,
    config: dict[str, Any] | None = None,
    force_document: bool = False,
) -> Any:
    max_retries = 3
    last_exception = None

    for attempt in range(max_retries):
        try:
            kwargs = {}
            if caption:
                kwargs["caption"] = caption
                kwargs["parse_mode"] = "HTML"

            if media_type == "image":
                if force_document:
                    return await bot.send_document(chat_id=chat_id, document=media_data, **kwargs)
                return await bot.send_photo(chat_id=chat_id, photo=media_data, **kwargs)
            elif media_type == "video":
                return await bot.send_video(chat_id=chat_id, video=media_data, **kwargs)
            elif media_type == "audio":
                return await bot.send_audio(chat_id=chat_id, audio=media_data, **kwargs)
            elif media_type == "file":
                return await bot.send_document(chat_id=chat_id, document=media_data, **kwargs)
            else:
                raise ValueError(f"Unsupported media type: {media_type}")
        except RetryAfter as e:
            await asyncio.sleep(e.retry_after)
            last_exception = e
        except (TimedOut, NetworkError, OSError) as e:
            last_exception = e
            await asyncio.sleep(1.0 * (2**attempt))

    if last_exception:
        raise last_exception
    raise TelegramError("Media send failed after retries")


def build_reply_parameters(
    response: Any,
    reply_mode: str,
) -> dict[str, Any] | None:
    reply_to_id = None
    if hasattr(response, "reply_to_message_id"):
        reply_to_id = response.reply_to_message_id
    elif isinstance(response, dict):
        reply_to_id = response.get("reply_to_message_id")

    if reply_mode == "off":
        return None
    if reply_mode == "first":
        if reply_to_id:
            return {"reply_to_message_id": reply_to_id, "message_thread_id": response.metadata.get("thread_id")}
        return None
    if reply_mode in ("all", "batched"):
        if reply_to_id:
            return {"reply_to_message_id": reply_to_id, "message_thread_id": response.metadata.get("thread_id")}
        return None
    return None


def build_quote_payload(response: Any) -> dict[str, Any] | None:
    quote_text = None
    if hasattr(response, "quote_text") and response.quote_text:
        quote_text = response.quote_text
    elif isinstance(response, dict) and response.get("quote_text"):
        quote_text = response["quote_text"]

    if not quote_text:
        return None

    reply_to_id = None
    if hasattr(response, "reply_to_message_id"):
        reply_to_id = response.reply_to_message_id
    elif isinstance(response, dict):
        reply_to_id = response.get("reply_to_message_id")

    if not reply_to_id:
        return None

    params: dict[str, Any] = {
        "message_id": int(reply_to_id),
    }
    if quote_text:
        params["quote"] = quote_text[:256]

    return {"reply_parameters": ReplyParameters(**params)}