# -*- coding: utf-8 -*-
"""
style_transfer.py：从用户上传的 docx 抽取视觉样式（theme1.xml + styles.xml 的默认字体引用），
套到内置骨架渲染结果上，产出"风格接近用户模板"的文档。

Level 1 (MVP)：整体替换 theme + 默认字体引用，不整体覆盖 styles（避免版式崩）。
"""
import os
import zipfile
import tempfile
import shutil
import xml.etree.ElementTree as ET

# Word OOXML 命名空间
_NS = {
    "a": "http://schemas.openxmlformats.org/drawingml/2006/main",
    "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
    "w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main",
}

# 主题字体引用中 major/minor 属性的前缀
_THEME_PREFIXES = ("majorHAnsi", "minorHAnsi", "majorEastAsia", "minorEastAsia",
                   "majorCs", "minorCs", "major", "minor")


def _unzip_docx(docx_path, target_dir):
    """解压 docx（zip）到目标目录。"""
    with zipfile.ZipFile(docx_path, "r") as z:
        z.extractall(target_dir)


def _zip_docx(source_dir, docx_path):
    """将目录内容打包为 docx（zip）。"""
    with zipfile.ZipFile(docx_path, "w", zipfile.ZIP_DEFLATED) as z:
        for root, _dirs, files in os.walk(source_dir):
            for fname in files:
                full = os.path.join(root, fname)
                arcname = os.path.relpath(full, source_dir)
                z.write(full, arcname)


def _safe_xml_read(xml_path):
    """安全读取 XML 文件，返回 ElementTree 和根元素。失败返回 (None, None)。"""
    try:
        tree = ET.parse(xml_path)
        return tree, tree.getroot()
    except Exception:
        return None, None


def _copy_theme(src_dir, dst_dir):
    """将 src 的 word/theme/theme1.xml 复制到 dst。"""
    src_theme = os.path.join(src_dir, "word", "theme", "theme1.xml")
    dst_theme = os.path.join(dst_dir, "word", "theme", "theme1.xml")
    if os.path.isfile(src_theme):
        os.makedirs(os.path.dirname(dst_theme), exist_ok=True)
        shutil.copy2(src_theme, dst_theme)
        return True
    return False


def _copy_font_table(src_dir, dst_dir):
    """将 src 的 word/fontTable.xml 复制到 dst。"""
    src_ft = os.path.join(src_dir, "word", "fontTable.xml")
    dst_ft = os.path.join(dst_dir, "word", "fontTable.xml")
    if os.path.isfile(src_ft):
        shutil.copy2(src_ft, dst_ft)
        return True
    return False


def _merge_default_fonts(src_styles_path, dst_styles_path):
    """合并默认字体引用（docDefaults 中的 rFonts）从 src 到 dst。

    不整体覆盖 styles.xml，只替换 docDefaults 里的字体引用。
    """
    if not os.path.isfile(src_styles_path) or not os.path.isfile(dst_styles_path):
        return False

    _, src_root = _safe_xml_read(src_styles_path)
    dst_tree, dst_root = _safe_xml_read(dst_styles_path)
    if src_root is None or dst_root is None:
        return False

    w_ns = _NS["w"]

    # 找 src 的 docDefaults
    src_doc_defaults = src_root.find(f".//{{{w_ns}}}docDefaults")
    if src_doc_defaults is None:
        return False

    src_rpr = src_doc_defaults.find(f"{{{w_ns}}}rPrDefault/{{{w_ns}}}rFonts")
    if src_rpr is None:
        return False

    # 找 dst 的 docDefaults
    dst_doc_defaults = dst_root.find(f".//{{{w_ns}}}docDefaults")
    if dst_doc_defaults is None:
        return False

    dst_rpr = dst_doc_defaults.find(f"{{{w_ns}}}rPrDefault/{{{w_ns}}}rFonts")
    if dst_rpr is None:
        # 如果 dst 没有 rFonts，创建并追加
        rpr_default = dst_doc_defaults.find(f"{{{w_ns}}}rPrDefault")
        if rpr_default is None:
            rpr_default = ET.SubElement(dst_doc_defaults, f"{{{w_ns}}}rPrDefault")
        rpr_default.append(src_rpr)
    else:
        # 替换 rFonts 的属性（主题字体引用）
        for attr_name, attr_val in src_rpr.attrib.items():
            if any(attr_name.endswith(prefix) for prefix in _THEME_PREFIXES):
                dst_rpr.set(attr_name, attr_val)

    # 保存
    dst_tree.write(dst_styles_path, xml_declaration=True, encoding="UTF-8", standalone="yes")
    return True


def transplant_style(result_path, user_template_path, out_path):
    """将用户模板的视觉样式移植到结果文档上。

    参数:
        result_path:        内置骨架渲染出的 docx 路径
        user_template_path: 用户上传的样式参考 docx
        out_path:           输出路径（可与 result_path 相同，原地覆盖）
    """
    # 校验用户模板可打开
    try:
        from docx import Document
        Document(user_template_path)
    except Exception as e:
        print(f"[warn] 用户上传模板无法打开，跳过样式迁移: {e}")
        if out_path != result_path:
            shutil.copy2(result_path, out_path)
        return out_path

    # 在打包覆盖前先记录原文档段落数（out_path 可能 == result_path，原地覆盖后就读不到原值了）
    src_count = None
    try:
        from docx import Document
        src_count = len(Document(result_path).paragraphs)
    except Exception as e:
        print(f"[warn] 读取原文档段落数失败，迁移后将跳过段落数校验: {e}")

    # 创建临时工作目录
    with tempfile.TemporaryDirectory() as tmpdir:
        src_dir = os.path.join(tmpdir, "src")   # 用户上传模板
        dst_dir = os.path.join(tmpdir, "dst")   # 渲染结果
        os.makedirs(src_dir)
        os.makedirs(dst_dir)

        # 解压
        _unzip_docx(user_template_path, src_dir)
        _unzip_docx(result_path, dst_dir)

        # 1) 替换 theme1.xml
        _copy_theme(src_dir, dst_dir)

        # 2) 替换 fontTable.xml
        _copy_font_table(src_dir, dst_dir)

        # 3) 合并默认字体引用
        src_styles = os.path.join(src_dir, "word", "styles.xml")
        dst_styles = os.path.join(dst_dir, "word", "styles.xml")
        _merge_default_fonts(src_styles, dst_styles)

        # 重新打包
        _zip_docx(dst_dir, out_path)

    # 校验：能正常打开且段落数与迁移前一致
    try:
        from docx import Document
        dst_count = len(Document(out_path).paragraphs)
        if src_count is None:
            print(f"[ok] 样式迁移完成（无原始段落数可比对），输出段落数: {dst_count}")
        elif dst_count == 0 or dst_count != src_count:
            print(f"[warn] 样式迁移后段落数不一致 (src={src_count}, dst={dst_count})，可能损坏")
        else:
            print(f"[ok] 样式迁移完成，段落数一致: {src_count}")
    except Exception as e:
        print(f"[error] 样式迁移后文档无法打开: {e}")

    return out_path
