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

"""docx 拆分 + 元数据提取模块

导入整卷 docx → 按题拆分 → 生成每题独立 docx
关键：通过 deep-copy 段落 XML 保留图片、公式、格式
"""

import hashlib
import json
import re
import io
import zipfile
from copy import deepcopy
from pathlib import Path
from docx import Document
from docx.oxml.ns import qn
from docx.shared import Pt


# 题型标题关键词映射
TYPE_KEYWORDS = {
    "单选题": "单选题",
    "多选题": "多选题",
    "填空题": "填空题",
    "解答题": "解答题",
}


def detect_type_header(text: str) -> str | None:
    for key, value in TYPE_KEYWORDS.items():
        if key in text:
            return value
    return None


def is_question_number(text: str) -> bool:
    text = text.strip()
    return bool(re.match(r'\d+[．.]', text))


def strip_question_number(text: str) -> str:
    return re.sub(r'^\d+[．.]\s*', '', text.strip())


def extract_metadata_from_text(text: str) -> dict:
    """从答案区文本中提取元数据"""
    result = {
        "answer": "", "difficulty": None, "knowledge_names": [],
        "analysis": "", "detail": "", "highlight": "",
    }
    sections = re.split(r'【(答案|难度|知识点|分析|详解|点睛)】', text)
    start = 0
    if sections and not sections[0].strip():
        start = 1
    for i in range(start, len(sections) - 1, 2):
        tag = sections[i].strip()
        content = sections[i + 1].strip()
        if tag == "答案":
            result["answer"] = content
        elif tag == "难度":
            try:
                result["difficulty"] = float(content.split()[0])
            except (ValueError, IndexError):
                pass
        elif tag == "知识点":
            names = re.split(r'[，,、\n]', content)
            result["knowledge_names"] = [n.strip() for n in names if n.strip()]
        elif tag == "分析":
            result["analysis"] = content
        elif tag == "详解":
            result["detail"] = content
        elif tag == "点睛":
            result["highlight"] = content
    return result


def _get_answer_boundary(paras, start_idx: int, end_idx: int) -> int:
    """在给定段落范围内找到【答案】分隔线的位置（绝对索引）"""
    for i in range(start_idx, end_idx + 1):
        text = paras[i].text.strip()
        if text.startswith("【答案】"):
            return i
    return end_idx + 1  # 没找到答案标记，全部算题干


def split_docx_to_questions(docx_path: str | Path) -> list[dict]:
    """将整卷 docx 拆分为题目段落范围列表

    Returns list of dicts:
      {question_type, question_num, para_start, para_end, answer_start, metadata}
    """
    docx_path = Path(docx_path)
    doc = Document(str(docx_path))
    paras = doc.paragraphs

    questions = []
    current_type = ""
    current_start = None
    question_num = 0

    for i, para in enumerate(paras):
        text = para.text.strip()
        if not text:
            continue

        # 检测题型标题
        detected = detect_type_header(text)
        if detected:
            current_type = detected
            continue

        # 检测题号：新题开始
        if is_question_number(text) and current_type:
            # 保存上一题
            if current_start is not None:
                answer_start = _get_answer_boundary(paras, current_start, i - 1)
                # 提取元数据
                answer_text = "\n".join(
                    paras[j].text.strip()
                    for j in range(answer_start, i)
                )
                questions.append({
                    "question_type": current_type,
                    "question_num": question_num,
                    "para_start": current_start,
                    "para_end": i - 1,
                    "answer_start": answer_start,
                    "metadata": extract_metadata_from_text(answer_text),
                })

            question_num += 1
            current_start = i

    # 保存最后一题
    if current_start is not None:
        answer_start = _get_answer_boundary(paras, current_start, len(paras) - 1)
        answer_text = "\n".join(
            paras[j].text.strip()
            for j in range(answer_start, len(paras))
        )
        questions.append({
            "question_type": current_type,
            "question_num": question_num,
            "para_start": current_start,
            "para_end": len(paras) - 1,
            "answer_start": answer_start,
            "metadata": extract_metadata_from_text(answer_text),
        })

    return questions


