"""Extract images from a DOCX package into a media folder."""

from __future__ import annotations

import hashlib
import os
import zipfile
from pathlib import Path

from docx.oxml.ns import qn


def _r_id_from_blip(blip) -> str | None:
    if blip is None:
        return None
    return blip.get(qn("r:embed")) or blip.get(
        "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed"
    )


def iter_blip_rids(paragraph_element) -> list[str]:
    """Collect relationship ids for drawings in a paragraph."""
    if paragraph_element is None:
        return []
    rids: list[str] = []
    for blip in paragraph_element.findall(".//" + qn("a:blip")):
        rid = _r_id_from_blip(blip)
        if rid:
            rids.append(rid)
    return rids


class MediaExtractor:
    """Copy word/media parts referenced by rIds into dest_dir."""

    def __init__(self, docx_path: str | Path, dest_dir: str | Path, *, rel_prefix: str):
        self.docx_path = Path(docx_path)
        self.dest_dir = Path(dest_dir)
        self.rel_prefix = rel_prefix.replace("\\", "/").rstrip("/")
        self.dest_dir.mkdir(parents=True, exist_ok=True)
        self._zf = zipfile.ZipFile(self.docx_path, "r")
        self._cache: dict[str, str] = {}  # rid -> relative path for md
        self._used_names: set[str] = set()

    def close(self) -> None:
        self._zf.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def extract_rid(self, document_part, rid: str) -> str | None:
        """Return path relative to the .md file, or None on failure."""
        if rid in self._cache:
            return self._cache[rid]
        try:
            rel = document_part.rels[rid]
        except KeyError:
            return None
        target = getattr(rel, "target_ref", None) or getattr(rel, "target_part", None)
        if target is None:
            return None
        # target_ref like "media/image1.png"
        if hasattr(rel, "target_part"):
            part = rel.target_part
            blob = part.blob
            name = os.path.basename(part.partname)
        else:
            ref = str(target).lstrip("/")
            zip_name = "word/" + ref if not ref.startswith("word/") else ref
            try:
                blob = self._zf.read(zip_name)
            except KeyError:
                return None
            name = os.path.basename(ref)

        stem, ext = os.path.splitext(name)
        if not ext:
            ext = ".png"
        safe = f"{stem}{ext}"
        if safe in self._used_names:
            digest = hashlib.md5(blob).hexdigest()[:8]
            safe = f"{stem}_{digest}{ext}"
        self._used_names.add(safe)
        out_path = self.dest_dir / safe
        out_path.write_bytes(blob)
        rel_path = f"{self.rel_prefix}/{safe}" if self.rel_prefix else safe
        self._cache[rid] = rel_path
        return rel_path