ForcePilot/backend/package/yuxi/channel/extensions/qqbot/format.py

from __future__ import annotations

import logging
import re
from typing import Any

logger = logging.getLogger(__name__)

QQ_MARKDOWN_CHUNK_LIMIT = 5000

QQ_MARKDOWN_SUPPORTED = {
    "**", "__",
    "*", "_",
    "`",
    "```",
    "~~",
    "[",
}


def chunk_markdown_text(text: str, limit: int = QQ_MARKDOWN_CHUNK_LIMIT) -> list[str]:
    if len(text) <= limit:
        return [text]

    chunks = []
    paragraphs = text.split("\n\n")
    current = ""

    for para in paragraphs:
        if len(current) + len(para) + 2 <= limit:
            current = f"{current}\n\n{para}" if current else para
        else:
            if current:
                chunks.append(current)
            if len(para) > limit:
                sub_chunks = _split_long_paragraph(para, limit)
                chunks.extend(sub_chunks[:-1])
                current = sub_chunks[-1]
            else:
                current = para

    if current:
        chunks.append(current)

    return chunks


def _split_long_paragraph(text: str, limit: int) -> list[str]:
    result = []
    sentences = re.split(r"(?<=[。！？.!?])\s*", text)
    current = ""

    for sent in sentences:
        if not sent.strip():
            continue
        if len(current) + len(sent) <= limit:
            current += sent
        else:
            if current:
                result.append(current)
            if len(sent) > limit:
                for i in range(0, len(sent), limit):
                    result.append(sent[i : i + limit])
                current = ""
            else:
                current = sent

    if current:
        result.append(current)

    return result or [text[:limit]]


def md_to_qq_markdown(text: str) -> str:
    placeholder_map: dict[str, str] = {}
    text = _extract_code_blocks(text, placeholder_map)
    text = _handle_headers(text)
    text = _handle_tables(text)
    text = _handle_images(text)
    text = _restore_code_blocks(text, placeholder_map)
    return text


def _extract_code_blocks(text: str, placeholder_map: dict[str, str]) -> str:
    pattern = re.compile(r"```(\w*)\n(.*?)```", re.DOTALL)
    parts: list[str] = []
    idx = 0
    last_end = 0
    for m in pattern.finditer(text):
        idx += 1
        lang = m.group(1)
        code = m.group(2)
        placeholder = f"\n__QBCB{idx}__\n"
        placeholder_map[placeholder] = f"\n```{lang}\n{code}```\n"
        parts.append(text[last_end : m.start()])
        parts.append(placeholder)
        last_end = m.end()
    parts.append(text[last_end:])
    return "".join(parts)


def _restore_code_blocks(text: str, placeholder_map: dict[str, str]) -> str:
    for placeholder, code_block in placeholder_map.items():
        text = text.replace(placeholder, code_block)
    return text


def _handle_headers(text: str) -> str:
    result = []
    for line in text.split("\n"):
        header_match = re.match(r"^(#{1,6})\s+(.+)", line)
        if header_match:
            result.append(f"**{header_match.group(2)}**")
        else:
            result.append(line)
    return "\n".join(result)


def _handle_tables(text: str) -> str:
    lines = text.split("\n")
    if not any("|" in line and line.strip().startswith("|") for line in lines):
        return text

    result = []
    i = 0
    while i < len(lines):
        if "|" in lines[i] and lines[i].strip().startswith("|"):
            table_lines = []
            while i < len(lines) and "|" in lines[i] and lines[i].strip().startswith("|"):
                table_lines.append(lines[i])
                i += 1
            result.append(_table_to_ascii(table_lines))
        else:
            result.append(lines[i])
            i += 1
    return "\n".join(result)


def _table_to_ascii(table_lines: list[str]) -> str:
    rows = []
    for line in table_lines:
        if re.match(r"^\|[\s\-:|]+\|$", line):
            continue
        cells = [cell.strip() for cell in line.strip("|").split("|")]
        rows.append(cells)

    if not rows:
        return ""

    col_widths = [max(len(row[col]) for row in rows) for col in range(len(rows[0]))]

    result = []
    for row in rows:
        padded = [cell.ljust(col_widths[i]) for i, cell in enumerate(row)]
        result.append("| " + " | ".join(padded) + " |")

    return "\n".join(result)


def _handle_images(text: str) -> str:
    return re.sub(r"!\[([^\]]*)\]\(([^)]+)\)", r'<qqimg>\2</qqimg>', text)


def strip_qqbot_mentions(text: str, bot_openid: str | None = None) -> str:
    if bot_openid:
        text = text.replace(f"<@!{bot_openid}>", "").replace(f"<@{bot_openid}>", "")
    text = re.sub(r"<@!\d+>", "", text)
    text = re.sub(r"<@\d+>", "", text)
    return text.strip()