tiku/tiku-server/app/services/docx_builder.py

"""docx 合并 + 导出排版模块 — 使用 deepcopy 保留图片和公式"""

from copy import deepcopy
from datetime import datetime
from pathlib import Path

from docx import Document
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.shared import Pt, Cm


TYPE_ORDER = ["单选题", "多选题", "填空题", "解答题"]
TYPE_HEADERS = {
    "单选题": "一、单选题",
    "多选题": "二、多选题",
    "填空题": "三、填空题",
    "解答题": "四、解答题",
}


def build_export_docx(
    problem_ids: list[str],
    problems_dir: str | Path,
    title: str = "试卷",
    include_solution: bool = True,
) -> Path:
    """从题目 ID 列表合并生成导出 docx（保留图片和公式）"""
    problems_dir = Path(problems_dir)
    doc = Document()

    # 设置默认样式
    style = doc.styles["Normal"]
    style.font.name = "宋体"
    style.font.size = Pt(10.5)
    style.element.rPr.rFonts.set(qn('w:eastAsia'), '宋体')

    # 设置页面边距
    for section in doc.sections:
        section.top_margin = Cm(2)
        section.bottom_margin = Cm(2)
        section.left_margin = Cm(2.5)
        section.right_margin = Cm(2.5)

    # 标题
    title_para = doc.add_paragraph()
    title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
    title_run = title_para.add_run(title)
    title_run.bold = True
    title_run.font.size = Pt(14)
    title_run.font.name = "宋体"
    title_run._element.rPr.rFonts.set(qn('w:eastAsia'), '宋体')

    subtitle = "教师版" if include_solution else "学生版"
    sub_para = doc.add_paragraph()
    sub_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
    sub_run = sub_para.add_run(subtitle)
    sub_run.font.size = Pt(10.5)

    doc.add_paragraph()

    # 设置 Word 自动编号
    _setup_auto_numbering(doc)

    # 从数据库获取题目并按题型分组
    from app.database import get_connection
    conn = get_connection(problems_dir.parent / "bank.db")
    try:
        problems = []
        for pid in problem_ids:
            row = conn.execute("SELECT * FROM problems WHERE id = ?", (pid,)).fetchone()
            if row:
                problems.append(dict(row))

        groups: dict[str, list] = {}
        for p in problems:
            qt = p.get("question_type", "")
            if qt not in groups:
                groups[qt] = []
            groups[qt].append(p)

        sorted_groups = [(qt, groups[qt]) for qt in TYPE_ORDER if qt in groups]
        for qt, items in groups.items():
            if qt not in TYPE_ORDER:
                sorted_groups.append((qt, items))

        question_num = 1
        for qt, items in sorted_groups:
            # 题型标题
            header_text = TYPE_HEADERS.get(qt, qt)
            header_para = doc.add_paragraph()
            header_run = header_para.add_run(header_text)
            header_run.bold = True
            header_run.font.size = Pt(12)
            header_run.font.name = "宋体"
            header_run._element.rPr.rFonts.set(qn('w:eastAsia'), '宋体')

            for item in items:
                pid = item["id"]
                docx_path = problems_dir / f"{pid}.docx"
                if not docx_path.exists():
                    continue

                _copy_question_deep(doc, docx_path, question_num, include_solution)
                question_num += 1

    finally:
        conn.close()

    # 保存临时文件 → 复制图片 → 最终保存
    papers_dir = problems_dir.parent / "papers"
    papers_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    suffix = "教师版" if include_solution else "学生版"
    safe_title = title.replace("/", "_").replace(" ", "_")
    output_path = papers_dir / f"{safe_title}_{ts}_{suffix}.docx"
    doc.save(str(output_path))

    # 从源 docx 复制图片到导出 docx
    _copy_images_for_export(output_path, problem_ids, problems_dir)

    return output_path


