rich_text.py

from __future__ import annotations

import html
from html.parser import HTMLParser
import re
from urllib.parse import urlparse


BLOCK_FORMULA_RE = re.compile(r"(\$\$.*?\$\$|\\\[.*?\\\])", re.DOTALL)
INLINE_FORMULA_RE = re.compile(r"(\\\(.*?\\\))", re.DOTALL)
URL_RE = re.compile(r"(?<![\]\"'=])(https?://[^\s<]+)")
CODE_FENCE_RE = re.compile(r"(```.*?```)", re.DOTALL)
PLAIN_MATH_PATTERNS = [
    re.compile(r"(?<![\\A-Za-z0-9_])([A-Za-z\u03c0][A-Za-z0-9\u03c0_{}\s+\-*/^.]*=\s*sqrt\([^,;，。；\n]+\)[A-Za-z0-9\u03c0_{}\s+\-*/^.()]*)"),
    re.compile(r"(?<![\\A-Za-z0-9_])([A-Za-z\u03c0][A-Za-z0-9\u03c0_{}\s+\-*/^.]*=\s*[^,;，。；\n]*\^[^,;，。；\n]*)"),
    re.compile(r"(?<![\\A-Za-z0-9_])([0-9A-Za-z\u03c0_{}\s+\-*/^.()]*\u222b[^,;，。；\n]+?d[A-Za-z])"),
    re.compile(r"(?<![\\A-Za-z0-9_])(sqrt\([^,;，。；\n]+\))"),
    re.compile(r"(?<![\\A-Za-z0-9_])(\u221a\([^,;，。；\n]+\))"),
]

ALLOWED_TAGS = {
    "a",
    "blockquote",
    "br",
    "code",
    "details",
    "div",
    "em",
    "h1",
    "h2",
    "h3",
    "h4",
    "h5",
    "h6",
    "hr",
    "li",
    "ol",
    "p",
    "pre",
    "span",
    "strong",
    "summary",
    "table",
    "tbody",
    "td",
    "th",
    "thead",
    "tr",
    "ul",
}
GLOBAL_ATTRS = {"class"}
TAG_ATTRS = {
    "a": {"href", "title", "target", "rel"},
    "code": {"class"},
    "div": {"class"},
    "pre": {"class"},
    "span": {"class"},
    "table": {"class"},
}
SAFE_CLASSES = {
    "codehilite",
    "formula-block",
    "hll",
    "language-python",
    "language-javascript",
    "language-js",
    "language-bash",
    "language-powershell",
    "language-json",
    "math-inline",
}


def render_message_html(content: str, text_color: str = "#111827") -> str:
    source = _strip_raw_html_tags(content.strip())
    normalized = _normalize_latex_blocks(_normalize_plain_math(source))
    try:
        import markdown
    except ModuleNotFoundError:
        return _sanitize_html(_fallback_render(normalized, text_color))

    rendered = markdown.markdown(
        _auto_link_markdown_urls(normalized),
        extensions=["extra", "fenced_code", "codehilite", "sane_lists", "tables"],
        extension_configs={
            "codehilite": {
                "guess_lang": False,
                "use_pygments": True,
                "noclasses": False,
            }
        },
        output_format="html5",
    )
    return _sanitize_html(rendered)


def _normalize_latex_blocks(content: str) -> str:
    def replace_block(match: re.Match[str]) -> str:
        formula = _strip_block_formula(match.group(1))
        return f"\n\n{_math_block_html(formula)}\n\n"

    content = BLOCK_FORMULA_RE.sub(replace_block, content)

    def replace_inline(match: re.Match[str]) -> str:
        formula = html.escape(match.group(1)[2:-2].strip())
        return _math_inline_html(formula)

    content = INLINE_FORMULA_RE.sub(replace_inline, content)
    return content


def _strip_raw_html_tags(content: str) -> str:
    return re.sub(r"</?[A-Za-z][^>]*>", "", content)


def _normalize_plain_math(content: str) -> str:
    parts = CODE_FENCE_RE.split(content)
    for index, part in enumerate(parts):
        if part.startswith("```"):
            continue
        parts[index] = _wrap_plain_math(part)
    return "".join(parts)


def _wrap_plain_math(text: str) -> str:
    for pattern in PLAIN_MATH_PATTERNS:
        text = _apply_plain_math_pattern(text, pattern)
    return text


def _apply_plain_math_pattern(text: str, pattern: re.Pattern[str]) -> str:
    parts = INLINE_FORMULA_RE.split(text)
    for index, part in enumerate(parts):
        if INLINE_FORMULA_RE.fullmatch(part):
            continue
        parts[index] = pattern.sub(_plain_math_replacement, part)
    return "".join(parts)


def _plain_math_replacement(match: re.Match[str]) -> str:
    expression = match.group(1).strip()
    if _already_math_delimited(expression):
        return expression
    return f"\\({_plain_expression_to_latex(expression)}\\)"


