ForcePilot/backend/package/yuxi/channel/render/style_map.py

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Callable

from yuxi.channel.render.ir import (
    BlockQuote,
    Bold,
    BoldItalic,
    CodeBlock,
    Document,
    Heading,
    Image,
    InlineCode,
    IrNode,
    Italic,
    LineBreak,
    Link,
    ListItem,
    OrderedList,
    Paragraph,
    Strikethrough,
    Table,
    TableCell,
    TableRow,
    Text,
    ThematicBreak,
    UnorderedList,
)

logger = logging.getLogger(__name__)

_PIPE_ESCAPE_RE = str.maketrans({"|": "\\|"})


def render_md_table(
    header_texts: list[str],
    body_texts: list[list[str]],
    aligns: list[str | None],
) -> str:
    if not header_texts and not body_texts:
        return ""
    col_count = len(header_texts) or (len(body_texts[0]) if body_texts else 0)
    if col_count == 0:
        return ""

    while len(header_texts) < col_count:
        header_texts.append("")
    while len(aligns) < col_count:
        aligns.append(None)

    align_map = {
        "left": ":--",
        "center": ":-:",
        "right": "--:",
    }

    def _fmt_row(cells: list[str]) -> str:
        parts = [cell.translate(_PIPE_ESCAPE_RE) for cell in cells[:col_count]]
        while len(parts) < col_count:
            parts.append("")
        return "| " + " | ".join(parts) + " |"

    def _fmt_sep() -> str:
        parts = [align_map.get(a or "", "---") for a in aligns[:col_count]]
        while len(parts) < col_count:
            parts.append("---")
        return "|" + "|".join(parts) + "|"

    lines: list[str] = []
    if header_texts:
        lines.append(_fmt_row(header_texts))
    lines.append(_fmt_sep())
    for row in body_texts:
        lines.append(_fmt_row(row))
    return "\n".join(lines)


def _pad_cell(text: str, width: int, align: str | None) -> str:
    if align == "right":
        return text.rjust(width)
    elif align == "center":
        return text.center(width)
    else:
        return text.ljust(width)


def render_plain_table(
    header_texts: list[str],
    body_texts: list[list[str]],
    aligns: list[str | None],
) -> str:
    if not header_texts and not body_texts:
        return ""
    all_rows = ([header_texts] if header_texts else []) + body_texts
    if not all_rows:
        return ""
    col_count = max(len(row) for row in all_rows)

    col_widths = [0] * col_count
    for row in all_rows:
        for j, cell in enumerate(row[:col_count]):
            col_widths[j] = max(col_widths[j], len(cell) + 2)

    while len(aligns) < col_count:
        aligns.append(None)

    def _fmt_row(cells: list[str]) -> str:
        padded = []
        for j in range(col_count):
            c = cells[j] if j < len(cells) else ""
            padded.append(_pad_cell(c, col_widths[j], aligns[j]))
        return "|" + "|".join(padded) + "|"

    lines: list[str] = []
    if header_texts:
        lines.append(_fmt_row(header_texts))
        sep_parts = ["-" * w for w in col_widths]
        lines.append("|" + "|".join(sep_parts) + "|")
    for row in body_texts:
        lines.append(_fmt_row(row))
    return "\n".join(lines)


@dataclass
class RenderStyleMap:
    bold: Callable[[str], str]
    italic: Callable[[str], str]
    bold_italic: Callable[[str], str]
    strikethrough: Callable[[str], str]
    inline_code: Callable[[str], str]
    code_block: Callable[[str, str | None], str]
    link: Callable[[str, str], str]
    image: Callable[[str, str], str]
    heading: Callable[[int, str], str]
    unordered_list_item: Callable[[str], str]
    ordered_list_item: Callable[[str, int], str]
    blockquote: Callable[[str], str]
    table: Callable[[list[str], list[list[str]], list[str | None]], str] | None = None
    thematic_break: str = ""
    paragraph_separator: str = "\n\n"

    @classmethod
    def markdown(cls) -> RenderStyleMap:
        return cls(
            bold=lambda s: f"**{s}**",
            italic=lambda s: f"*{s}*",
            bold_italic=lambda s: f"***{s}***",
            strikethrough=lambda s: f"~~{s}~~",
            inline_code=lambda s: f"`{s}`",
            code_block=lambda code, lang: f"```{lang or ''}\n{code}\n```",
            link=lambda text, url: f"[{text}]({url})",
            image=lambda alt, url: f"![{alt}]({url})",
            heading=lambda level, text: f"{'#' * level} {text}",
            unordered_list_item=lambda text: f"- {text}",
            ordered_list_item=lambda text, idx: f"{idx}. {text}",
            blockquote=lambda s: "\n".join(f"> {line}" for line in s.split("\n")),
            table=render_md_table,
            thematic_break="---",
        )

    @classmethod
    def plain_text(cls) -> RenderStyleMap:
        return cls(
            bold=lambda s: s,
            italic=lambda s: s,
            bold_italic=lambda s: s,
            strikethrough=lambda s: s,
            inline_code=lambda s: s,
            code_block=lambda code, lang: code,
            link=lambda text, url: f"{text} ({url})" if text else url,
            image=lambda alt, url: f"[{alt or 'Image'}]({url})",
            heading=lambda level, text: text,
            unordered_list_item=lambda text: f"- {text}",
            ordered_list_item=lambda text, idx: f"{idx}. {text}",
            blockquote=lambda s: s,
            table=render_plain_table,
            thematic_break="---",
        )


