#!/usr/bin/env python3
"""Extract puzzle 5 map pages from the printables PDF and stitch into one image."""

from __future__ import annotations

import argparse
import io
import sys

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, ImageOps


def render_page(doc: fitz.Document, page_number: int, dpi: float) -> Image.Image:
    page = doc[page_number - 1]
    zoom = dpi / 72.0
    matrix = fitz.Matrix(zoom, zoom)
    pixmap = page.get_pixmap(matrix=matrix, alpha=False)
    return Image.open(io.BytesIO(pixmap.tobytes("png")))


def trim_white(image: Image.Image, *, pad: int = 8, threshold: int = 245) -> Image.Image:
    gray = image.convert("L")
    mask = gray.point(lambda value: 255 if value > threshold else 0)
    bbox = mask.getbbox()
    if bbox is None:
        return image

    left, top, right, bottom = bbox
    left = max(0, left - pad)
    top = max(0, top - pad)
    right = min(image.width, right + pad)
    bottom = min(image.height, bottom + pad)
    return image.crop((left, top, right, bottom))


def match_height(left: Image.Image, right: Image.Image) -> tuple[Image.Image, Image.Image]:
    height = max(left.height, right.height)
    if left.height != height:
        canvas = Image.new("RGB", (left.width, height), "white")
        canvas.paste(left, (0, 0))
        left = canvas
    if right.height != height:
        canvas = Image.new("RGB", (right.width, height), "white")
        canvas.paste(right, (0, 0))
        right = canvas
    return left, right


def stitch_map(
    pdf_path: str,
    *,
    left_page: int = 9,
    right_page: int = 10,
    dpi: float = 200,
) -> Image.Image:
    doc = fitz.open(pdf_path)
    try:
        if doc.page_count < max(left_page, right_page):
            raise ValueError(
                f"PDF has {doc.page_count} pages; need at least page {max(left_page, right_page)}"
            )

        left = trim_white(render_page(doc, left_page, dpi))
        right = trim_white(render_page(doc, right_page, dpi))
        left, right = match_height(left, right)

        combined = Image.new("RGB", (left.width + right.width, left.height), "white")
        combined.paste(left, (0, 0))
        combined.paste(right, (left.width, 0))
        return combined
    finally:
        doc.close()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("pdf", help="Path to EgyptEscapeRoomPrintables.pdf")
    parser.add_argument("output", help="Output image path (.jpg or .png)")
    parser.add_argument("--left-page", type=int, default=9, help="Left half page number (default: 9)")
    parser.add_argument("--right-page", type=int, default=10, help="Right half page number (default: 10)")
    parser.add_argument("--dpi", type=float, default=200, help="Render DPI (default: 200)")
    args = parser.parse_args()

    image = stitch_map(
        args.pdf,
        left_page=args.left_page,
        right_page=args.right_page,
        dpi=args.dpi,
    )

    output = args.output.lower()
    if output.endswith(".png"):
        image.save(args.output, optimize=True)
    else:
        image.save(args.output, quality=90, optimize=True)

    print(f"Wrote {args.output} ({image.width}x{image.height})")


if __name__ == "__main__":
    main()