#!/usr/bin/env python3
from __future__ import annotations

import argparse
import shutil
import tempfile
import zipfile
from pathlib import Path
from urllib.parse import unquote, urlparse
from xml.etree import ElementTree as ET


PKG_NS = "http://schemas.openxmlformats.org/package/2006/relationships"
R_NS = "http://schemas.openxmlformats.org/officeDocument/2006/relationships"
A_NS = "http://schemas.openxmlformats.org/drawingml/2006/main"
CT_NS = "http://schemas.openxmlformats.org/package/2006/content-types"

ET.register_namespace("", PKG_NS)
ET.register_namespace("a", A_NS)
ET.register_namespace("r", R_NS)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Embed externally linked local images into a DOCX file."
    )
    parser.add_argument("input_docx", type=Path)
    parser.add_argument("output_docx", type=Path)
    return parser.parse_args()


def local_path_from_target(target: str) -> Path | None:
    parsed = urlparse(target)
    if parsed.scheme != "file":
        return None
    return Path(unquote(parsed.path))


def ensure_content_type(defaults_root: ET.Element, extension: str) -> None:
    ext = extension.lower().lstrip(".")
    existing = defaults_root.find(f"{{{CT_NS}}}Default[@Extension='{ext}']")
    if existing is not None:
        return

    content_types = {
        "png": "image/png",
        "jpg": "image/jpeg",
        "jpeg": "image/jpeg",
        "gif": "image/gif",
        "bmp": "image/bmp",
        "svg": "image/svg+xml",
    }
    content_type = content_types.get(ext)
    if not content_type:
        return

    ET.SubElement(
        defaults_root,
        f"{{{CT_NS}}}Default",
        {"Extension": ext, "ContentType": content_type},
    )


def main() -> int:
    args = parse_args()
    input_docx = args.input_docx.resolve()
    output_docx = args.output_docx.resolve()

    with tempfile.TemporaryDirectory() as tmpdir:
        workdir = Path(tmpdir)
        with zipfile.ZipFile(input_docx) as zin:
            zin.extractall(workdir)

        rels_path = workdir / "word" / "_rels" / "document.xml.rels"
        doc_path = workdir / "word" / "document.xml"
        types_path = workdir / "[Content_Types].xml"
        media_dir = workdir / "word" / "media"
        media_dir.mkdir(parents=True, exist_ok=True)

        rels_tree = ET.parse(rels_path)
        rels_root = rels_tree.getroot()
        doc_tree = ET.parse(doc_path)
        doc_root = doc_tree.getroot()
        types_tree = ET.parse(types_path)
        types_root = types_tree.getroot()

        rel_to_target: dict[str, str] = {}
        next_image_num = 1

        existing_media = sorted(media_dir.glob("image*.*"))
        if existing_media:
            next_image_num = len(existing_media) + 1

        for rel in rels_root.findall(f"{{{PKG_NS}}}Relationship"):
            rel_type = rel.attrib.get("Type", "")
            target_mode = rel.attrib.get("TargetMode")
            if not rel_type.endswith("/image") or target_mode != "External":
                continue

            target = rel.attrib.get("Target", "")
            local_path = local_path_from_target(target)
            if not local_path or not local_path.exists():
                continue

            image_name = f"image{next_image_num}{local_path.suffix.lower()}"
            next_image_num += 1
            embedded_target = f"media/{image_name}"
            shutil.copy2(local_path, media_dir / image_name)

            rel.attrib["Target"] = embedded_target
            rel.attrib.pop("TargetMode", None)
            rel_to_target[rel.attrib["Id"]] = embedded_target
            ensure_content_type(types_root, local_path.suffix)

        if not rel_to_target:
            raise SystemExit("No external local images found in DOCX relationships.")

        for blip in doc_root.findall(f".//{{{A_NS}}}blip"):
            link_key = f"{{{R_NS}}}link"
            embed_key = f"{{{R_NS}}}embed"
            rel_id = blip.attrib.get(link_key)
            if rel_id in rel_to_target:
                blip.attrib[embed_key] = rel_id
                blip.attrib.pop(link_key, None)

        rels_tree.write(rels_path, encoding="UTF-8", xml_declaration=True)
        doc_tree.write(doc_path, encoding="UTF-8", xml_declaration=True)
        types_tree.write(types_path, encoding="UTF-8", xml_declaration=True)

        if output_docx.exists():
            output_docx.unlink()

        with zipfile.ZipFile(output_docx, "w", zipfile.ZIP_DEFLATED) as zout:
            for path in sorted(workdir.rglob("*")):
                if path.is_dir():
                    continue
                zout.write(path, path.relative_to(workdir).as_posix())

    return 0


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