ax-backend/app/services/export_service.py

"""export_service.py — 将文档 Markdown 内容转换为 .doc 文件并返回永久下载链接。"""

import base64
import io
import json
import time
import unicodedata
from datetime import date
from pathlib import Path
from typing import Optional

import mistune
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, RGBColor
from lxml import etree

from app.config import settings
from app.core.exceptions import ExportError


# ------------------------------------------------------------------ #
# 样式文件加载
# ------------------------------------------------------------------ #

def load_style_file(style_id: Optional[str] = None) -> dict:
    """加载样式 JSON；style_id=None 时使用默认样式文件。"""
    if style_id is not None:
        # 阶段 1 占位
        raise ExportError(f"样式 ID 暂不支持: {style_id}（阶段 1 功能）")

    path = Path(settings.default_style_file)
    if not path.exists():
        raise ExportError(f"默认样式文件不存在: {path}")

    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError) as exc:
        raise ExportError(f"样式文件解析失败: {exc}") from exc


def build_style_map(style_data: dict) -> dict[str, dict]:
    """将样式列表转为双键映射（style_id 和 name 均可命中）。"""
    mapping: dict[str, dict] = {}
    for s in style_data.get("styles", []):
        if s.get("style_id"):
            mapping[s["style_id"]] = s
        if s.get("name"):
            mapping[s["name"]] = s
    return mapping


# ------------------------------------------------------------------ #
# JSON ↔ lxml 互转
# ------------------------------------------------------------------ #

def dict_to_element(d: dict) -> etree._Element:
    elem = etree.Element(d["@tag"], attrib=dict(d.get("@attrib", {})))
    if d.get("#text"):
        elem.text = d["#text"]
    if d.get("#tail"):
        elem.tail = d["#tail"]
    for child_tag, child_val in d.get("@children", {}).items():
        items = child_val if isinstance(child_val, list) else [child_val]
        for item in items:
            if isinstance(item, dict):
                elem.append(dict_to_element(item))
    return elem


def inject_styles_from_json(doc: Document, style_data: dict) -> None:
    """将 JSON 中所有样式的 full_xml_definition upsert 到文档 <w:styles> 节点。"""
    styles_element = doc.styles.element
    for style_entry in style_data.get("styles", []):
        xml_def = style_entry.get("full_xml_definition")
        if not xml_def:
            continue
        try:
            new_elem = dict_to_element(xml_def)
        except Exception:
            continue
        style_id_key = qn("w:styleId")
        new_style_id = new_elem.get(style_id_key)
        if new_style_id:
            existing = styles_element.find(
                f'.//{qn("w:style")}[@{qn("w:styleId")}="{new_style_id}"]'
            )
            if existing is not None:
                styles_element.remove(existing)
        styles_element.append(new_elem)


def inject_numbering_from_json(doc: Document, style_data: dict) -> None:
    """
    将 JSON 中的 numbering 定义注入到文档中。
    这样可以恢复标题的编号格式。
    """
    numbering_def = style_data.get("numbering")
    if not numbering_def:
        return  # 没有编号定义，跳过
    
    try:
        # 将字典转换为 lxml Element
        numbering_elem = dict_to_element(numbering_def)
        
        # 获取文档的 numbering part
        # python-docx 可能没有 numbering part，需要创建
        if doc.part.numbering_part is None:
            # 创建 numbering part
            from docx.opc.constants import CONTENT_TYPE as CT
            from docx.opc.part import XmlPart
            from docx.opc.packuri import PackURI
            
            numbering_part = XmlPart(
                PackURI('/word/numbering.xml'),
                CT.WML_NUMBERING,
                numbering_elem,
                doc.part.package
            )
            doc.part.relate_to(numbering_part, 'http://schemas.openxmlformats.org/officeDocument/2006/relationships/numbering')
        else:
            # 替换现有的 numbering part 内容
            doc.part.numbering_part._element = numbering_elem
    except Exception as e:
        # 编号注入失败，不影响其他功能
        print(f"警告: 编号格式注入失败: {e}")


