ForcePilot/backend/package/yuxi/channels/adapters/signal/format.py

from __future__ import annotations

import re
from dataclasses import dataclass, field


@dataclass
class StyleRange:
    start: int
    length: int
    style: str


BOLD = "BOLD"
ITALIC = "ITALIC"
STRIKETHROUGH = "STRIKETHROUGH"
MONOSPACE = "MONOSPACE"
SPOILER = "SPOILER"

_MARKDOWN_PATTERNS: list[tuple[str, str]] = [
    (r"\*\*(.+?)\*\*", BOLD),
    (r"__(.+?)__", BOLD),
    (r"\*(.+?)\*", ITALIC),
    (r"_(.+?)_", ITALIC),
    (r"~~(.+?)~~", STRIKETHROUGH),
    (r"`(.+?)`", MONOSPACE),
    (r"\|\|(.+?)\|\|", SPOILER),
]

_HEADING_PATTERN = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)
_BLOCKQUOTE_PATTERN = re.compile(r"^>\s?(.+)$", re.MULTILINE)
_TABLE_SEP_PATTERN = re.compile(r"^\|?[-:|\s]+\|?$")
_LINK_PATTERN = re.compile(r"\[([^\]]+)\]\(([^)]+)\)")
_URL_PATTERN = re.compile(r"https?://\S+")


@dataclass
class FormattedText:
    body: str
    styles: list[StyleRange] = field(default_factory=list)


def markdown_to_signal_styles(
    text: str,
    table_mode: str = "bullets",
    heading_style: str = "bold",
    blockquote_prefix: str = "> ",
) -> FormattedText:
    result = text
    styles: list[StyleRange] = []

    result = _convert_tables(result, table_mode)

    result, link_ranges = _process_links(result)

    for pattern, style_name in _MARKDOWN_PATTERNS:
        result, styles = _apply_markdown_pattern(result, pattern, style_name, styles, link_ranges)

    result = _convert_headings(result, heading_style)
    result = _convert_blockquotes(result, blockquote_prefix)

    styles = _merge_adjacent_styles(sorted(styles, key=lambda s: s.start))

    return FormattedText(body=result, styles=styles)


def _convert_tables(text: str, table_mode: str) -> str:
    if table_mode != "bullets":
        return text

    lines = text.split("\n")
    result: list[str] = []
    in_table = False
    table_rows: list[list[str]] = []

    for line in lines:
        stripped = line.strip()
        if stripped.startswith("|") and stripped.endswith("|"):
            if _TABLE_SEP_PATTERN.match(stripped):
                continue
            cells = [c.strip() for c in stripped[1:-1].split("|")]
            table_rows.append(cells)
            in_table = True
            continue
        else:
            if in_table and table_rows:
                result.extend(_table_rows_to_bullets(table_rows))
                table_rows = []
            in_table = False
            result.append(line)

    if table_rows:
        result.extend(_table_rows_to_bullets(table_rows))

    return "\n".join(result)


def _table_rows_to_bullets(rows: list[list[str]]) -> list[str]:
    if not rows:
        return []
    bullets: list[str] = []
    is_header = True
    for row in rows:
        if is_header and len(rows) > 1:
            bullets.append(" • ".join(row))
            is_header = False
        else:
            for cell in row:
                bullets.append(f"• {cell}")
    return bullets


def _process_links(text: str) -> tuple[str, list[tuple[int, int]]]:
    link_replacements: list[tuple[int, int, str, str]] = []
    for m in _LINK_PATTERN.finditer(text):
        label = m.group(1)
        url = m.group(2)
        start, end = m.span()
        if label == url or _url_text_equivalent(label, url):
            replacement = url
        else:
            replacement = f"{label} ({url})"
        link_replacements.append((start, end, replacement, url))

    result = text
    offset = 0
    link_ranges: list[tuple[int, int]] = []
    for start, end, replacement, _url in sorted(link_replacements, key=lambda x: x[0]):
        old_len = end - start
        new_len = len(replacement)
        actual_start = start + offset
        result = result[:actual_start] + replacement + result[actual_start + old_len :]
        link_ranges.append((actual_start, actual_start + new_len))
        offset += new_len - old_len

    return result, link_ranges


