mytoolkit/mytoolkit/md_to_pdf.py

"""Markdown to PDF conversion core logic."""

import re
import uuid
from pathlib import Path

import click

from mytoolkit import templates_registry
from mytoolkit.utils import run_command

# Typst callouts exported by mytoolkit/templates/md-to-pdf/{textbook,manual}/template.typ.
# Keep this tuple and _extra_imports (below) in sync when adding or removing `#let` callouts.
# Pandoc escapes leading `#`, so these names are wrapped in a long-enough `{=typst}` fence
# (length chosen so inner ``` from Pandoc does not terminate the fence early).
_TY_CALL_NAMES: tuple[str, ...] = (
    "info-box",
    "tip-box",
    "warn-box",
    "danger-box",
    "theorem",
    "lemma",
    "proof",
    "example",
    "exercise",
)
_TY_CALL_ALT = "|".join(
    re.escape(n) for n in sorted(_TY_CALL_NAMES, key=len, reverse=True)
)
_TY_CALL_START = re.compile(
    rf"^(\s{{0,3}})(#(?:{_TY_CALL_ALT}))(?![A-Za-z0-9-])",
    re.MULTILINE,
)


def _odd_fenced_code_before(md: str, pos: int) -> bool:
    """True if ``pos`` lies inside an odd number of ``` fenced regions (heuristic)."""
    in_fence = False
    for line in md[:pos].split("\n"):
        stripped = line.lstrip()
        if stripped.startswith("```"):
            in_fence = not in_fence
    return in_fence


def _skip_quoted(md: str, i: int) -> int:
    quote = md[i]
    i += 1
    n = len(md)
    while i < n:
        ch = md[i]
        if ch == "\\" and i + 1 < n:
            i += 2
            continue
        if ch == quote:
            return i + 1
        i += 1
    return i


def _consume_balanced(md: str, i: int, open_c: str, close_c: str) -> int | None:
    """If ``md[i]`` is ``open_c``, return index one past matching ``close_c``."""
    if i >= len(md) or md[i] != open_c:
        return None
    depth = 0
    pos = i
    n = len(md)
    while pos < n:
        ch = md[pos]
        if ch in ('"', "'"):
            pos = _skip_quoted(md, pos)
            continue
        if ch == open_c:
            depth += 1
        elif ch == close_c:
            depth -= 1
            if depth == 0:
                return pos + 1
        pos += 1
    return None


def _md_fragment_to_typst(fragment: str) -> str:
    fragment = fragment.strip("\n")
    if not fragment.strip():
        return ""
    result = run_command(
        ["pandoc", "-f", "markdown-citations", "-t", "typst", "-"],
        input=fragment,
        capture_output=True,
        text=True,
    )
    return result.stdout.rstrip()


def _longest_backtick_run(s: str) -> int:
    longest = 0
    i = 0
    n = len(s)
    while i < n:
        if s[i] == "`":
            j = i
            while j < n and s[j] == "`":
                j += 1
            longest = max(longest, j - i)
            i = j
        else:
            i += 1
    return longest


def _typst_raw_fence(inner: str) -> str:
    """Fence long enough to contain Pandoc typst that may include ``` code blocks."""
    run = _longest_backtick_run(inner)
    return "`" * max(3, run + 1)


def preprocess_typst_callouts(md: str) -> str:
    """Wrap template Typst callouts in Pandoc raw Typst blocks so they compile.

    Bodies are converted with Pandoc separately so Markdown and math still work.
    Markdown links that use ``[...](...)`` inside a callout can break bracket matching.
    """
    out: list[str] = []
    last = 0
    for m in _TY_CALL_START.finditer(md):
        if _odd_fenced_code_before(md, m.start()):
            continue
        idx = m.start(2)
        pos = m.end(2)
        while pos < len(md) and md[pos] in " \t":
            pos += 1
        if pos < len(md) and md[pos] == "(":
            end_paren = _consume_balanced(md, pos, "(", ")")
            if end_paren is None:
                continue
            pos = end_paren
            while pos < len(md) and md[pos] in " \t":
                pos += 1
        if pos >= len(md) or md[pos] != "[":
            continue
        bracket_open = pos
        after_body = _consume_balanced(md, bracket_open, "[", "]")
        if after_body is None:
            continue
        body_inner = md[bracket_open + 1 : after_body - 1]
        prefix = md[idx : bracket_open + 1]
        inner_typ = _md_fragment_to_typst(body_inner)
        fence = _typst_raw_fence(prefix + "\n" + inner_typ)
        replacement = f"{fence}{{=typst}}\n{prefix}\n{inner_typ}\n]\n{fence}\n"
        out.append(md[last : m.start()])
        out.append(replacement)
        last = after_body
    out.append(md[last:])
    return "".join(out)


