ai-tg/utils/formatter.py

"""
Convert AI markdown output → Telegram HTML.
Supports: **bold**, *italic*, `code`, ```blocks```, ~~strike~~, ||spoiler||,
          [links](url), # headers, bullet lists, horizontal rules, markdown tables.
"""
import re
from html import escape as _esc


def _split_table_row(line: str) -> list[str]:
    return [_plain_table_cell(cell.strip()) for cell in line.strip().strip("|").split("|")]


def _plain_table_cell(cell: str) -> str:
    cell = re.sub(r"\[([^\]]+)\]\(([^)]+)\)", r"\1 (\2)", cell)
    cell = re.sub(r"(```|`|\*\*|__|\*|_|~~|\|\|)", "", cell)
    return cell.strip()


def _is_table_separator(line: str) -> bool:
    cells = _split_table_row(line)
    return bool(cells) and all(re.fullmatch(r":?-{3,}:?", cell.strip()) for cell in cells)


def _render_table(lines: list[str]) -> str:
    rows = [_split_table_row(line) for line in lines if not _is_table_separator(line)]
    if not rows:
        return "\n".join(lines)

    width = max(len(row) for row in rows)
    normalized = [row + [""] * (width - len(row)) for row in rows]
    col_widths = [max(len(row[i]) for row in normalized) for i in range(width)]

    rendered = []
    for idx, row in enumerate(normalized):
        rendered.append("  ".join(row[i].ljust(col_widths[i]) for i in range(width)).rstrip())
        if idx == 0 and len(normalized) > 1:
            rendered.append("  ".join("-" * col_widths[i] for i in range(width)).rstrip())
    return "```\n" + "\n".join(rendered) + "\n```"


def _normalize_tables(text: str) -> str:
    lines = text.splitlines()
    out: list[str] = []
    i = 0
    while i < len(lines):
        if (
            "|" in lines[i]
            and i + 1 < len(lines)
            and "|" in lines[i + 1]
            and _is_table_separator(lines[i + 1])
        ):
            block = [lines[i], lines[i + 1]]
            i += 2
            while i < len(lines) and "|" in lines[i] and lines[i].strip():
                block.append(lines[i])
                i += 1
            out.append(_render_table(block))
            continue
        out.append(lines[i])
        i += 1
    return "\n".join(out)


def _protect(text: str) -> tuple[str, list[str]]:
    """Extract code blocks/inline code and replace with placeholders."""
    saved: list[str] = []

    def _save_block(m: re.Match) -> str:
        lang = (m.group(1) or "").strip()
        code = _esc(m.group(2))
        tag = f'<pre><code class="language-{lang}">{code}</code></pre>' if lang else f"<pre>{code}</pre>"
        saved.append(tag)
        return f"\x00BLK{len(saved)-1}\x00"

    def _save_inline(m: re.Match) -> str:
        saved.append(f"<code>{_esc(m.group(1))}</code>")
        return f"\x00INL{len(saved)-1}\x00"

    text = re.sub(r"```(\w*)\n?(.*?)```", _save_block, text, flags=re.DOTALL)
    text = re.sub(r"`([^`\n]+)`", _save_inline, text)
    return text, saved


def _restore(text: str, saved: list[str]) -> str:
    for i, block in enumerate(saved):
        text = text.replace(f"\x00BLK{i}\x00", block)
        text = text.replace(f"\x00INL{i}\x00", block)
    return text


def md_to_tg(text: str) -> str:
    """Convert AI markdown response to Telegram-safe HTML."""
    text = _normalize_tables(text)
    text, saved = _protect(text)

    # Escape HTML special chars in non-code content
    text = _esc(text)

    # ── Block-level ──────────────────────────────────────────
    # Headers → bold line
    text = re.sub(r"^#{1,6}\s+(.+)$", r"<b>\1</b>", text, flags=re.MULTILINE)

    # Horizontal rule
    text = re.sub(r"^[-*_]{3,}$", "────────────────", text, flags=re.MULTILINE)

    # ── Inline ───────────────────────────────────────────────
    # Bold **text** or __text__
    text = re.sub(r"\*\*(.+?)\*\*", r"<b>\1</b>", text, flags=re.DOTALL)
    text = re.sub(r"__(.+?)__", r"<b>\1</b>", text)

    # Italic *text* or _text_ (don't match ** or __)
    text = re.sub(r"(?<!\*)\*(?!\*)([^*\n]+?)(?<!\*)\*(?!\*)", r"<i>\1</i>", text)
    text = re.sub(r"(?<!_)_(?!_)([^_\n]+?)(?<!_)_(?!_)", r"<i>\1</i>", text)

    # Strikethrough ~~text~~
    text = re.sub(r"~~(.+?)~~", r"<s>\1</s>", text, flags=re.DOTALL)

    # Spoiler ||text||
    text = re.sub(r"\|\|(.+?)\|\|", r"<tg-spoiler>\1</tg-spoiler>", text, flags=re.DOTALL)

    # Links [text](url) — must use HTML escape for the URL too
    def _link(m: re.Match) -> str:
        label = m.group(1)
        url = m.group(2).replace('"', "&quot;")
        return f'<a href="{url}">{label}</a>'
    text = re.sub(r"\[([^\]]+)\]\(([^)]+)\)", _link, text)

    text = _restore(text, saved)
    return text


def safe_send(text: str) -> tuple[str, str | None]:
    """Return (formatted_text, parse_mode). Falls back to plain if formatting looks empty."""
    formatted = md_to_tg(text)
    from telegram.constants import ParseMode
    return formatted, ParseMode.HTML


# Matches raw XML-style tool call blocks that some models leak into text output.
# Covers: <tool_call>...</tool_call>, <longcattoolcall>...</longcattoolcall>,
#         <function_calls>...</function_calls>, <function=name>...</function>, etc.
_RAW_TOOL_CALL_RE = re.compile(
    r"<\w*tool\w*>.*?</\w*tool\w*>"
    r"|<function[^>]*>.*?</function[^>]*>"
    r"|<function_calls>.*?</function_calls>"
    r"|<\w*toolcall\w*>.*?</\w*toolcall\w*>",
    re.DOTALL | re.IGNORECASE,
)


def strip_raw_tool_calls(text: str) -> str:
    """Remove raw XML tool-call blocks that models sometimes leak into responses."""
    cleaned = _RAW_TOOL_CALL_RE.sub("", text).strip()
    # Collapse any runs of blank lines left behind
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
    return cleaned