def _resolve_style_id(style_map: dict, *keys: str) -> Optional[str]:
    for key in keys:
        entry = style_map.get(key)
        if entry and entry.get("style_id"):
            return entry["style_id"]
    return None


# ------------------------------------------------------------------ #
# Markdown → python-docx 渲染器
# ------------------------------------------------------------------ #

class DocxRenderer(mistune.BaseRenderer):
    _HEADING_ALIASES = {i: [f"Heading {i}", f"heading {i}"] for i in range(1, 7)}

    def __init__(self, style_map: dict, style_data: dict) -> None:
        super().__init__()
        self.style_map = style_map
        self.doc = Document()
        inject_styles_from_json(self.doc, style_data)
        inject_numbering_from_json(self.doc, style_data)  # 新增：注入编号格式
        self._normal_id: Optional[str] = _resolve_style_id(style_map, "Normal", "1")
        self.pending_style: Optional[str] = None  # 待应用的样式名
        self.pending_image_style: Optional[dict] = None  # 待应用的图片样式

    def _get_style_by_id(self, style_id: str):
        for style in self.doc.styles:
            if style.style_id == style_id:
                return style
        raise KeyError(style_id)
    
    def _apply_numbering_from_style(self, paragraph):
        """
        从段落的样式中提取编号属性并应用到段落。
        这是必需的，因为 python-docx 不会自动继承样式的编号格式。
        """
        if not paragraph.style:
            return
        
        try:
            # 获取样式的 XML 元素
            style_elem = paragraph.style.element
            
            # 查找样式中的编号定义
            pPr = style_elem.find(qn('w:pPr'))
            if pPr is None:
                return
            
            numPr = pPr.find(qn('w:numPr'))
            if numPr is None:
                return
            
            # 复制编号属性到段落
            para_pPr = paragraph._p.get_or_add_pPr()
            
            # 移除现有的 numPr（如果有）
            existing_numPr = para_pPr.find(qn('w:numPr'))
            if existing_numPr is not None:
                para_pPr.remove(existing_numPr)
            
            # 深度复制样式的 numPr 到段落
            from copy import deepcopy
            new_numPr = deepcopy(numPr)
            para_pPr.append(new_numPr)
        except Exception:
            # 编号应用失败，不影响其他功能
            pass

    @staticmethod
    def _extract_text(children: list) -> str:
        parts: list[str] = []
        for child in children:
            if isinstance(child, dict):
                if child.get("raw"):
                    parts.append(child["raw"])
                if child.get("children"):
                    parts.append(DocxRenderer._extract_text(child["children"]))
        return "".join(parts)

    def heading(self, token: dict, state) -> str:
        level = token["attrs"]["level"]
        text = self._extract_text(token.get("children", []))
        aliases = self._HEADING_ALIASES.get(level, [f"Heading {level}"])
        style_id = _resolve_style_id(self.style_map, *aliases)
        if style_id:
            para = self.doc.add_paragraph(text)
            try:
                para.style = self._get_style_by_id(style_id)
                # 应用样式后，复制编号属性到段落
                self._apply_numbering_from_style(para)
            except KeyError:
                pass
        else:
            self.doc.add_heading(text, level=level)
        return ""

    def paragraph(self, token: dict, state) -> str:
        # 检查是否包含图片
        children = token.get("children", [])
        has_image = any(child.get("type") == "image" for child in children)
        
        if has_image:
            # 如果包含图片，直接调用 image 处理
            for child in children:
                if child.get("type") == "image":
                    self.image(child, state)
            return ""
        
        p = self.doc.add_paragraph()
        
        # 尝试应用待定样式
        style_applied = False
        if self.pending_style:
            style_id = _resolve_style_id(self.style_map, self.pending_style)
            if style_id:
                try:
                    p.style = self._get_style_by_id(style_id)
                    style_applied = True
                except KeyError:
                    pass  # 样式不存在，静默忽略
            self.pending_style = None
        
        # 如果没有应用样式，使用 Normal
        if not style_applied and self._normal_id:
            try:
                p.style = self._get_style_by_id(self._normal_id)
            except Exception:
                pass
        
        self._render_inline_children(p, token.get("children", []))
        return ""
    
    def html(self, token: dict, state) -> str:
        """处理内联 HTML 注释（表格单元格中的样式标记）"""
        raw = token.get("raw", "")
        if "<!-- style:" in raw and "-->" in raw:
            try:
                start = raw.index("<!-- style:") + 11
                end = raw.index("-->", start)
                self.pending_style = raw[start:end].strip()
            except (ValueError, IndexError):
                pass
        return ""
    
    def block_html(self, token: dict, state) -> str:
        """处理块级 HTML（样式注释+文本在同一行）"""
        raw = token.get("raw", "")
        
        # 处理图片样式注释
        if "<!-- img-style:" in raw and "-->" in raw:
            try:
                start = raw.index("{")
                end = raw.rindex("}") + 1
                self.pending_image_style = json.loads(raw[start:end])
            except (ValueError, json.JSONDecodeError):
                pass
            return ""
        
        # 处理文本样式注释
        if "<!-- style:" in raw and "-->" in raw:
            try:
                # 提取样式名
                style_start = raw.index("<!-- style:") + 11
                style_end = raw.index("-->", style_start)
                style_name = raw[style_start:style_end].strip()
                
                # 提取文本（注释后面的内容）
                text_start = style_end + 3  # "-->".length = 3
                text = raw[text_start:].strip()
                
                # 创建段落并应用样式
                p = self.doc.add_paragraph(text)
                style_id = _resolve_style_id(self.style_map, style_name)
                if style_id:
                    try:
                        p.style = self._get_style_by_id(style_id)
                        # 应用样式后，复制编号属性到段落
                        self._apply_numbering_from_style(p)
                    except KeyError:
                        pass  # 样式不存在，使用默认
            except (ValueError, IndexError):
                # 解析失败，当作普通 HTML 处理（忽略）
                pass
        return ""

    def blank_line(self, token: dict, state) -> str:
        return ""

    def image(self, token: dict, state) -> str:
        """处理图片 token（支持 Base64 Data URL）"""
        url = token['attrs']['url']
        alt = token['attrs'].get('alt', '图片')
        
        # 只处理 Data URL
        if not url.startswith('data:'):
            return ""
        
        try:
            # 解析 data:image/png;base64,xxxxx
            if ',' not in url:
                return ""
            
            header, b64_data = url.split(',', 1)
            image_bytes = base64.b64decode(b64_data)
            
            # 获取样式（来自前面的 HTML 注释）
            style = self.pending_image_style or {}
            self.pending_image_style = None
            
            # 创建段落并设置对齐
            paragraph = self.doc.add_paragraph()
            
            # 应用段落样式：优先使用保存的样式，否则使用 Normal
            para_style = style.get('para_style', 'Normal')
            style_id = _resolve_style_id(self.style_map, para_style)
            
            if style_id:
                try:
                    paragraph.style = self._get_style_by_id(style_id)
                except KeyError:
                    # 如果样式不存在，回退到 Normal
                    if self._normal_id:
                        try:
                            paragraph.style = self._get_style_by_id(self._normal_id)
                        except Exception:
                            pass
            elif self._normal_id:
                # 如果没有找到样式 ID，使用 Normal
                try:
                    paragraph.style = self._get_style_by_id(self._normal_id)
                except Exception:
                    pass
            
            align = style.get('align', 'left')
            if align == 'center':
                paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
            elif align == 'right':
                paragraph.alignment = WD_ALIGN_PARAGRAPH.RIGHT
            
            # 插入图片
            run = paragraph.add_run()
            width = style.get('width', 10.0)
            height = style.get('height', 7.0)
            unit = style.get('unit', 'cm')
            
            # 转换为磅（Word内部单位：1厘米 = 28.35磅，1英寸 = 72磅）
            if unit == 'cm':
                width_pt = width * 28.35
                height_pt = height * 28.35
            else:  # inches
                width_pt = width * 72
                height_pt = height * 72
            
            run.add_picture(
                io.BytesIO(image_bytes),
                width=Pt(width_pt),
                height=Pt(height_pt)
            )
            
        except Exception as e:
            # 失败时添加占位文本
            p = self.doc.add_paragraph(f"[图片加载失败: {alt}]")
            if self._normal_id:
                try:
                    p.style = self._get_style_by_id(self._normal_id)
                except Exception:
                    pass
            p.runs[0].font.color.rgb = RGBColor(255, 0, 0)
        
        return ""

    def thematic_break(self, token: dict, state) -> str:
        p = self.doc.add_paragraph()
        pPr = p._p.get_or_add_pPr()
        pBdr = OxmlElement("w:pBdr")
        bottom = OxmlElement("w:bottom")
        bottom.set(qn("w:val"), "single")
        bottom.set(qn("w:sz"), "6")
        bottom.set(qn("w:space"), "1")
        bottom.set(qn("w:color"), "auto")
        pBdr.append(bottom)
        pPr.append(pBdr)
        return ""

    def block_quote(self, token: dict, state) -> str:
        for child in token.get("children", []):
            text = self._extract_text(child.get("children", []))
            quote_id = _resolve_style_id(self.style_map, "Quote", "Quote Char")
            p = self.doc.add_paragraph(text)
            if quote_id:
                try:
                    p.style = self._get_style_by_id(quote_id)
                except Exception:
                    p.style = "Quote"
            else:
                p.style = "Quote"
        return ""

    def block_code(self, token: dict, state) -> str:
        p = self.doc.add_paragraph(style="No Spacing")
        run = p.add_run(token.get("raw", ""))
        run.font.name = "Courier New"
        run.font.size = Pt(10)
        run.font.color.rgb = RGBColor(0x33, 0x33, 0x33)
        return ""

    def list(self, token: dict, state) -> str:
        ordered = token["attrs"].get("ordered", False)
        depth = token["attrs"].get("depth", 0)  # mistune的depth从0开始
        self._render_list_items(token.get("children", []), ordered, depth + 1)  # 转换为从1开始
        return ""

    def _render_list_items(self, items: list, ordered: bool, depth: int) -> None:
        for item in items:
            for child in item.get("children", []):
                if child["type"] == "list":
                    self._render_list_items(
                        child.get("children", []),
                        child["attrs"].get("ordered", False),
                        depth + 1,
                    )
                else:
                    # 使用内联格式渲染（保留粗体、斜体等）
                    # 注意：depth=1 用 "List Bullet"，depth=2用"List Bullet 2"
                    if ordered:
                        if depth == 1:
                            style = "List Number"
                        elif depth == 2:
                            style = "List Number 2"
                        else:
                            style = "List Number 3"
                    else:
                        if depth == 1:
                            style = "List Bullet"
                        elif depth == 2:
                            style = "List Bullet 2"
                        else:
                            style = "List Bullet 3"
                    
                    para = self.doc.add_paragraph(style=style)
                    self._render_inline_children(para, child.get("children", [child]))

    def table(self, token: dict, state) -> str:
        children = token.get("children", [])
        head_token = next((c for c in children if c["type"] == "table_head"), None)
        body_token = next((c for c in children if c["type"] == "table_body"), None)

        head_cells = head_token.get("children", []) if head_token else []
        cols = len(head_cells)
        if cols == 0:
            return ""

        body_rows: list[list[dict]] = []
        if body_token:
            for row in body_token.get("children", []):
                if row["type"] == "table_row":
                    body_rows.append(row.get("children", []))

        tbl = self.doc.add_table(rows=1 + len(body_rows), cols=cols)
        tbl.style = "Table Grid"

        # 表头行（保留内联格式，不强制加粗）
        for c, cell_token in enumerate(head_cells):
            cell = tbl.rows[0].cells[c]
            # 清空默认段落
            cell.text = ""
            para = cell.paragraphs[0]
            # 渲染内联内容（样式由 _render_inline_children 处理）
            self._render_inline_children(para, cell_token.get("children", []))

        # 数据行（保留内联格式）
        for r, row_cells in enumerate(body_rows):
            for c, cell_token in enumerate(row_cells):
                if c >= cols:
                    break
                cell = tbl.rows[r + 1].cells[c]
                cell.text = ""
                para = cell.paragraphs[0]
                # 渲染内联内容（样式注释会在 _render_inline_children 中处理）
                self._render_inline_children(para, cell_token.get("children", []))

        return ""

    def _render_inline_children(self, paragraph, children: list) -> None:
        """渲染内联子元素，处理粗体、斜体等格式"""
        for child in children:
            ctype = child.get("type", "")
            raw = child.get("raw", "")
            
            if ctype == "inline_html":
                # 处理图片样式注释
                if "<!-- img-style:" in raw and "-->" in raw:
                    try:
                        start = raw.index("{")
                        end = raw.rindex("}") + 1
                        self.pending_image_style = json.loads(raw[start:end])
                    except (ValueError, json.JSONDecodeError):
                        pass
                    # 注释不输出
                    continue
                
                # 处理文本样式注释
                if "<!-- style:" in raw and "-->" in raw:
                    try:
                        start = raw.index("<!-- style:") + 11
                        end = raw.index("-->", start)
                        self.pending_style = raw[start:end].strip()
                    except (ValueError, IndexError):
                        pass
                # 注释不输出
                continue
            
            elif ctype == "text":
                # 应用待定样式（来自前一个 inline_html）
                if self.pending_style:
                    style_id = _resolve_style_id(self.style_map, self.pending_style)
                    if style_id:
                        try:
                            paragraph.style = self._get_style_by_id(style_id)
                        except KeyError:
                            pass
                    self.pending_style = None
                paragraph.add_run(raw)
            
            elif ctype == "strong":
                paragraph.add_run(self._extract_text(child.get("children", []))).bold = True
            elif ctype == "emphasis":
                paragraph.add_run(self._extract_text(child.get("children", []))).italic = True
            elif ctype == "strikethrough":
                paragraph.add_run(self._extract_text(child.get("children", []))).font.strike = True
            elif ctype == "codespan":
                run = paragraph.add_run(raw)
                run.font.name = "Courier New"
                run.font.size = Pt(10)
            elif ctype == "linebreak":
                paragraph.add_run().add_break()
            elif ctype == "softlinebreak":
                paragraph.add_run(" ")
            elif ctype == "image":
                # 处理内联图片
                # 注意：这里的图片是在段落中的，需要特殊处理
                # 我们需要跳过这个段落，让 image() 方法来处理
                pass
            else:
                sub = child.get("children")
                if sub:
                    self._render_inline_children(paragraph, sub)
                elif raw:
                    paragraph.add_run(raw)

    def render_token(self, token: dict, state) -> str:
        func = getattr(self, token["type"], None)
        if func:
            return func(token, state)
        for child in token.get("children", []):
            self.render_token(child, state)
        return ""

    def __call__(self, tokens: list, state) -> str:
        for token in tokens:
            self.render_token(token, state)
        return ""


# ------------------------------------------------------------------ #
# 公共工具
# ------------------------------------------------------------------ #

def markdown_to_docx_bytes(content: str, style_map: dict, style_data: dict) -> bytes:
    renderer = DocxRenderer(style_map=style_map, style_data=style_data)
    md = mistune.create_markdown(
        renderer=renderer,
        plugins=["table", "strikethrough", "url"],
    )
    md(content)
    buf = io.BytesIO()
    renderer.doc.save(buf)
    return buf.getvalue()


def _safe_filename(name: str) -> str:
    name = unicodedata.normalize("NFKC", name)
    for ch in r'\/:*?"<>|':
        name = name.replace(ch, "_")
    return name.strip() or "document"


def _make_filename(content: str) -> str:
    """取内容第一行文本 + 毫秒时间戳，生成文件名（不含扩展名）。"""
    first_line = content.lstrip().split("\n")[0].lstrip("#").strip()
    safe = _safe_filename(first_line) if first_line else "document"
    ts = int(time.time() * 1000)
    return f"{safe}_{ts}"