bot/telegram_api.py

import os
import re
import html
import logging
import httpx

log = logging.getLogger("telegram")

BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN", "")
PROXY = os.environ.get("TELEGRAM_PROXY", "")

TELEGRAM_MAX_LEN = 4096
SAFE_CHUNK = 3900


def _base():
    if PROXY:
        return PROXY.rstrip("/") + f"/bot{BOT_TOKEN}"
    return f"https://api.telegram.org/bot{BOT_TOKEN}"


async def _post(method: str, payload: dict, timeout: float = 20.0):
    if not BOT_TOKEN:
        log.error("TELEGRAM_BOT_TOKEN missing")
        return {"ok": False, "description": "no token"}

    # Strip None values — Telegram rejects `null` for optional fields
    clean = {k: v for k, v in payload.items() if v is not None}

    url = f"{_base()}/{method}"
    async with httpx.AsyncClient(timeout=timeout) as client:
        r = await client.post(url, json=clean)
        try:
            return r.json()
        except Exception:
            return {"ok": False, "description": r.text}


# ============================================================
# HTML detection + Markdown → Telegram HTML
# ============================================================
_ALLOWED_TAGS_RE = re.compile(
    r"</?(b|strong|i|em|u|s|strike|del|code|pre|a|tg-spoiler)(\s+[^>]*)?>",
    re.IGNORECASE,
)
_CODE_FENCE_RE = re.compile(r"```([a-zA-Z0-9_+\-]*)\n?(.*?)```", re.DOTALL)
_INLINE_CODE_RE = re.compile(r"`([^`\n]+)`")
_BOLD_RE = re.compile(r"\*\*([^\*\n]+?)\*\*")
_ITALIC_RE = re.compile(r"(?<!\*)\*([^\*\n]+?)\*(?!\*)")
_STRIKE_RE = re.compile(r"~~([^~\n]+?)~~")
_LINK_RE = re.compile(r"\[([^\]]+)\]\((https?://[^\s\)]+)\)")
_HEADER_RE = re.compile(r"^(#{1,6})\s+(.*)$", re.MULTILINE)


def _looks_like_html(text: str) -> bool:
    return bool(_ALLOWED_TAGS_RE.search(text))


def md_to_telegram_html(text: str) -> str:
    if text is None:
        return ""

    # Idempotency: if it already has Telegram HTML tags, trust it.
    if _looks_like_html(text):
        if "```" in text:
            text = _CODE_FENCE_RE.sub(
                lambda m: _render_code_block(m.group(1), m.group(2)), text
            )
        return text

    code_blocks = []

    def _stash_code(m):
        lang = (m.group(1) or "").strip()
        body = m.group(2)
        code_blocks.append(_render_code_block(lang, body))
        return f"\x00CODE{len(code_blocks) - 1}\x00"

    text = _CODE_FENCE_RE.sub(_stash_code, text)
    text = html.escape(text, quote=False)
    text = _INLINE_CODE_RE.sub(lambda m: f"<code>{m.group(1)}</code>", text)
    text = _LINK_RE.sub(lambda m: f'<a href="{m.group(2)}">{m.group(1)}</a>', text)
    text = _BOLD_RE.sub(r"<b>\1</b>", text)
    text = _ITALIC_RE.sub(r"<i>\1</i>", text)
    text = _STRIKE_RE.sub(r"<s>\1</s>", text)
    text = _HEADER_RE.sub(lambda m: f"<b>{m.group(2)}</b>", text)

    for i, block in enumerate(code_blocks):
        text = text.replace(f"\x00CODE{i}\x00", block)

    return text


def _render_code_block(lang: str, body: str) -> str:
    body = body.rstrip("\n")
    body_escaped = html.escape(body, quote=False)
    if lang:
        return (
            f'<pre><code class="language-{html.escape(lang)}">'
            f"{body_escaped}</code></pre>"
        )
    return f"<pre><code>{body_escaped}</code></pre>"


# ============================================================
# Splitting
# ============================================================
def split_for_telegram(text: str, max_len: int = SAFE_CHUNK):
    if len(text) <= max_len:
        return [text]

    chunks = []
    remaining = text

    while len(remaining) > max_len:
        cut = remaining.rfind("\n\n", 0, max_len)
        if cut == -1 or cut < max_len * 0.5:
            cut = remaining.rfind("\n", 0, max_len)
        if cut == -1 or cut < max_len * 0.5:
            cut = remaining.rfind(" ", 0, max_len)
        if cut == -1:
            cut = max_len

        piece = remaining[:cut]
        rest = remaining[cut:].lstrip("\n")

        open_pre = piece.count("<pre>")
        close_pre = piece.count("</pre>")
        if open_pre > close_pre:
            unclosed = open_pre - close_pre
            piece += "</code></pre>" * unclosed
            rest = ("<pre><code>" * unclosed) + rest

        chunks.append(piece)
        remaining = rest

    if remaining:
        chunks.append(remaining)

    return chunks


# ============================================================
# Public API
# ============================================================
async def send_message(chat_id, text, **kw):
    """
    Send a message. If `parse_mode` is not passed, defaults to HTML.
    If `parse_mode=None` is explicitly passed, no parse_mode is sent
    (correct for raw streaming previews).
    """
    explicit_no_parse = ("parse_mode" in kw and kw["parse_mode"] is None)

    if explicit_no_parse:
        kw.pop("parse_mode")
        html_text = text
    else:
        kw.setdefault("parse_mode", "HTML")
        html_text = md_to_telegram_html(text) if kw["parse_mode"] == "HTML" else text

    kw.setdefault("disable_web_page_preview", True)

    results = []
    for chunk in split_for_telegram(html_text):
        payload = {"chat_id": chat_id, "text": chunk, **kw}
        r = await _post("sendMessage", payload)
        results.append(r)
        if not r.get("ok"):
            log.warning(f"sendMessage chunk failed: {r}")
            break
    return results[-1] if results else {"ok": False}


async def edit_message(chat_id, message_id, text, **kw):
    explicit_no_parse = ("parse_mode" in kw and kw["parse_mode"] is None)

    if explicit_no_parse:
        kw.pop("parse_mode")
        html_text = text
    else:
        kw.setdefault("parse_mode", "HTML")
        html_text = md_to_telegram_html(text) if kw["parse_mode"] == "HTML" else text

    kw.setdefault("disable_web_page_preview", True)

    chunks = split_for_telegram(html_text)
    first = chunks[0]
    payload = {
        "chat_id": chat_id,
        "message_id": message_id,
        "text": first,
        **kw,
    }
    r = await _post("editMessageText", payload)

    for extra in chunks[1:]:
        await send_message(chat_id, extra,
                           **({"parse_mode": None} if explicit_no_parse else {}))

    return r


async def delete_message(chat_id, message_id):
    return await _post("deleteMessage", {
        "chat_id": chat_id, "message_id": message_id
    })


async def set_webhook(public_url: str, secret: str = ""):
    url = public_url.rstrip("/") + "/webhook"
    payload = {
        "url": url,
        "allowed_updates": ["message", "edited_message"],
        "drop_pending_updates": False,
    }
    if secret:
        payload["secret_token"] = secret
    return await _post("setWebhook", payload)


async def delete_webhook():
    return await _post("deleteWebhook", {"drop_pending_updates": False})