def create_single_docx(
    original_docx_path: str | Path,
    question_data: dict,
    import_type: str,
    source_name: str,
    knowledge_map: dict,
    problems_dir: str | Path = "data/problems",
) -> tuple[Path, str, dict]:
    """从原始 docx 的段落范围创建单题 docx（保留图片和格式）

    关键：deep-copy 原始段落的 XML 元素，完全保留图片、公式、格式
    """
    original_docx_path = Path(original_docx_path)
    problems_dir = Path(problems_dir)
    problems_dir.mkdir(parents=True, exist_ok=True)

    src_doc = Document(str(original_docx_path))
    src_paras = src_doc.paragraphs
    src_body = src_doc.element.body

    para_start = question_data["para_start"]
    para_end = question_data["para_end"]
    answer_start = question_data["answer_start"]
    meta = question_data["metadata"]

    # ── 创建新文档 ──
    dst_doc = Document()

    # 复制默认样式
    style = dst_doc.styles["Normal"]
    style.font.name = "宋体"
    style.font.size = Pt(10.5)

    # ── 复制题干段落（para_start → answer_start - 1）──
    for i in range(para_start, answer_start):
        _copy_paragraph(src_body, dst_doc.element.body, i)

    # ── 复制答案区（跳过【难度】【知识点】段）──
    for i in range(answer_start, para_end + 1):
        text = src_paras[i].text.strip()
        if text.startswith("【难度】") or text.startswith("【知识点】"):
            continue
        _copy_paragraph(src_body, dst_doc.element.body, i)

    # ── 处理知识点编码 ──
    knowledge_codes = []
    knowledge_names = meta.get("knowledge_names", [])
    for name in knowledge_names:
        from app.services.knowledge import find_code_by_name_fuzzy
        code = find_code_by_name_fuzzy(name, knowledge_map)
        if code:
            knowledge_codes.append(code)

    # ── 写入元数据到自定义属性 ──
    _set_custom_props(dst_doc, {
        "question_type": question_data["question_type"],
        "difficulty": meta["difficulty"] or 0.5,
        "knowledge_codes": json.dumps(knowledge_codes, ensure_ascii=False),
        "import_type": import_type,
        "comment": "",
        "source": source_name,
    })

    # ── 计算 ID ──
    # 先把新文档保存为临时文件来计算 ID
    tmp_bytes = io.BytesIO()
    dst_doc.save(tmp_bytes)
    checksum = hashlib.sha256(tmp_bytes.getvalue()).hexdigest()
    problem_id = checksum[:8]

    # ── 处理图片：从原始 docx 复制到新 docx ──
    docx_path = problems_dir / f"{problem_id}.docx"
    dst_doc.save(str(docx_path))

    # ── 嵌入图片（从原始 docx 的 ZIP 中提取，写入新 docx 的 ZIP）──
    _copy_images_to_docx(docx_path, original_docx_path)

    # ── 章节 ──
    chapters = _get_chapters(knowledge_codes, knowledge_map)

    info = {
        "id": problem_id,
        "file_path": str(docx_path),
        "checksum": checksum,
        "chapters": chapters,
        "difficulty": meta["difficulty"] or 0.5,
        "question_type": question_data["question_type"],
        "import_type": import_type,
        "knowledge_codes": knowledge_codes,
        "comment": "",
        "source": source_name,
    }

    return docx_path, problem_id, info


def _copy_paragraph(src_body, dst_body, para_index: int):
    """Deep-copy 一个段落从 src_body 到 dst_body"""
    paragraphs = src_body.findall(qn('w:p'))
    if para_index >= len(paragraphs):
        return
    src_p = paragraphs[para_index]
    # 找到插入位置（dst_body 的最后一个元素之后）
    dst_body.append(deepcopy(src_p))


