ax-backend-v1/app/services/export_service.py

"""export_service.py — 将 Blocks 转换为 .doc 文件并返回永久下载链接"""

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

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:
    """递归将字典转为 lxml 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)
    
    # 注入编号格式定义
    inject_numbering_from_json(doc, style_data)


def inject_numbering_from_json(doc: Document, style_data: dict) -> None:
    """将 JSON 中的编号格式定义注入到文档
    
    由于 python-docx 对 numbering part 的支持有限，
    我们需要在保存后通过修改 ZIP 文件来注入编号格式。
    这个函数主要是为了记录编号定义，实际注入在 blocks_to_docx_bytes 中完成。
    """
    # 暂时不在这里注入，而是在生成文档后通过 ZIP 修改
    pass


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


def _apply_paragraph_style(para, style: dict):
    """应用段落级样式（对齐方式）
    
    Args:
        para: python-docx 段落对象
        style: 样式字典
    """
    # 对齐方式
    align = style.get('align')
    if align == 'center':
        para.alignment = WD_ALIGN_PARAGRAPH.CENTER
    elif align == 'right':
        para.alignment = WD_ALIGN_PARAGRAPH.RIGHT
    elif align == 'justify':
        para.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY
    elif align == 'left':
        para.alignment = WD_ALIGN_PARAGRAPH.LEFT


def _apply_run_style(run, style: dict):
    """应用 run 级样式（字符级格式）
    
    Args:
        run: python-docx run 对象
        style: 样式字典
    """
    # 粗体
    if style.get('bold'):
        run.bold = True
    
    # 斜体
    if style.get('italic'):
        run.italic = True
    
    # 下划线
    if style.get('underline'):
        run.underline = True
    
    # 删除线
    if style.get('strike'):
        run.font.strike = True
    
    # 颜色
    if style.get('color'):
        try:
            # 移除可能的 # 前缀
            color = style['color'].lstrip('#')
            if len(color) == 6:
                run.font.color.rgb = RGBColor(
                    int(color[0:2], 16),
                    int(color[2:4], 16),
                    int(color[4:6], 16)
                )
        except (ValueError, AttributeError):
            pass
    
    # 字体名称（支持中文字体 eastAsia）
    if style.get('font_name'):
        font_name = style['font_name']
        run.font.name = font_name
        
        # 对于中文字体，需要设置 eastAsia 属性
        try:
            r = run._element
            rPr = r.get_or_add_rPr()
            rFonts = rPr.get_or_add_rFonts()
            rFonts.set(qn('w:eastAsia'), font_name)
        except Exception:
            pass
    
    # 字号
    if style.get('font_size'):
        run.font.size = Pt(style['font_size'])


# ------------------------------------------------------------------ #
# Blocks → Word 转换
# ------------------------------------------------------------------ #

def blocks_to_docx_bytes(blocks: list[dict], style_map: dict, style_data: dict) -> bytes:
    """将 Blocks 列表转换为 Word 文档字节流
    
    Args:
        blocks: Block 列表
        style_map: 样式映射
        style_data: 样式数据
        
    Returns:
        Word 文档字节流
    """
    doc = Document()
    inject_styles_from_json(doc, style_data)
    
    # 检查 blocks 中是否有常用的字体和字号，用于修改 Normal 样式
    # 这样可以确保空行在 Word 中显示正确的字体
    _update_normal_style_if_needed(doc, blocks)
    
    for block in blocks:
        block_type = block['type']
        
        if block_type == 'heading':
            _render_heading_block(doc, block, style_map)
        elif block_type == 'paragraph':
            _render_paragraph_block(doc, block, style_map)
        elif block_type == 'table':
            _render_table_block(doc, block, style_map)
        elif block_type == 'image':
            _render_image_block(doc, block)
    
    # 先保存到临时缓冲区
    buf = io.BytesIO()
    doc.save(buf)
    
    # 通过 ZIP 操作注入编号格式
    docx_bytes = _inject_numbering_via_zip(buf.getvalue(), style_data)
    
    return docx_bytes


def _update_normal_style_if_needed(doc: Document, blocks: list[dict]):
    """更新 Normal 样式以匹配 blocks 中最常用的字体
    
    这样可以确保空行在 Word 中显示正确的字体和字号
    """
    # 统计段落中最常用的字体和字号
    font_counts = {}
    size_counts = {}
    
    for block in blocks:
        if block['type'] == 'paragraph':
            style = block.get('style', {})
            font_name = style.get('font_name')
            font_size = style.get('font_size')
            
            if font_name:
                font_counts[font_name] = font_counts.get(font_name, 0) + 1
            if font_size:
                size_counts[font_size] = size_counts.get(font_size, 0) + 1
    
    # 找到最常用的字体和字号
    most_common_font = max(font_counts.items(), key=lambda x: x[1])[0] if font_counts else None
    most_common_size = max(size_counts.items(), key=lambda x: x[1])[0] if size_counts else None
    
    # 如果找到了常用字体或字号，更新 Normal 样式
    if most_common_font or most_common_size:
        try:
            normal_style = doc.styles['Normal']
            
            if most_common_font:
                # 修改 Normal 样式的字体
                style_element = normal_style.element
                rPr = style_element.find(qn('w:rPr'))
                if rPr is None:
                    rPr = OxmlElement('w:rPr')
                    # 插入到第一个子元素之前
                    if len(style_element):
                        style_element.insert(0, rPr)
                    else:
                        style_element.append(rPr)
                
                rFonts = rPr.find(qn('w:rFonts'))
                if rFonts is None:
                    rFonts = OxmlElement('w:rFonts')
                    rPr.append(rFonts)
                
                # 设置所有字体属性
                rFonts.set(qn('w:ascii'), most_common_font)
                rFonts.set(qn('w:hAnsi'), most_common_font)
                rFonts.set(qn('w:eastAsia'), most_common_font)
            
            if most_common_size:
                # 修改 Normal 样式的字号
                style_element = normal_style.element
                rPr = style_element.find(qn('w:rPr'))
                if rPr is None:
                    rPr = OxmlElement('w:rPr')
                    if len(style_element):
                        style_element.insert(0, rPr)
                    else:
                        style_element.append(rPr)
                
                # 删除旧的字号元素
                old_sz = rPr.find(qn('w:sz'))
                if old_sz is not None:
                    rPr.remove(old_sz)
                old_szCs = rPr.find(qn('w:szCs'))
                if old_szCs is not None:
                    rPr.remove(old_szCs)
                
                # 添加新的字号元素
                sz = OxmlElement('w:sz')
                sz.set(qn('w:val'), str(int(most_common_size * 2)))  # Word 使用半磅
                rPr.append(sz)
                
                szCs = OxmlElement('w:szCs')
                szCs.set(qn('w:val'), str(int(most_common_size * 2)))
                rPr.append(szCs)
                
        except Exception:
            # 如果修改样式失败，继续（不影响文档生成）
            pass


def _inject_numbering_via_zip(docx_bytes: bytes, style_data: dict) -> bytes:
    """通过 ZIP 操作注入编号格式到 Word 文档
    
    Args:
        docx_bytes: 原始 Word 文档字节流
        style_data: 样式数据（包含 numbering 定义）
        
    Returns:
        注入编号格式后的 Word 文档字节流
    """
    numbering_def = style_data.get("numbering")
    if not numbering_def:
        # 没有编号定义，直接返回原文档
        return docx_bytes
    
    try:
        from zipfile import ZipFile
        from lxml import etree
        
        # 读取原文档
        input_buf = io.BytesIO(docx_bytes)
        output_buf = io.BytesIO()
        
        with ZipFile(input_buf, 'r') as zip_read:
            with ZipFile(output_buf, 'w') as zip_write:
                # 复制所有文件
                for item in zip_read.infolist():
                    data = zip_read.read(item.filename)
                    
                    # 跳过 numbering.xml，我们会重新写入
                    if item.filename == 'word/numbering.xml':
                        continue
                    
                    zip_write.writestr(item, data)
                
                # 将 JSON 格式的编号定义转换为 XML
                numbering_element = dict_to_element(numbering_def)
                numbering_xml = etree.tostring(
                    numbering_element,
                    encoding='UTF-8',
                    xml_declaration=True,
                    standalone=True
                )
                
                # 写入 numbering.xml
                zip_write.writestr('word/numbering.xml', numbering_xml)
                
                # 确保 _rels/document.xml.rels 中有 numbering 的关系
                # 读取 document.xml.rels
                try:
                    rels_data = zip_read.read('word/_rels/document.xml.rels')
                    rels_root = etree.fromstring(rels_data)
                    
                    # 检查是否已有 numbering 关系
                    ns = {'r': 'http://schemas.openxmlformats.org/package/2006/relationships'}
                    numbering_rels = rels_root.xpath(
                        '//r:Relationship[@Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/numbering"]',
                        namespaces=ns
                    )
                    
                    if not numbering_rels:
                        # 添加 numbering 关系
                        from docx.oxml.ns import qn
                        rel_elem = etree.SubElement(rels_root, qn('r:Relationship'))
                        # 找到最大的 rId
                        existing_ids = [int(r.get('Id')[3:]) for r in rels_root.findall(qn('r:Relationship')) if r.get('Id', '').startswith('rId')]
                        next_id = max(existing_ids) + 1 if existing_ids else 1
                        rel_elem.set('Id', f'rId{next_id}')
                        rel_elem.set('Type', 'http://schemas.openxmlformats.org/officeDocument/2006/relationships/numbering')
                        rel_elem.set('Target', 'numbering.xml')
                        
                        # 写回 rels 文件
                        rels_xml = etree.tostring(rels_root, encoding='UTF-8', xml_declaration=True)
                        zip_write.writestr('word/_rels/document.xml.rels', rels_xml)
                    
                except KeyError:
                    # 如果没有 rels 文件，忽略
                    pass
        
        return output_buf.getvalue()
        
    except Exception as e:
        # 如果注入失败，返回原文档
        print(f"警告: 通过 ZIP 注入编号格式失败: {e}")
        return docx_bytes


def _render_heading_block(doc: Document, block: dict, style_map: dict):
    """渲染标题块（支持编号格式和自定义样式）"""
    level = block['level']
    content = block['content']
    style_name = block.get('word_style', f'Heading {level}')
    block_style = block.get('style', {})
    
    # 创建段落
    para = doc.add_paragraph()
    
    # 应用 Word 样式
    style_id = _resolve_style_id(style_map, style_name, f'Heading {level}')
    if style_id:
        try:
            para.style = doc.styles[style_id]
        except KeyError:
            para.style = f'Heading {level}'
    else:
        para.style = f'Heading {level}'
    
    # 应用 Block 级自定义样式（段落级）
    _apply_paragraph_style(para, block_style)
    
    # 渲染内容
    if isinstance(content, list):
        # 富文本：应用 run 级样式
        _render_rich_text(para, content, block_style)
    else:
        # 纯文本：应用 block 级样式到 run
        run = para.add_run(str(content))
        _apply_run_style(run, block_style)
    
    # 尝试应用编号格式（如果样式中包含编号定义）
    try:
        # 检查样式是否有编号定义
        style_element = para.style.element
        if style_element is not None:
            from docx.oxml.ns import qn
            # 查找样式中的编号属性
            pPr = style_element.find(qn('w:pPr'))
            if pPr is not None:
                numPr = pPr.find(qn('w:numPr'))
                if numPr is not None:
                    # 样式中有编号定义，复制到段落
                    para_pPr = para._element.get_or_add_pPr()
                    # 移除可能存在的旧编号属性
                    old_numPr = para_pPr.find(qn('w:numPr'))
                    if old_numPr is not None:
                        para_pPr.remove(old_numPr)
                    # 复制编号属性
                    import copy
                    para_pPr.append(copy.deepcopy(numPr))
    except Exception as e:
        # 如果应用编号失败，继续（标题仍然会显示，只是没有编号）
        pass


def _render_paragraph_block(doc: Document, block: dict, style_map: dict):
    """渲染段落块（支持富文本和自定义样式）"""
    content = block['content']
    style_name = block.get('word_style', 'Normal')
    block_style = block.get('style', {})
    
    para = doc.add_paragraph()
    
    # 应用 Word 样式
    style_id = _resolve_style_id(style_map, style_name, 'Normal')
    if style_id:
        try:
            para.style = doc.styles[style_id]
        except KeyError:
            para.style = 'Normal'
    else:
        para.style = 'Normal'
    
    # 应用 Block 级自定义样式（段落级）
    _apply_paragraph_style(para, block_style)
    
    # 渲染内容
    if isinstance(content, list):
        # 富文本：应用 run 级样式
        _render_rich_text(para, content, block_style)
    else:
        # 纯文本或空内容
        # 即使是空内容，如果有 block 级样式（字体、字号等），也需要添加空 run 来保存样式
        # 这样当用户在 Word 中输入文本时，会自动应用这些样式
        run = para.add_run(str(content) if content else '')
        _apply_run_style(run, block_style)


def _render_rich_text(para, segments: list, block_style: dict = None):
    """渲染富文本格式
    
    Args:
        para: python-docx 段落对象
        segments: 富文本片段列表，每个片段包含 text 和 style
        block_style: Block 级样式，作为默认样式（可选）
    """
    for seg in segments:
        text = seg.get('text', '')
        seg_style = seg.get('style', {})
        
        run = para.add_run(text)
        
        # 合并样式：block_style 作为默认，seg_style 覆盖
        merged_style = {}
        if block_style:
            merged_style.update(block_style)
        merged_style.update(seg_style)
        
        # 应用合并后的样式
        _apply_run_style(run, merged_style)


def _render_table_block(doc: Document, block: dict, style_map: dict):
    """渲染表格块（支持合并单元格、列宽、行高等）"""
    table_data = block['content']
    if isinstance(table_data, str):
        try:
            table_data = json.loads(table_data)
        except json.JSONDecodeError:
            return
    
    rows = table_data.get('rows', [])
    if not rows:
        return
    
    # 使用 col_widths 确定真实列数（而不是第一行的单元格数）
    col_widths = table_data.get('col_widths', [])
    if col_widths:
        num_cols = len(col_widths)
    else:
        # 回退：扫描所有行，找到最大的列索引
        num_cols = 0
        for row_data in rows:
            col_index = 0
            for cell_data in row_data.get('cells', []):
                colspan = cell_data.get('colspan', 1)
                col_index += colspan
            num_cols = max(num_cols, col_index)
        
        if num_cols == 0:
            return
    
    num_rows = len(rows)
    
    # 创建表格
    table = doc.add_table(rows=num_rows, cols=num_cols)
    
    # 应用表格样式
    table_style = block.get('word_style', 'Table Grid')
    try:
        table.style = table_style
    except KeyError:
        table.style = 'Table Grid'
    
    # 设置列宽
    if col_widths:
        for col_idx, width in enumerate(col_widths):
            if col_idx < len(table.columns):
                table.columns[col_idx].width = Pt(width)
    
    # 填充内容并处理合并单元格
    merge_map = {}  # {(row, col): (end_row, end_col)} 记录合并区域
    
    for r_idx, row_data in enumerate(rows):
        # 设置行高
        row_height = row_data.get('height')
        if row_height:
            table.rows[r_idx].height = Pt(row_height)
        
        cells_data = row_data.get('cells', [])
        col_offset = 0  # 当前列偏移（考虑 colspan）
        
        for cell_data in cells_data:
            # 跳过被合并的单元格（rowspan=0 表示这个单元格被上面的单元格合并了）
            rowspan = cell_data.get('rowspan', 1)
            if rowspan == 0:
                col_offset += 1
                continue
            
            colspan = cell_data.get('colspan', 1)
            
            # 确保不越界
            if col_offset >= num_cols:
                break
            
            # 获取起始单元格
            start_cell = table.rows[r_idx].cells[col_offset]
            
            # 处理合并单元格
            if colspan > 1 or rowspan > 1:
                # 计算结束位置
                end_col = min(col_offset + colspan - 1, num_cols - 1)
                end_row = min(r_idx + rowspan - 1, num_rows - 1)
                
                # 合并单元格
                if end_col > col_offset or end_row > r_idx:
                    try:
                        end_cell = table.rows[end_row].cells[end_col]
                        start_cell.merge(end_cell)
                        merge_map[(r_idx, col_offset)] = (end_row, end_col)
                    except Exception:
                        pass  # 合并失败，继续
            
            # 设置单元格宽度（如果有）
            cell_width = cell_data.get('width')
            if cell_width:
                try:
                    start_cell.width = Pt(cell_width)
                except Exception:
                    pass
            
            # 填充单元格内容
            cell_text = cell_data.get('text', '')
            cell_style = cell_data.get('style', {})
            
            # 清空默认段落
            start_cell.text = ''
            para = start_cell.paragraphs[0]
            
            # 应用单元格段落级样式（对齐）
            _apply_paragraph_style(para, cell_style)
            
            # 渲染单元格内容（支持富文本）
            if isinstance(cell_text, list):
                # 富文本格式
                _render_rich_text(para, cell_text, cell_style)
            else:
                # 纯文本格式
                run = para.add_run(str(cell_text))
                # 应用单元格 run 级样式
                _apply_run_style(run, cell_style)
            
            # 更新列偏移
            col_offset += colspan


def _render_image_block(doc: Document, block: dict):
    """渲染图片块（支持 Base64 Data URL）"""
    content = block['content']
    style = block.get('style', {})
    
    # 只处理 Data URL
    if not isinstance(content, str) or not content.startswith('data:'):
        return
    
    try:
        # 解析 data:image/png;base64,xxxxx
        if ',' not in content:
            return
        
        header, b64_data = content.split(',', 1)
        image_bytes = base64.b64decode(b64_data)
        
        # 创建段落并设置对齐
        paragraph = doc.add_paragraph()
        
        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 = doc.add_paragraph(f"[图片加载失败]")
        p.runs[0].font.color.rgb = RGBColor(255, 0, 0)


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

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(blocks: list[dict]) -> str:
    """从 blocks 中提取第一个标题或段落作为文件名 + 时间戳
    
    Args:
        blocks: Block 列表
        
    Returns:
        文件名（不含扩展名）
    """
    # 查找第一个标题或段落
    first_text = ""
    for block in blocks:
        if block['type'] in ('heading', 'paragraph'):
            content = block['content']
            if isinstance(content, list):
                # 富文本：拼接所有片段
                first_text = "".join(seg.get("text", "") for seg in content)
            else:
                first_text = str(content)
            
            if first_text.strip():
                break
    
    # 提取第一行
    first_line = first_text.split('\n')[0].strip()
    safe = _safe_filename(first_line) if first_line else "document"
    
    # 限制长度
    if len(safe) > 50:
        safe = safe[:50]
    
    ts = int(time.time() * 1000)
    return f"{safe}_{ts}"