#!/usr/bin/env python3
"""Embed regenerated draw.io figure PNGs back into the source DOCX.

The source Word file already contains figure placeholders and captions. This
script replaces media/image3.png through media/image9.png with the regenerated
draw.io PNG exports, then adjusts each drawing extent to preserve the new image
aspect ratio while keeping the original displayed width.
"""

from __future__ import annotations

import argparse
import shutil
import tempfile
import zipfile
from pathlib import Path
from xml.etree import ElementTree as ET

try:
    from PIL import Image
except ImportError as exc:  # pragma: no cover - environment guard
    raise SystemExit(f"Pillow is required to read PNG dimensions: {exc}")


NS = {
    "a": "http://schemas.openxmlformats.org/drawingml/2006/main",
    "pic": "http://schemas.openxmlformats.org/drawingml/2006/picture",
    "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
    "wp": "http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing",
}

for prefix, uri in NS.items():
    ET.register_namespace(prefix, uri)


FIGURE_MEDIA = {
    "rId13": ("word/media/image3.png", "fig-01-system-overall-architecture.png"),
    "rId14": ("word/media/image4.png", "fig-02-network-deployment-security-zone.png"),
    "rId15": ("word/media/image5.png", "fig-03-business-loop.png"),
    "rId16": ("word/media/image6.png", "fig-04-key-technology-route.png"),
    "rId17": ("word/media/image7.png", "fig-05-risk-warning-loop.png"),
    "rId18": ("word/media/image8.png", "fig-06-application-device-collaboration.png"),
    "rId19": ("word/media/image9.png", "fig-07-network-security-defense.png"),
}


def image_ratio(path: Path) -> float:
    with Image.open(path) as image:
        width, height = image.size
    return height / width


def find_parent(root: ET.Element, child: ET.Element) -> ET.Element | None:
    for parent in root.iter():
        for node in list(parent):
            if node is child:
                return parent
    return None


def find_ancestor(root: ET.Element, child: ET.Element, tag: str) -> ET.Element | None:
    node = child
    while True:
        parent = find_parent(root, node)
        if parent is None:
            return None
        if parent.tag == tag:
            return parent
        node = parent


def update_extents(document_xml: Path, ratio_by_rid: dict[str, float]) -> None:
    tree = ET.parse(document_xml)
    root = tree.getroot()
    wp_extent_tag = f"{{{NS['wp']}}}extent"
    a_ext_tag = f"{{{NS['a']}}}ext"
    wp_inline_tag = f"{{{NS['wp']}}}inline"
    wp_anchor_tag = f"{{{NS['wp']}}}anchor"

    for blip in root.findall(".//a:blip", NS):
        rid = blip.attrib.get(f"{{{NS['r']}}}embed")
        if rid not in ratio_by_rid:
            continue

        container = find_ancestor(root, blip, wp_inline_tag)
        if container is None:
            container = find_ancestor(root, blip, wp_anchor_tag)
        if container is None:
            continue

        wp_extent = container.find("wp:extent", NS)
        if wp_extent is None:
            continue
        width = int(wp_extent.attrib["cx"])
        height = max(1, round(width * ratio_by_rid[rid]))
        wp_extent.set("cy", str(height))

        for a_ext in container.findall(".//a:ext", NS):
            if "cx" in a_ext.attrib and "cy" in a_ext.attrib:
                a_ext.set("cx", str(width))
                a_ext.set("cy", str(height))

    tree.write(document_xml, encoding="UTF-8", xml_declaration=True)


def replace_docx_images(source_docx: Path, exports_dir: Path, output_docx: Path) -> None:
    with tempfile.TemporaryDirectory() as tmp:
        work = Path(tmp) / "docx"
        work.mkdir()
        with zipfile.ZipFile(source_docx) as zin:
            zin.extractall(work)

        ratio_by_rid: dict[str, float] = {}
        for rid, (media_name, export_name) in FIGURE_MEDIA.items():
            export_path = exports_dir / export_name
            target_path = work / media_name
            if not export_path.exists():
                raise SystemExit(f"missing export: {export_path}")
            if not target_path.exists():
                raise SystemExit(f"missing target media in DOCX: {media_name}")
            shutil.copyfile(export_path, target_path)
            ratio_by_rid[rid] = image_ratio(export_path)

        update_extents(work / "word" / "document.xml", ratio_by_rid)

        output_docx.parent.mkdir(parents=True, exist_ok=True)
        if output_docx.exists():
            output_docx.unlink()
        with zipfile.ZipFile(output_docx, "w", compression=zipfile.ZIP_DEFLATED) as zout:
            for path in sorted(work.rglob("*")):
                if path.is_file():
                    zout.write(path, path.relative_to(work).as_posix())


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-docx", required=True, type=Path)
    parser.add_argument("--exports-dir", required=True, type=Path)
    parser.add_argument("--output-docx", required=True, type=Path)
    args = parser.parse_args()
    replace_docx_images(args.source_docx, args.exports_dir, args.output_docx)
    print(args.output_docx)


if __name__ == "__main__":
    main()