def resolve_markdown_image_paths(md_content: str, base_dir: Path) -> str:
    """Resolve relative image paths in Markdown to absolute paths.

    Supports Markdown image syntax ``![alt](src)`` and HTML ``<img>`` tags.
    Skips URLs, data URIs, and anchor-only references. Missing files are left
    unchanged so pandoc/typst can report the original path in error messages.
    """

    def _replace_md_image(match: re.Match) -> str:
        alt = match.group(1)
        src = match.group(2)
        if src.startswith(("http://", "https://", "data:", "#")):
            return match.group(0)
        p = Path(src)
        if p.is_absolute():
            return match.group(0)
        abs_path = (base_dir / p).resolve()
        if abs_path.exists():
            return f"![{alt}]({abs_path})"
        return match.group(0)

    md_content = re.sub(r"!\[(.*?)\]\((.*?)\)", _replace_md_image, md_content)

    def _replace_html_img(match: re.Match) -> str:
        pre = match.group(1)
        src = match.group(2)
        post = match.group(3)
        if src.startswith(("http://", "https://", "data:", "#")):
            return match.group(0)
        p = Path(src)
        if p.is_absolute():
            return match.group(0)
        abs_path = (base_dir / p).resolve()
        if abs_path.exists():
            return f"{pre}{abs_path}{post}"
        return match.group(0)

    md_content = re.sub(
        r'(<img[^>]+src=["\'])([^"\']+)(["\'][^>]*>)',
        _replace_html_img,
        md_content,
        flags=re.IGNORECASE,
    )
    return md_content


def md_to_typ(inputs: list[Path]) -> str:
    """Convert Markdown files to Typst via pandoc."""
    md_content = ""
    for inp in inputs:
        content = inp.read_text(encoding="utf-8")
        content = resolve_markdown_image_paths(content, inp.parent)
        md_content += content + "\n\n"
    md_content = preprocess_typst_callouts(md_content)
    result = run_command(
        ["pandoc", "-f", "markdown-citations", "-t", "typst", "-"],
        input=md_content,
        capture_output=True,
        text=True,
    )
    typ = result.stdout
    # Pandoc outputs #horizontalrule which Typst does not define; match template rule-blue
    typ = typ.replace(
        "#horizontalrule",
        '#line(length: 100%, stroke: 0.65pt + rgb("#3467A8").lighten(46%))',
    )
    # Remove Pandoc auto-generated heading labels to avoid named-destination
    # links that break with non-ASCII characters in some PDF readers.
    typ = re.sub(r"(^=+ .*?)\n<[^>\n]+>\n", r"\1\n", typ, flags=re.MULTILINE)
    return typ


def expand_bookmarks(pdf_path: Path, max_level: int) -> None:
    """Expand PDF outline nodes up to ``max_level`` without rebuilding."""
    if max_level <= 1:
        return

    from pypdf import PdfReader, PdfWriter
    from pypdf.generic import NumberObject

    reader = PdfReader(str(pdf_path))
    writer = PdfWriter()
    writer.clone_document_from_reader(reader)

    outline_root = writer._root_object.get("/Outlines")
    if not outline_root:
        writer.write(str(pdf_path))
        return

    def update_count(item_ref, depth=0):
        item = item_ref.get_object()
        count = item.get("/Count")
        if count is not None and depth < max_level:
            val = int(count)
            if val < 0:
                item.update({"/Count": NumberObject(-val)})
        first = item.get("/First")
        if first:
            child = first
            while child:
                update_count(child, depth=depth + 1)
                child = child.get_object().get("/Next")

    first = outline_root.get("/First")
    if first:
        update_count(first, depth=0)

    writer.write(str(pdf_path))


def build_pdf(
    template: str,
    inputs: list[Path],
    output: Path,
    bookmark_depth: int = 2,
) -> None:
    """Build PDF from Markdown using Typst.

    ``bookmark_depth`` controls outline expansion: ``1`` keeps only top-level
    bookmarks visible, ``2`` expands to second level (default), and so on.
    """
    templates_dir = templates_registry.get_md_to_pdf_root()
    template_file = templates_dir / template / "template.typ"
    if not template_file.exists():
        click.secho(f"Template not found: {template_file}", fg="red", err=True)
        raise SystemExit(1)

    typ_body = md_to_typ(inputs)

    # Convert pandoc tables to three-line tables
    typ_body = typ_body.replace("#table(", "#three-line-table(")

    tmp_name = f".{uuid.uuid4().hex}.typ"
    tmp_path = templates_dir / tmp_name
    rel_template = f"{template}/template.typ"

    # Pandoc raw Typst callouts (#theorem, #info-box, …) live in the driver file; they must
    # resolve to symbols exported from the template module (not only inside #show: …).
    _extra_imports = {
        "textbook": "theorem, lemma, proof, example, exercise",
        "manual": "info-box, tip-box, warn-box, danger-box",
    }
    extra = _extra_imports.get(template, "")
    import_list = f"{template}, three-line-table"
    if extra:
        import_list += f", {extra}"

    with open(tmp_path, "w", encoding="utf-8") as f:
        f.write(f'#import "{rel_template}": {import_list}\n\n')
        f.write(f"#show: {template}.with()\n\n")
        f.write(typ_body)

    output = output.resolve()
    output.parent.mkdir(parents=True, exist_ok=True)

    try:
        run_command(
            ["typst", "compile", "--root", "/", tmp_name, str(output)],
            cwd=str(templates_dir),
        )
    finally:
        tmp_path.unlink(missing_ok=True)

    if output.exists() and bookmark_depth > 1:
        expand_bookmarks(output, bookmark_depth)