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

from __future__ import annotations

import re

from .constants import MAX_TEXT_LENGTH


def _code_block_handler(match: re.Match) -> str:
    lang = match.group(1) or ""
    code = match.group(2)
    if lang:
        return f'<pre lang="{lang}">{code}</pre>'
    return f"<pre>{code}</pre>"


def markdown_to_flockml(text: str) -> str:
    if not text:
        return ""

    code_blocks: dict[str, str] = {}
    placeholder_prefix = "__FCB_"

    def _store_code(m: re.Match) -> str:
        idx = len(code_blocks)
        key = f"{placeholder_prefix}{idx}__"
        code_blocks[key] = _code_block_handler(m)
        return key

    result = re.sub(r"```(\w*)\n?(.*?)```", _store_code, text, flags=re.DOTALL)

    result = re.sub(r"`([^`]+)`", r"<code>\1</code>", result)
    result = re.sub(r"\*\*(.+?)\*\*", r"<b>\1</b>", result)
    result = re.sub(r"\*(.+?)\*", r"<i>\1</i>", result)
    result = re.sub(r"__(.+?)__", r"<u>\1</u>", result)
    result = re.sub(r"~~(.+?)~~", r"<s>\1</s>", result)
    result = re.sub(r"\[(.+?)\]\((.+?)\)", r'<a href="\2">\1</a>', result)
    result = re.sub(r"^---$", r"<hr/>", result, flags=re.MULTILINE)

    lines = result.split("\n")
    result_lines = _process_list_lines(lines)

    result = "\n".join(result_lines)

    result = re.sub(r"^> (.+)$", r"<blockquote>\1</blockquote>", result, flags=re.MULTILINE)

    for key, value in code_blocks.items():
        result = result.replace(key, value)

    return f"<flockml>{result}</flockml>"


def _process_list_lines(lines: list[str]) -> list[str]:
    result: list[str] = []
    i = 0
    while i < len(lines):
        line = lines[i]
        ul_match = re.match(r"^[\-\*]\s+(.+)$", line)
        ol_match = re.match(r"^(\d+)\.\s+(.+)$", line)

        if ul_match:
            items = [ul_match.group(1)]
            j = i + 1
            while j < len(lines) and re.match(r"^[\-\*]\s+(.+)$", lines[j]):
                items.append(re.match(r"^[\-\*]\s+(.+)$", lines[j]).group(1))
                j += 1
            ul_items = "".join(f"<li>{item}</li>" for item in items)
            result.append(f"<ul>{ul_items}</ul>")
            i = j
            continue

        if ol_match:
            items = [ol_match.group(2)]
            j = i + 1
            while j < len(lines) and re.match(r"^\d+\.\s+(.+)$", lines[j]):
                items.append(re.match(r"^\d+\.\s+(.+)$", lines[j]).group(1))
                j += 1
            ol_items = "".join(f"<li>{item}</li>" for item in items)
            result.append(f"<ol>{ol_items}</ol>")
            i = j
            continue

        result.append(line)
        i += 1

    return result


def sanitize_flockml(flockml: str) -> str:
    from .constants import ALLOWED_FLOCKML_TAGS

    tag_pattern = re.compile(r"</?([a-zA-Z]+)[^>]*>")

    def _sanitize_match(m: re.Match) -> str:
        tag_name = m.group(1).lower()
        if tag_name not in ALLOWED_FLOCKML_TAGS:
            return re.sub(r"[<>]", lambda c: {"<": "&lt;", ">": "&gt;"}[c.group()], m.group(0))
        return m.group(0)

    return tag_pattern.sub(_sanitize_match, flockml)


def strip_markdown_for_plain(text: str) -> str:
    if not text:
        return ""

    text = re.sub(r"```.*?```", "", text, flags=re.DOTALL)
    text = re.sub(r"`([^`]+)`", r"\1", text)
    text = re.sub(r"\*\*(.+?)\*\*", r"\1", text)
    text = re.sub(r"\*(.+?)\*", r"\1", text)
    text = re.sub(r"~~(.+?)~~", r"\1", text)
    text = re.sub(r"\[(.+?)\]\(.+?\)", r"\1", text)

    lines = text.split("\n")
    cleaned = []
    for line in lines:
        stripped = line.strip()
        if re.match(r"^[\-\*]\s+", stripped):
            stripped = re.sub(r"^[\-\*]\s+", "- ", stripped)
        elif re.match(r"^\d+\.\s+", stripped):
            stripped = re.sub(r"^\d+\.\s+", "", stripped)
        elif re.match(r"^>\s+", stripped):
            stripped = re.sub(r"^>\s+", "", stripped)
        elif stripped in ("---", "---", "***"):
            continue
        cleaned.append(stripped)

    return "\n".join(cleaned)


def chunk_text(text: str, limit: int = MAX_TEXT_LENGTH) -> list[str]:
    if not text:
        return []

    if len(text) <= limit:
        return [text]

    chunks: list[str] = []
    paragraphs = re.split(r"\n\s*\n", text)

    for para in paragraphs:
        if len(para) <= limit:
            chunks.append(para)
            continue

        sentences = re.split(r"(?<=[。！？.!?])\s*", para)
        current = ""
        for sentence in sentences:
            if len(current) + len(sentence) <= limit:
                current += sentence
            else:
                if current:
                    chunks.append(current)
                if len(sentence) > limit:
                    words = sentence.split(" ")
                    current = ""
                    for word in words:
                        if len(current) + len(word) + 1 <= limit:
                            current += (" " if current else "") + word
                        else:
                            if current:
                                chunks.append(current)
                            current = word
                    if current:
                        chunks.append(current)
                    current = ""
                else:
                    current = sentence
        if current:
            chunks.append(current)

    return chunks if chunks else [text[:limit]]