#!/usr/bin/env python3
"""Extract every embedded image from a PDF, named by page number.

Writes <outdir>/page_<XXXX>_<N>.<ext> where <N> is the image's index on the
page. The original compressed stream (e.g. JPEG2000) is preserved as-is.
"""

import argparse
import os

import pymupdf


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("pdf", help="input PDF")
    ap.add_argument("-o", "--outdir", default="images", help="output directory")
    args = ap.parse_args()

    doc = pymupdf.open(args.pdf)
    os.makedirs(args.outdir, exist_ok=True)
    total = 0
    for i in range(doc.page_count):
        page = doc[i]
        seen = set()
        n = 0
        for img in page.get_images(full=True):
            xref = img[0]
            if xref in seen:
                continue
            seen.add(xref)
            info = doc.extract_image(xref)
            name = os.path.join(args.outdir, f"page_{i+1:04d}_{n}.{info['ext']}")
            with open(name, "wb") as f:
                f.write(info["image"])
            n += 1
            total += 1
        if i % 100 == 0:
            print(f"  page {i+1}/{doc.page_count} ...", flush=True)
    print(f"Extracted {total} images to {args.outdir}", flush=True)


if __name__ == "__main__":
    main()