#!/usr/bin/env python3
"""Extract the 16 museum-map artifact images from EgyptEscapeRoomPrintables.pdf."""

from __future__ import annotations

import argparse
import io
import json
import sys
from pathlib import Path

try:
    import fitz  # PyMuPDF
except ImportError:
    print("PyMuPDF is required: pip install pymupdf", file=sys.stderr)
    raise SystemExit(1) from None

from PIL import Image

# Identified by xref on pages 9–10 (puzzle 5 map). Slugs match the answer sheet.
ARTIFACTS: list[dict[str, int | str]] = [
    {"page": 9, "xref": 56, "slug": "double-crown", "name": "Double Crown"},
    {"page": 9, "xref": 57, "slug": "obelisk", "name": "Obelisk"},
    {"page": 9, "xref": 63, "slug": "goddess-isis", "name": "Goddess Isis"},
    {"page": 9, "xref": 53, "slug": "canopic-jars", "name": "Canopic Jars"},
    {"page": 9, "xref": 61, "slug": "howard-carter-excavating-king-tut", "name": "Howard Carter excavating King Tut"},
    {"page": 9, "xref": 54, "slug": "anubis", "name": "Anubis"},
    {"page": 9, "xref": 58, "slug": "bust-of-nefertiti", "name": "Bust of Nefertiti"},
    {"page": 9, "xref": 60, "slug": "god-osiris", "name": "God Osiris"},
    {"page": 10, "xref": 68, "slug": "pyramid-of-giza", "name": "Pyramid of Giza"},
    {"page": 10, "xref": 75, "slug": "temple-of-hatshepsut", "name": "Temple of Hatshepsut"},
    {"page": 10, "xref": 72, "slug": "map-of-the-nile", "name": "Map of the Nile"},
    {"page": 10, "xref": 67, "slug": "egyptian-painting", "name": "Egyptian painting"},
    {"page": 10, "xref": 74, "slug": "sarcophagus", "name": "Sarcophagus"},
    {"page": 10, "xref": 70, "slug": "statue-of-hatshepsut", "name": "Statue of Hatshepsut"},
    {"page": 10, "xref": 76, "slug": "step-pyramid", "name": "Step Pyramid"},
    {"page": 10, "xref": 71, "slug": "mastaba", "name": "Mastaba"},
]


def find_image_info(page: fitz.Page, xref: int) -> dict:
    for info in page.get_image_info(xrefs=True):
        if info["xref"] == xref:
            return info
    raise ValueError(f"Image xref {xref} not found on page {page.number + 1}")


def render_clip(page: fitz.Page, bbox: tuple[float, float, float, float], dpi: float) -> Image.Image:
    matrix = fitz.Matrix(dpi / 72.0, dpi / 72.0)
    pixmap = page.get_pixmap(matrix=matrix, clip=fitz.Rect(bbox), alpha=False)
    return Image.open(io.BytesIO(pixmap.tobytes("png")))


def extract_embedded(doc: fitz.Document, xref: int) -> tuple[Image.Image, str]:
    data = doc.extract_image(xref)
    ext = data["ext"]
    image = Image.open(io.BytesIO(data["image"]))

    smask = data.get("smask")
    if smask:
        mask_data = doc.extract_image(smask)
        mask = Image.open(io.BytesIO(mask_data["image"])).convert("L")
        if mask.size != image.size:
            mask = mask.resize(image.size, Image.Resampling.LANCZOS)
        rgba = image.convert("RGBA")
        rgba.putalpha(mask)
        background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
        image = Image.alpha_composite(background, rgba).convert("RGB")
        ext = "png"

    return image, ext


def extract_artifacts(
    pdf_path: str,
    output_dir: Path,
    *,
    dpi: float = 300,
    mode: str = "render",
) -> list[dict[str, object]]:
    doc = fitz.open(pdf_path)
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest: list[dict[str, object]] = []

    try:
        for artifact in ARTIFACTS:
            page_number = int(artifact["page"])
            xref = int(artifact["xref"])
            slug = str(artifact["slug"])
            name = str(artifact["name"])
            page = doc[page_number - 1]
            info = find_image_info(page, xref)

            if mode == "render":
                image = render_clip(page, info["bbox"], dpi)
                filename = f"{slug}.png"
                image.save(output_dir / filename, optimize=True)
            else:
                image, ext = extract_embedded(doc, xref)
                filename = f"{slug}.{ext}"
                if ext == "jpeg":
                    image.save(output_dir / filename, quality=92, optimize=True)
                else:
                    image.save(output_dir / filename, optimize=True)

            manifest.append(
                {
                    "slug": slug,
                    "name": name,
                    "page": page_number,
                    "xref": xref,
                    "file": filename,
                    "width": image.width,
                    "height": image.height,
                    "bbox": info["bbox"],
                }
            )
            print(f"Wrote {filename} ({image.width}x{image.height})")
    finally:
        doc.close()

    manifest_path = output_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {manifest_path}")
    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("pdf", help="Path to EgyptEscapeRoomPrintables.pdf")
    parser.add_argument(
        "output_dir",
        nargs="?",
        default="public/media/artifacts",
        help="Output directory (default: public/media/artifacts)",
    )
    parser.add_argument("--dpi", type=float, default=300, help="Render DPI for --mode render (default: 300)")
    parser.add_argument(
        "--mode",
        choices=("render", "embedded"),
        default="render",
        help="render: clip from page (default, best for illustrations); embedded: raw image bytes",
    )
    args = parser.parse_args()

    extract_artifacts(args.pdf, Path(args.output_dir), dpi=args.dpi, mode=args.mode)


if __name__ == "__main__":
    main()