def _copy_images_to_docx(dst_path: Path, src_path: Path):
    """将原始 docx 中被引用的图片复制到新 docx 中"""
    from lxml import etree
    import shutil

    # 读取目标 docx 中引用的所有图片 rId
    with zipfile.ZipFile(str(dst_path), 'r') as dst_zip:
        doc_xml = dst_zip.read('word/document.xml')
        try:
            dst_rels_xml = dst_zip.read('word/_rels/document.xml.rels')
        except KeyError:
            dst_rels_xml = None

    # 从目标 document.xml 中提取所有 a:blip 的 r:embed
    ns_a = 'http://schemas.openxmlformats.org/drawingml/2006/main'
    ns_r = 'http://schemas.openxmlformats.org/officeDocument/2006/relationships'
    ns_rel = 'http://schemas.openxmlformats.org/package/2006/relationships'

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

    if not needed_embeds:
        return

    # 从原始 docx 中获取图片数据
    with zipfile.ZipFile(str(src_path), 'r') as src_zip:
        # 读取原始关系文件
        src_rels_xml = src_zip.read('word/_rels/document.xml.rels')
        src_rels_tree = etree.fromstring(src_rels_xml)

        # 找到每个 embed 对应的 target
        embed_targets = {}
        for rel in src_rels_tree:
            r_id = rel.get('Id')
            target = rel.get('Target', '')
            if r_id in needed_embeds:
                # target like "media/image1.png"
                embed_targets[r_id] = target

        # 读写目标 docx，添加图片和关系
        tmp_path = dst_path.with_suffix('.tmp.docx')

        with zipfile.ZipFile(str(dst_path), 'r') as dst_zip:
            with zipfile.ZipFile(str(tmp_path), 'w', zipfile.ZIP_DEFLATED) as out_zip:
                # 复制所有现有条目（跳过将被更新的文件）
                skip = {'word/document.xml', 'word/_rels/document.xml.rels', '[Content_Types].xml'}
                for item in dst_zip.infolist():
                    if item.filename in skip:
                        continue
                    data = dst_zip.read(item.filename)
                    out_zip.writestr(item, data)

                # 添加图片文件并记录扩展名
                next_rid = _get_max_rid(dst_zip) + 1
                new_rels = []
                content_types_additions = []

                for embed, target in embed_targets.items():
                    try:
                        img_data = src_zip.read(f'word/{target}')
                        out_zip.writestr(f'word/{target}', img_data)
                        new_rels.append((f'rId{next_rid}', target, embed))
                        next_rid += 1
                        # 记录 content type
                        ext = target.rsplit('.', 1)[-1].lower()
                        mime = {'png': 'image/png', 'jpg': 'image/jpeg', 'jpeg': 'image/jpeg',
                                'gif': 'image/gif', 'bmp': 'image/bmp'}.get(ext, 'image/png')
                        content_types_additions.append((f'/word/{target}', mime))
                    except KeyError:
                        pass

                # 更新关系文件
                if new_rels and dst_rels_xml is not None:
                    rels_tree = etree.fromstring(dst_rels_xml)
                    for new_rid, target, old_rid in new_rels:
                        rel_elem = etree.SubElement(rels_tree, f'{{{ns_rel}}}Relationship')
                        rel_elem.set('Id', new_rid)
                        rel_elem.set('Type', 'http://schemas.openxmlformats.org/officeDocument/2006/relationships/image')
                        rel_elem.set('Target', target)
                        for blip in tree.findall(f'.//{{{ns_a}}}blip'):
                            if blip.get(f'{{{ns_r}}}embed') == old_rid:
                                blip.set(f'{{{ns_r}}}embed', new_rid)
                    rels_out = etree.tostring(rels_tree, xml_declaration=True, encoding='UTF-8', standalone=True)
                    doc_out = etree.tostring(tree, xml_declaration=True, encoding='UTF-8', standalone=True)
                else:
                    rels_out = dst_rels_xml if dst_rels_xml is not None else b''
                    doc_out = doc_xml

                out_zip.writestr('word/_rels/document.xml.rels', rels_out)
                out_zip.writestr('word/document.xml', doc_out)

                # 更新 [Content_Types].xml
                ct_xml = dst_zip.read('[Content_Types].xml')
                ct_tree = etree.fromstring(ct_xml)
                ct_ns = 'http://schemas.openxmlformats.org/package/2006/content-types'
                for partname, mime in content_types_additions:
                    # Check if override already exists
                    exists = False
                    for ov in ct_tree.findall(f'{{{ct_ns}}}Override'):
                        if ov.get('PartName') == partname:
                            exists = True
                            break
                    if not exists:
                        ov = etree.SubElement(ct_tree, f'{{{ct_ns}}}Override')
                        ov.set('PartName', partname)
                        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(dst_path))


