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

from __future__ import annotations

import re
from enum import StrEnum


class SignalTextStyle(StrEnum):
    BOLD = "BOLD"
    ITALIC = "ITALIC"
    STRIKETHROUGH = "STRIKETHROUGH"
    MONOSPACE = "MONOSPACE"
    SPOILER = "SPOILER"


SIGNAL_STYLE_MAP = {
    "**": SignalTextStyle.BOLD,
    "*": SignalTextStyle.ITALIC,
    "~~": SignalTextStyle.STRIKETHROUGH,
    "`": SignalTextStyle.MONOSPACE,
    "```": SignalTextStyle.MONOSPACE,
    "||": SignalTextStyle.SPOILER,
}

STYLE_RE = re.compile(r"(\*\*|~~|\*|`|```|\|\|)")
LINK_RE = re.compile(r"\[([^\]]*)\]\(([^)]*)\)")
URL_RE = re.compile(r"https?://[^\s)]+")
TRIM_URL_RE = re.compile(r"^(?:https?://)?(?:www\.)?|/$")

_INVALID_CONTROL_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
_ZERO_WIDTH_RE = re.compile(r"[\u200b\u200c\u200d\u200e\u200f\ufeff]")


def sanitize_text(text: str, payload: object | None = None) -> str:
    text = _INVALID_CONTROL_RE.sub("", text)
    text = _ZERO_WIDTH_RE.sub("", text)
    return text.strip()


def _trim_url(url: str) -> str:
    return TRIM_URL_RE.sub("", url).rstrip("/")


def _inline_spans(text: str) -> list[dict]:
    spans: list[dict] = []
    stack: list[tuple[int, str]] = []
    i = 0
    while i < len(text):
        m = STYLE_RE.search(text, i)
        if not m:
            break

        marker = m.group(1)
        start = m.start()

        if marker == "```":
            end = text.find("```", start + 3)
            if end != -1:
                spans.append({"start": start, "length": end + 3 - start, "style": SignalTextStyle.MONOSPACE})
                i = end + 3
                continue
            i = start + 3
            continue

        if stack and stack[-1][1] == marker:
            open_start, _ = stack.pop()
            length = start - open_start - len(marker)
            if length > 0:
                style = SIGNAL_STYLE_MAP.get(marker, SignalTextStyle.BOLD)
                spans.append(
                    {
                        "start": open_start + len(marker),
                        "length": length,
                        "style": style,
                    }
                )
            i = start + len(marker)
        else:
            stack.append((start, marker))
            i = start + len(marker)

    return spans


def markdown_to_signal_formatted_text(
    text: str,
) -> tuple[str, list[dict]]:
    processed = _process_links(text)
    spans = _inline_spans(processed)
    spans = _clamp_styles(processed, spans)
    return processed, spans


def markdown_to_plain_bodies(text: str) -> list[str]:
    processed = _process_links(text)
    items = _convert_headings_and_quotes(processed)
    return [item for item in items if item.strip()]


def _process_links(text: str) -> str:
    def _replace_link(m):
        label = m.group(1)
        url = m.group(2)
        if _trim_url(label) == _trim_url(url):
            return url
        return f"{label} ({url})"

    return LINK_RE.sub(_replace_link, text)


def _convert_headings_and_quotes(text: str) -> list[str]:
    lines = text.split("\n")
    result: list[str] = []
    for line in lines:
        stripped = line.strip()
        if stripped.startswith("#"):
            level = 0
            for ch in stripped:
                if ch == "#":
                    level += 1
                else:
                    break
            heading_text = stripped[level:].strip()
            if heading_text:
                result.append(f"**{heading_text}**")
        elif stripped.startswith("> "):
            result.append(stripped)
        else:
            result.append(line)
    return result


def _clamp_styles(text: str, spans: list[dict]) -> list[dict]:
    text_len = len(text)
    clamped = []
    for span in spans:
        start = span["start"]
        length = span["length"]
        if start < 0:
            length += start
            start = 0
        if start >= text_len:
            continue
        if start + length > text_len:
            length = text_len - start
        if length <= 0:
            continue
        clamped.append({"start": start, "length": length, "style": span["style"]})
    return clamped


def format_text_styles_to_rpc(spans: list[dict]) -> list[str]:
    return [f"{s['start']}:{s['length']}:{s['style']}" for s in spans]


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

    chunks = []
    start = 0
    while start < len(text):
        end = min(start + limit, len(text))
        if end < len(text):
            break_point = text.rfind("\n", start, end)
            if break_point > start:
                end = break_point + 1
            else:
                break_point = text.rfind(" ", start, end)
                if break_point > start:
                    end = break_point + 1
        chunks.append(text[start:end])
        start = end
    return chunks


def _generate_media_placeholder(mime_type: str | None) -> str:
    if not mime_type:
        return ""
    kind = _kind_from_mime(mime_type)
    if kind == "image":
        return "<media:image>"
    return f"<media:{kind}>"


def _kind_from_mime(mime_type: str) -> str:
    main = mime_type.split("/")[0].lower() if mime_type else ""
    known = {"image", "video", "audio", "document"}
    return main if main in known else "document"