def _copy_question_deep(dst_doc: Document, src_path: Path, num: int, include_solution: bool):
    """Deep-copy 题目段落到目标文档，保留图片和公式，并加上自动编号"""
    src_doc = Document(str(src_path))
    src_body = src_doc.element.body
    dst_body = dst_doc.element.body

    in_answer = False
    first_para = True  # 标记是否为该题的第一个段落（需要编号）
    src_paras = src_body.findall(qn('w:p'))

    for i, para_elem in enumerate(src_paras):
        # Get text from the paragraph element
        all_text = ''
        for run_elem in para_elem.findall(qn('w:r')):
            for t_elem in run_elem.findall(qn('w:t')):
                all_text += (t_elem.text or '')
        text = all_text.strip()

        if not text:
            dst_body.append(deepcopy(para_elem))
            continue

        if text.startswith("【答案】"):
            in_answer = True
            if not include_solution:
                break

        if in_answer and not include_solution:
            break

        copied = deepcopy(para_elem)

        # 给第一个非空段落加自动编号
        if first_para and text:
            _apply_numbering(copied)
            first_para = False

        dst_body.append(copied)


def _apply_numbering(para_element):
    """给段落应用自动编号 (numId=1, ilvl=0)"""
    pPr = para_element.find(qn('w:pPr'))
    if pPr is None:
        pPr = OxmlElement('w:pPr')
        para_element.insert(0, pPr)

    numPr = OxmlElement('w:numPr')
    ilvl = OxmlElement('w:ilvl')
    ilvl.set(qn('w:val'), '0')
    numPr.append(ilvl)

    numId = OxmlElement('w:numId')
    numId.set(qn('w:val'), '1')
    numPr.append(numId)

    pPr.append(numPr)


def _setup_auto_numbering(doc: Document):
    """在文档中设置 Word 自动编号（每道题加题号）"""
    numbering_part = doc.part.numbering_part
    numbering_xml = numbering_part._element

    # 定义抽象编号 (abstractNum)
    abstract_num = OxmlElement('w:abstractNum')
    abstract_num.set(qn('w:abstractNumId'), '0')

    lvl = OxmlElement('w:lvl')
    lvl.set(qn('w:ilvl'), '0')
    lvl.set(qn('w:tentative'), '0')
    start = OxmlElement('w:start')
    start.set(qn('w:val'), '1')
    lvl.append(start)

    num_fmt = OxmlElement('w:numFmt')
    num_fmt.set(qn('w:val'), 'decimal')
    lvl.append(num_fmt)

    lvl_text = OxmlElement('w:lvlText')
    lvl_text.set(qn('w:val'), '%1.')
    lvl.append(lvl_text)

    lvl_jc = OxmlElement('w:lvlJc')
    lvl_jc.set(qn('w:val'), 'left')
    lvl.append(lvl_jc)

    p_pr = OxmlElement('w:pPr')
    ind = OxmlElement('w:ind')
    ind.set(qn('w:left'), '420')
    ind.set(qn('w:hanging'), '420')
    p_pr.append(ind)
    lvl.append(p_pr)

    abstract_num.append(lvl)
    numbering_xml.append(abstract_num)

    # 定义编号实例 (num)
    num = OxmlElement('w:num')
    num.set(qn('w:numId'), '1')
    num_ref = OxmlElement('w:abstractNumId')
    num_ref.set(qn('w:val'), '0')
    num.append(num_ref)
    numbering_xml.append(num)