def _plain_expression_to_latex(expression: str) -> str:
    expression = expression.replace("\u03c0", r"\pi ")
    expression = expression.replace("\u222b", r"\int ")
    expression = re.sub(r"sqrt\(([^()]+)\)", r"\\sqrt{\1}", expression)
    expression = re.sub(r"\u221a\(([^()]+)\)", r"\\sqrt{\1}", expression)
    return expression


def _already_math_delimited(expression: str) -> bool:
    return expression.startswith(("\\(", "\\[", "$$")) or "math-inline" in expression


def _auto_link_markdown_urls(text: str) -> str:
    def replace(match: re.Match[str]) -> str:
        raw = match.group(1)
        url = raw.rstrip(".,;:!?)")
        trailing = raw[len(url) :]
        if not _is_safe_url(url):
            return html.escape(raw)
        escaped_url = url.replace(")", "%29")
        return f"[{url}]({escaped_url}){trailing}"

    return URL_RE.sub(replace, text)


def _fallback_render(content: str, text_color: str) -> str:
    parts = BLOCK_FORMULA_RE.split(content)
    rendered = []
    for part in parts:
        if not part:
            continue
        if BLOCK_FORMULA_RE.fullmatch(part):
            rendered.append(_render_formula_block(_strip_block_formula(part)))
        else:
            rendered.append(_render_markdownish_text(part, text_color))
    return "\n".join(rendered)


def _render_markdownish_text(text: str, text_color: str) -> str:
    lines = text.splitlines()
    html_parts: list[str] = []
    ordered_items: list[str] = []
    unordered_items: list[str] = []
    code_lines: list[str] = []
    in_code_block = False
    code_lang = ""

    def flush_list() -> None:
        if ordered_items:
            html_parts.append("<ol>" + "".join(ordered_items) + "</ol>")
            ordered_items.clear()
        if unordered_items:
            html_parts.append("<ul>" + "".join(unordered_items) + "</ul>")
            unordered_items.clear()

    index = 0
    while index < len(lines):
        raw_line = lines[index]
        code_fence = re.match(r"^```(\w+)?\s*$", raw_line.strip())
        if code_fence:
            flush_list()
            if in_code_block:
                language_class = f" language-{html.escape(code_lang)}" if code_lang else ""
                code = html.escape("\n".join(code_lines))
                html_parts.append(f"<pre><code class='{language_class}'>{code}</code></pre>")
                code_lines.clear()
                code_lang = ""
                in_code_block = False
            else:
                in_code_block = True
                code_lang = code_fence.group(1) or ""
            index += 1
            continue

        if in_code_block:
            code_lines.append(raw_line)
            index += 1
            continue

        line = raw_line.strip()
        if not line:
            flush_list()
            html_parts.append("<div class='spacer'></div>")
            index += 1
            continue

        if re.fullmatch(r"-{3,}|\*{3,}", line):
            flush_list()
            html_parts.append("<hr>")
            index += 1
            continue

        if line.startswith(">"):
            flush_list()
            quote = line.lstrip(">").strip()
            html_parts.append(f"<blockquote>{_render_inline(quote, text_color)}</blockquote>")
            index += 1
            continue

        if _is_table_start(lines, index):
            flush_list()
            table_html, index = _render_table(lines, index, text_color)
            html_parts.append(table_html)
            continue

        heading = re.match(r"^(#{1,6})\s+(.+)$", line)
        if heading:
            flush_list()
            level = min(len(heading.group(1)) + 1, 4)
            html_parts.append(f"<h{level}>{_render_inline(heading.group(2), text_color)}</h{level}>")
            index += 1
            continue

        numbered = re.match(r"^\d+[.)]\s+(.+)$", line)
        bullet = re.match(r"^[-*]\s+(.+)$", line)
        if numbered:
            if unordered_items:
                flush_list()
            ordered_items.append(f"<li>{_render_inline(numbered.group(1), text_color)}</li>")
            index += 1
            continue
        if bullet:
            if ordered_items:
                flush_list()
            unordered_items.append(f"<li>{_render_inline(bullet.group(1), text_color)}</li>")
            index += 1
            continue

        flush_list()
        html_parts.append(f"<p>{_render_inline(line, text_color)}</p>")
        index += 1

    flush_list()
    if in_code_block:
        language_class = f" language-{html.escape(code_lang)}" if code_lang else ""
        code = html.escape("\n".join(code_lines))
        html_parts.append(f"<pre><code class='{language_class}'>{code}</code></pre>")
    return "\n".join(html_parts)


def _render_inline(text: str, text_color: str) -> str:
    escaped = html.escape(text)
    escaped = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", escaped)

    def replace_formula(match: re.Match[str]) -> str:
        return _math_inline_html(html.escape(match.group(1)[2:-2].strip()))

    escaped = INLINE_FORMULA_RE.sub(replace_formula, escaped)
    return _auto_link_escaped_urls(escaped)