def _get_max_rid(zf: zipfile.ZipFile) -> int:
    """获取 ZIP 中关系文件的最大 rId 编号"""
    try:
        rels_xml = zf.read('word/_rels/document.xml.rels')
        tree = etree.fromstring(rels_xml)
        max_id = 0
        for rel in tree:
            r_id = rel.get('Id', '')
            if r_id.startswith('rId'):
                try:
                    num = int(r_id[3:])
                    max_id = max(max_id, num)
                except ValueError:
                    pass
        return max_id
    except Exception:
        return 100  # fallback


# Import for etree
from lxml import etree


def _set_custom_props(doc: Document, props: dict):
    """设置 docx 自定义属性"""
    for key, value in props.items():
        try:
            doc.custom_document_properties[key] = value
        except Exception:
            pass


def _get_chapters(knowledge_codes: list[str], knowledge_map: dict) -> list[str]:
    """从知识点编码反推章节名"""
    from app.services.knowledge import get_chapter_from_code
    chapters = set()
    for code in knowledge_codes:
        ch = get_chapter_from_code(code, knowledge_map)
        if ch:
            chapters.add(ch)
    return sorted(chapters)


def read_custom_props(docx_path: str | Path) -> dict:
    """读取 docx 自定义属性"""
    doc = Document(str(docx_path))
    props = {}
    try:
        for name in doc.custom_document_properties:
            props[name] = doc.custom_document_properties[name]
    except Exception:
        pass
    if not props:
        try:
            from lxml import etree
            custom_xml = doc.part.package.parts
            for part in custom_xml:
                if part.partname.endswith('/custom.xml'):
                    root = etree.fromstring(part.blob)
                    for prop in root:
                        tag = prop.tag.split('}')[-1] if '}' in prop.tag else prop.tag
                        value = prop.text or ""
                        try:
                            if value:
                                props[tag] = float(value)
                        except ValueError:
                            props[tag] = value
        except Exception:
            pass
    return props


def extract_plain_text(docx_path: str | Path) -> str:
    """提取 docx 中所有段落的纯文本（用于 FTS 索引）"""
    doc = Document(str(docx_path))
    texts = []
    for para in doc.paragraphs:
        text = para.text.strip()
        if text:
            texts.append(text)
    return "\n".join(texts)


def compute_id_from_docx(file_path: str | Path) -> str:
    """从 docx 计算题目 ID"""
    import zipfile as zf
    try:
        with zf.ZipFile(str(file_path), 'r') as z:
            if 'word/document.xml' in z.namelist():
                return hashlib.sha256(z.read('word/document.xml')).hexdigest()[:8]
    except Exception:
        pass
    with open(file_path, 'rb') as f:
        return hashlib.sha256(f.read()).hexdigest()[:8]


def compute_checksum(file_path: str | Path) -> str:
    try:
        with zipfile.ZipFile(str(file_path), 'r') as z:
            if 'word/document.xml' in z.namelist():
                return hashlib.sha256(z.read('word/document.xml')).hexdigest()
    except Exception:
        pass
    with open(file_path, 'rb') as f:
        return hashlib.sha256(f.read()).hexdigest()