def _copy_images_for_export(output_path: Path, problem_ids: list[str], problems_dir: Path):
    """从各个单题 docx 中复制图片到导出 docx"""
    import zipfile
    from lxml import etree
    import shutil

    # 收集导出 docx 中引用的所有图片 rId，以及需要从哪些源文件复制
    ns_a = 'http://schemas.openxmlformats.org/drawingml/2006/main'
    ns_r_embed = 'http://schemas.openxmlformats.org/officeDocument/2006/relationships'
    ns_rel = 'http://schemas.openxmlformats.org/package/2006/relationships'

    with zipfile.ZipFile(str(output_path), 'r') as dst_zip:
        doc_xml = dst_zip.read('word/document.xml')
        try:
            rels_xml = dst_zip.read('word/_rels/document.xml.rels')
        except KeyError:
            rels_xml = None
        try:
            ct_xml = dst_zip.read('[Content_Types].xml')
        except KeyError:
            ct_xml = None

    tree = etree.fromstring(doc_xml)
    needed_embeds = set()
    for blip in tree.findall(f'.//{{{ns_a}}}blip'):
        embed = blip.get(f'{{{ns_r_embed}}}embed')
        if embed:
            needed_embeds.add(embed)

    if not needed_embeds:
        return

    # 从源文件复制图片
    img_data_map = {}  # {target_filename: image_bytes}
    embed_target_map = {}  # {rId: target_filename}

    for pid in problem_ids:
        src_docx = problems_dir / f'{pid}.docx'
        if not src_docx.exists():
            continue
        with zipfile.ZipFile(str(src_docx)) as src_zip:
            try:
                src_rels = etree.fromstring(src_zip.read('word/_rels/document.xml.rels'))
            except KeyError:
                continue
            for rel in src_rels:
                r_id = rel.get('Id')
                if r_id in needed_embeds:
                    target = rel.get('Target', '')
                    if target not in img_data_map:
                        try:
                            img_data_map[target] = src_zip.read(f'word/{target}')
                            embed_target_map[r_id] = target
                        except KeyError:
                            pass

    if not img_data_map:
        return

    # 写入导出 docx
    tmp_path = output_path.with_suffix('.tmp.docx')
    skip = {'word/_rels/document.xml.rels', 'word/document.xml', '[Content_Types].xml'}

    with zipfile.ZipFile(str(output_path), 'r') as dst_zip:
        with zipfile.ZipFile(str(tmp_path), 'w', zipfile.ZIP_DEFLATED) as out_zip:
            for item in dst_zip.infolist():
                if item.filename in skip:
                    continue
                out_zip.writestr(item, dst_zip.read(item.filename))

            # 添加图片
            for target, data in img_data_map.items():
                out_zip.writestr(f'word/{target}', data)

            # 更新关系
            if rels_xml:
                rels_tree = etree.fromstring(rels_xml)
                next_rid = max(
                    (int(rel.get('Id', 'rId0')[3:]) for rel in rels_tree if rel.get('Id', '').startswith('rId')),
                    default=100
                ) + 1
                for old_rid, target in embed_target_map.items():
                    rel = etree.SubElement(rels_tree, f'{{{ns_rel}}}Relationship')
                    rel.set('Id', f'rId{next_rid}')
                    rel.set('Type', 'http://schemas.openxmlformats.org/officeDocument/2006/relationships/image')
                    rel.set('Target', target)
                    # 更新 document.xml 中的引用
                    for blip in tree.findall(f'.//{{{ns_a}}}blip'):
                        if blip.get(f'{{{ns_r_embed}}}embed') == old_rid:
                            blip.set(f'{{{ns_r_embed}}}embed', f'rId{next_rid}')
                    next_rid += 1
                out_zip.writestr('word/_rels/document.xml.rels',
                    etree.tostring(rels_tree, xml_declaration=True, encoding='UTF-8', standalone=True))

            # 更新 document.xml
            out_zip.writestr('word/document.xml',
                etree.tostring(tree, xml_declaration=True, encoding='UTF-8', standalone=True))

            # 更新 Content_Types
            if ct_xml:
                ct_tree = etree.fromstring(ct_xml)
                ct_ns = 'http://schemas.openxmlformats.org/package/2006/content-types'
                for target in img_data_map:
                    ext = target.rsplit('.', 1)[-1].lower()
                    mime_map = {'png': 'image/png', 'jpg': 'image/jpeg', 'jpeg': 'image/jpeg',
                                'gif': 'image/gif', 'bmp': 'image/bmp', 'webp': 'image/webp'}
                    mime = mime_map.get(ext, 'image/png')
                    exists = any(ov.get('PartName') == f'/word/{target}' for ov in ct_tree.findall(f'{{{ct_ns}}}Override'))
                    if not exists:
                        ov = etree.SubElement(ct_tree, f'{{{ct_ns}}}Override')
                        ov.set('PartName', f'/word/{target}')
                        ov.set('ContentType', mime)
                out_zip.writestr('[Content_Types].xml',
                    etree.tostring(ct_tree, xml_declaration=True, encoding='UTF-8', standalone=True))

    shutil.move(str(tmp_path), str(output_path))