def _render_inline(nodes: list[IrNode], style: RenderStyleMap) -> str:
    parts: list[str] = []
    for node in nodes:
        if isinstance(node, Text):
            parts.append(node.content)
        elif isinstance(node, Bold):
            parts.append(style.bold(_render_inline(node.children, style)))
        elif isinstance(node, Italic):
            parts.append(style.italic(_render_inline(node.children, style)))
        elif isinstance(node, BoldItalic):
            parts.append(style.bold_italic(_render_inline(node.children, style)))
        elif isinstance(node, Strikethrough):
            parts.append(style.strikethrough(_render_inline(node.children, style)))
        elif isinstance(node, InlineCode):
            parts.append(style.inline_code(node.code))
        elif isinstance(node, Link):
            text = _render_inline(node.children, style)
            parts.append(style.link(text, node.url))
        elif isinstance(node, Image):
            parts.append(style.image(node.alt, node.url))
        elif isinstance(node, LineBreak):
            parts.append("\n")
    return "".join(parts)


def render(doc: Document, style: RenderStyleMap) -> str:
    parts: list[str] = []
    for i, node in enumerate(doc.children):
        if i > 0:
            parts.append(style.paragraph_separator)
        parts.append(_render_block(node, style))
    return "".join(parts)


def _render_block(node: IrNode, style: RenderStyleMap) -> str:
    if isinstance(node, Paragraph):
        return _render_inline(node.children, style)
    elif isinstance(node, Heading):
        text = _render_inline(node.children, style)
        return style.heading(node.level, text)
    elif isinstance(node, CodeBlock):
        return style.code_block(node.code, node.language)
    elif isinstance(node, BlockQuote):
        inner = _render_blocks(node.children, style)
        return style.blockquote(inner)
    elif isinstance(node, UnorderedList):
        lines = []
        for item in node.items:
            text = _render_list_item(item, style)
            lines.append(style.unordered_list_item(text))
        return "\n".join(lines)
    elif isinstance(node, OrderedList):
        lines = []
        for idx, item in enumerate(node.items, start=node.start):
            text = _render_list_item(item, style)
            lines.append(style.ordered_list_item(text, idx))
        return "\n".join(lines)
    elif isinstance(node, ThematicBreak):
        return style.thematic_break
    elif isinstance(node, Table):
        if style.table is not None:
            header_texts = [_render_inline(cell.children, style) for cell in (node.header.cells if node.header else [])]
            aligns = [
                cell.align for cell in (node.header.cells if node.header else node.rows[0].cells if node.rows else [])
            ]
            body_texts = [[_render_inline(cell.children, style) for cell in row.cells] for row in node.rows]
            return style.table(header_texts, body_texts, aligns)
        return ""
    else:
        logger.warning("Unknown IR node type during render: %s", type(node).__name__)
        return ""


def _render_list_item(item: ListItem, style: RenderStyleMap) -> str:
    inline_parts: list[str] = []
    block_parts: list[str] = []
    for child in item.children:
        if isinstance(child, (UnorderedList, OrderedList)):
            block_parts.append(_render_block(child, style))
        elif isinstance(child, Paragraph):
            block_parts.append(_render_block(child, style))
        else:
            inline_parts.append(_render_inline([child], style))
    result = "".join(inline_parts)
    if block_parts:
        indent = "  "
        result += "\n" + "\n".join(indent + line for bp in block_parts for line in bp.split("\n"))
    return result


def _render_blocks(nodes: list[IrNode], style: RenderStyleMap) -> str:
    parts: list[str] = []
    for node in nodes:
        parts.append(_render_block(node, style))
    return "\n".join(parts)