deep_research/scripts/lib/markdown_chunker.py

"""Markdown 切块 / 合并工具。

把一篇 Markdown 按标题层级（# / ##）切成"翻译单元"或"润色单元"，
每个单元带稳定的 ID，便于断点续传和按需重跑。

核心约定：
- H1（`# `）是一级块，通常对应 Chapter / 封面 / 前置件
- H2（`## `）是二级块，对应一个 section 或独立前置件（Disclaimer / Executive Summary / Abstract / Glossary / References）
- 没有任何标题的文件头（frontmatter 区）归到第 0 块

切块粒度默认到 H2；如果某个 H2 下的正文特别长可以再按 H3 切，但这是 polish 的事情，
translate 一般不需要。
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable


HEADER_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*$", re.MULTILINE)


@dataclass
class MarkdownBlock:
    """一个翻译/润色单元。"""

    order: int                 # 在原文中的顺序（0-based）
    level: int                 # 0 = frontmatter; 1/2/... = H1/H2/...
    title: str                 # 标题原文（不含 # 号）；frontmatter 为空串
    anchor: str                # 稳定 ID，用于断点续传（order + title hash）
    content: str               # 完整内容（包含标题行本身，除 frontmatter 块外）
    parent_order: int | None = None  # H2 的父 H1 order；H1 为 None
    word_count: int = 0        # 英文 word count 估算（只含 a-z）
    char_count: int = 0        # 字符数（含中文）
    meta: dict = field(default_factory=dict)

    @property
    def short_title(self) -> str:
        t = self.title.strip()
        if len(t) <= 50:
            return t
        return t[:47] + "..."


def _count_words(text: str) -> int:
    return len(re.findall(r"[A-Za-z]+(?:[-'][A-Za-z]+)*", text))


def _stable_anchor(order: int, title: str) -> str:
    """order + title 生成稳定短 ID。同一文件改顺序不变，改标题重算。"""
    h = hashlib.sha1(title.strip().encode("utf-8")).hexdigest()[:8]
    return f"b{order:03d}-{h}"


def split_by_headers(text: str, max_level: int = 2) -> list[MarkdownBlock]:
    """把 Markdown 按 H1..H{max_level} 切块。

    返回的 block 按出现顺序排列。第 0 块可能是 frontmatter（level=0，无 title）。
    """
    lines = text.splitlines(keepends=False)
    # 先找出所有 header 位置
    header_positions: list[tuple[int, int, str]] = []  # (line_idx, level, title)
    in_code_block = False
    for i, ln in enumerate(lines):
        if ln.lstrip().startswith("```"):
            in_code_block = not in_code_block
            continue
        if in_code_block:
            continue
        m = re.match(r"^(#{1,6})\s+(.+?)\s*$", ln)
        if m:
            level = len(m.group(1))
            if level <= max_level:
                header_positions.append((i, level, m.group(2)))

    blocks: list[MarkdownBlock] = []
    order = 0

    # frontmatter：第一个 header 前的所有内容
    first_header_line = header_positions[0][0] if header_positions else len(lines)
    frontmatter_text = "\n".join(lines[:first_header_line]).rstrip()
    if frontmatter_text.strip():
        blocks.append(
            MarkdownBlock(
                order=order,
                level=0,
                title="",
                anchor=_stable_anchor(order, "__frontmatter__"),
                content=frontmatter_text,
                word_count=_count_words(frontmatter_text),
                char_count=len(frontmatter_text),
            )
        )
        order += 1

    # 为每个 header 建一个 block：内容 = 本 header 行 + 下一 header 行前所有内容
    last_h1_order: int | None = None
    for idx, (line_idx, level, title) in enumerate(header_positions):
        end_line = (
            header_positions[idx + 1][0]
            if idx + 1 < len(header_positions)
            else len(lines)
        )
        content = "\n".join(lines[line_idx:end_line]).rstrip()
        block = MarkdownBlock(
            order=order,
            level=level,
            title=title,
            anchor=_stable_anchor(order, title),
            content=content,
            parent_order=last_h1_order if level > 1 else None,
            word_count=_count_words(content),
            char_count=len(content),
        )
        blocks.append(block)
        if level == 1:
            last_h1_order = order
        order += 1
    return blocks


def merge_blocks(blocks: Iterable[MarkdownBlock], separator: str = "\n\n") -> str:
    """按 order 拼回完整 Markdown。"""
    return separator.join(b.content for b in sorted(blocks, key=lambda x: x.order))


def chapter_stem_from_title(title: str) -> str:
    """从 `# Chapter 1 — foo bar` 生成文件名 stem 如 `ch01`。找不到就 fallback。"""
    m = re.search(r"chapter\s+(\d+)", title, re.IGNORECASE)
    if m:
        return f"ch{int(m.group(1)):02d}"
    # Executive Summary / Abstract / Glossary / Disclaimer / References / Version History
    slug = re.sub(r"[^a-z0-9]+", "-", title.lower()).strip("-")
    return slug[:40] or "untitled"


def count_chinese_chars(text: str) -> int:
    return sum(1 for c in text if "\u4e00" <= c <= "\u9fff")


if __name__ == "__main__":
    import sys

    src = Path(sys.argv[1])
    for b in split_by_headers(src.read_text(encoding="utf-8")):
        print(f"[{b.order:3d}] L{b.level} {b.word_count:>5}w {b.char_count:>6}c  {b.short_title}")