"""将解压后的 3MF 文件夹重新打包为 .3mf 文件的工具脚本。

用法:
    python repack_3mf.py "解压目录" "output.3mf"

会自动更新 Metadata/plate_N.gcode.md5 并确保 [Content_Types].xml 为 ZIP 第一条目。

如果省略输出路径，自动在解压目录同级生成同名 .3mf。
"""

from __future__ import annotations

import hashlib
import os
import re
import sys
import zipfile
from pathlib import Path

GCODE_MEMBER_RE = re.compile(r"^Metadata[\\/]plate_(\d+)\.gcode$")


def update_md5(src_dir: Path) -> None:
    """更新所有 plate_N.gcode 对应的 .md5 文件。"""
    metadata = src_dir / "Metadata"
    if not metadata.is_dir():
        return
    for gcode_file in sorted(metadata.glob("plate_*.gcode")):
        if not GCODE_MEMBER_RE.match(str(gcode_file.relative_to(src_dir)).replace("\\", "/")):
            continue
        h = hashlib.md5()
        with open(gcode_file, "rb") as f:
            h.update(f.read())
        md5_path = Path(str(gcode_file) + ".md5")
        md5_path.write_text(h.hexdigest(), encoding="ascii")
        print(f"  MD5 已更新: {md5_path.name} → {h.hexdigest()}")


def repack(src_dir: Path, output: Path) -> None:
    """将解压目录重新打包为 3MF。"""
    ct_path = src_dir / "[Content_Types].xml"
    if not ct_path.exists():
        raise FileNotFoundError(f"未找到 [Content_Types].xml，请确认 {src_dir} 是有效的 3MF 解压目录")

    print(f"源目录: {src_dir}")
    update_md5(src_dir)

    with zipfile.ZipFile(output, "w", zipfile.ZIP_DEFLATED) as z:
        z.write(ct_path, "[Content_Types].xml")

        count = 1
        for root, _dirs, files in os.walk(str(src_dir)):
            for fn in files:
                fp = Path(root) / fn
                arcname = fp.relative_to(src_dir).as_posix()
                if arcname == "[Content_Types].xml":
                    continue
                z.write(str(fp), arcname)
                count += 1

    size_mb = output.stat().st_size / (1024 * 1024)
    print(f"输出文件: {output} ({size_mb:.1f} MB, {count} 条目)")
    print("打包完成。")


def main() -> int:
    if len(sys.argv) < 2:
        print(__doc__)
        return 1

    src = Path(sys.argv[1])
    if not src.is_dir():
        print(f"错误: 目录不存在 — {src}")
        return 1

    if len(sys.argv) >= 3:
        output = Path(sys.argv[2])
    else:
        output = src.parent / (src.name.rstrip("/\\") + ".3mf")

    try:
        repack(src, output)
    except Exception as exc:
        print(f"错误: {exc}")
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())