def _url_text_equivalent(label: str, url: str) -> bool:
    clean_label = label.strip().rstrip("/").lower()
    clean_url = url.strip().rstrip("/").lower()
    if clean_label == clean_url:
        return True
    if clean_url.startswith("https://"):
        if clean_label == clean_url.removeprefix("https://"):
            return True
    if clean_url.startswith("http://"):
        if clean_label == clean_url.removeprefix("http://"):
            return True
    return False


def _apply_markdown_pattern(
    text: str,
    pattern: str,
    style_name: str,
    existing_styles: list[StyleRange],
    link_ranges: list[tuple[int, int]],
) -> tuple[str, list[StyleRange]]:
    compiled = re.compile(pattern)
    replacements: list[tuple[int, int, str]] = []
    new_styles: list[StyleRange] = []

    for m in compiled.finditer(text):
        inner = m.group(1)
        start, end = m.span()
        if _overlaps_any(start, end, link_ranges):
            continue
        replacement = inner
        replacements.append((start, end, replacement))
        new_styles.append(StyleRange(start=start, length=len(inner), style=style_name))

    result = text
    offset = 0
    final_styles = list(existing_styles)

    replacements.sort(key=lambda x: x[0])
    for start, end, replacement in replacements:
        actual_start = start + offset
        old_len = end - start
        new_len = len(replacement)
        result = result[:actual_start] + replacement + result[actual_start + old_len :]

        for ns in new_styles:
            if ns.start == start:
                ns.start = actual_start
        offset += new_len - old_len

    for ns in new_styles:
        if ns.start + ns.length <= len(result):
            final_styles.append(ns)

    for existing in final_styles:
        if existing in new_styles:
            continue
        for rep_start, rep_end, _ in replacements:
            if existing.start >= rep_start and existing.start < rep_end:
                existing.start += len(replacement) - (rep_end - rep_start)

    return result, final_styles


def _overlaps_any(start: int, end: int, ranges: list[tuple[int, int]]) -> bool:
    for r_start, r_end in ranges:
        if start < r_end and end > r_start:
            return True
    return False


def _convert_headings(text: str, heading_style: str) -> str:
    if heading_style == "bold":

        def _replace(m: re.Match) -> str:
            content = m.group(2)
            return f"**{content}**"

        return _HEADING_PATTERN.sub(_replace, text)
    return text


def _convert_blockquotes(text: str, blockquote_prefix: str) -> str:
    lines = text.split("\n")
    result = []
    in_blockquote = False
    prefix = blockquote_prefix or "> "

    for line in lines:
        m = _BLOCKQUOTE_PATTERN.match(line)
        if m:
            result.append(f"{prefix}{m.group(1)}")
            in_blockquote = True
        else:
            if in_blockquote and line.strip() == "":
                in_blockquote = False
            result.append(line)

    return "\n".join(result)


def _merge_adjacent_styles(styles: list[StyleRange]) -> list[StyleRange]:
    if not styles:
        return []
    merged: list[StyleRange] = []
    for style in styles:
        if merged and merged[-1].style == style.style and merged[-1].start + merged[-1].length == style.start:
            merged[-1].length += style.length
        else:
            merged.append(StyleRange(start=style.start, length=style.length, style=style.style))
    return merged


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

    if chunk_mode == "length":
        return [text[i : i + limit] for i in range(0, len(text), limit)]

    chunks: list[str] = []
    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:
                for i in range(0, len(para), limit):
                    chunks.append(para[i : i + limit])
                current = ""
            else:
                current = para

    if current:
        chunks.append(current)
    return chunks or [text]


def clamp_styles_to_length(styles: list[StyleRange], body: str) -> list[StyleRange]:
    body_len = len(body)
    result: list[StyleRange] = []
    for s in styles:
        if s.start >= body_len:
            continue
        end = s.start + s.length
        if end > body_len:
            s.length = body_len - s.start
        if s.length > 0:
            result.append(s)
    return result