def _auto_link_escaped_urls(text: str) -> str:
    def replace(match: re.Match[str]) -> str:
        raw = match.group(1)
        url = raw.rstrip(".,;:!?)")
        trailing = raw[len(url) :]
        if not _is_safe_url(url):
            return html.escape(raw)
        safe_url = html.escape(url, quote=True)
        return f'<a href="{safe_url}" target="_blank" rel="noopener noreferrer">{safe_url}</a>{trailing}'

    return URL_RE.sub(replace, text)


def _is_table_start(lines: list[str], index: int) -> bool:
    if index + 1 >= len(lines):
        return False
    header = lines[index].strip()
    separator = lines[index + 1].strip()
    return "|" in header and re.fullmatch(r"\|?\s*:?-{3,}:?\s*(\|\s*:?-{3,}:?\s*)+\|?", separator) is not None


def _render_table(lines: list[str], index: int, text_color: str) -> tuple[str, int]:
    header = _split_table_row(lines[index])
    index += 2
    rows: list[list[str]] = []
    while index < len(lines) and "|" in lines[index].strip():
        rows.append(_split_table_row(lines[index]))
        index += 1

    head_html = "".join(f"<th>{_render_inline(cell, text_color)}</th>" for cell in header)
    body_html = ""
    for row in rows:
        cells = row + [""] * max(0, len(header) - len(row))
        body_html += "<tr>" + "".join(f"<td>{_render_inline(cell, text_color)}</td>" for cell in cells[: len(header)]) + "</tr>"
    return f"<table><thead><tr>{head_html}</tr></thead><tbody>{body_html}</tbody></table>", index


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


def _render_formula_block(formula: str) -> str:
    return _math_block_html(formula)


def _math_inline_html(formula: str) -> str:
    return f'<span class="math-inline">&#92;({formula}&#92;)</span>'


def _math_block_html(formula: str) -> str:
    return f'<div class="formula-block">&#92;[{html.escape(formula)}&#92;]</div>'


def _strip_block_formula(text: str) -> str:
    stripped = text.strip()
    if stripped.startswith("$$") and stripped.endswith("$$"):
        return stripped[2:-2].strip()
    if stripped.startswith("\\[") and stripped.endswith("\\]"):
        return stripped[2:-2].strip()
    return stripped


def _is_safe_url(url: str) -> bool:
    parsed = urlparse(url)
    return parsed.scheme in {"http", "https"} and bool(parsed.netloc)


class _HTMLSanitizer(HTMLParser):
    def __init__(self) -> None:
        super().__init__(convert_charrefs=False)
        self.parts: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        if tag not in ALLOWED_TAGS:
            return
        clean_attrs = self._clean_attrs(tag, attrs)
        attr_text = "".join(f' {name}="{html.escape(value, quote=True)}"' for name, value in clean_attrs)
        self.parts.append(f"<{tag}{attr_text}>")

    def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        self.handle_starttag(tag, attrs)
        if tag not in {"br", "hr"}:
            self.handle_endtag(tag)

    def handle_endtag(self, tag: str) -> None:
        if tag in ALLOWED_TAGS:
            self.parts.append(f"</{tag}>")

    def handle_data(self, data: str) -> None:
        self.parts.append(html.escape(data))

    def handle_entityref(self, name: str) -> None:
        self.parts.append(f"&{name};")

    def handle_charref(self, name: str) -> None:
        self.parts.append(f"&#{name};")

    def _clean_attrs(self, tag: str, attrs: list[tuple[str, str | None]]) -> list[tuple[str, str]]:
        allowed = TAG_ATTRS.get(tag, set()) | GLOBAL_ATTRS
        clean: list[tuple[str, str]] = []
        for name, value in attrs:
            if value is None or name not in allowed:
                continue
            if name == "href" and not _is_safe_url(value):
                continue
            if name == "class":
                classes = [item for item in value.split() if item in SAFE_CLASSES or item.startswith(("k", "s", "c", "m", "n", "o", "p", "bp", "nf", "nc"))]
                if not classes:
                    continue
                value = " ".join(classes)
            if tag == "a" and name == "target":
                value = "_blank"
            if tag == "a" and name == "rel":
                value = "noopener noreferrer"
            clean.append((name, value))
        if tag == "a":
            names = {name for name, _ in clean}
            if "target" not in names:
                clean.append(("target", "_blank"))
            if "rel" not in names:
                clean.append(("rel", "noopener noreferrer"))
        return clean


def _sanitize_html(markup: str) -> str:
    sanitizer = _HTMLSanitizer()
    sanitizer.feed(markup)
    sanitizer.close()
    return "".join(sanitizer.parts)