# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "numpy",
#     "opencv-python-headless",
#     "pymupdf",
# ]
# ///

import argparse
import hashlib
import os
import warnings

import cv2
import numpy as np
import pymupdf


def generate_marker(marker_id: int, marker_size: int = 200, white_offset: int = 15):
    apriltag_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_APRILTAG_36h11)
    if marker_id >= apriltag_dict.bytesList.shape[0]:
        warnings.warn(
            f"Marker ID {marker_id} exceeds the available range for this dictionary "
            f"({apriltag_dict.bytesList.shape[0]} IDs). Markers may not be unique.",
            UserWarning,
            stacklevel=2,
        )
    marker_data = cv2.aruco.generateImageMarker(
        dictionary=apriltag_dict,
        id=marker_id % apriltag_dict.bytesList.shape[0],
        sidePixels=marker_size,
        borderBits=1,
    )
    total_size = marker_size + white_offset * 2
    marker_image = np.full((total_size, total_size), 255, dtype=np.uint8)
    marker_image[
        white_offset : (white_offset + marker_size),
        white_offset : white_offset + marker_size,
    ] = marker_data
    _, jpg_bytes = cv2.imencode(".jpg", marker_image)
    return jpg_bytes.tobytes()


def add_markers_to_pdf(
    input_pdf_path: str, marker_size: int, padding: int = 10, white_margin: int = 15
):
    """Adds markers to each page of a PDF at specified positions and saves the PDF.

    Args:
        input_pdf_path (str): Path to the input PDF file.
        marker_size (int): Size (in pixels) of the markers.
        marker_ids (list): List of marker IDs to generate and place.
        padding (int): How many pixels to move the markers inwards
        white_margin (int): How many white pixels as margin
    """
    pdf_document = pymupdf.open(input_pdf_path)

    markers, hashes = {}, {}
    for marker_id in range(len(pdf_document) * 4):
        marker = generate_marker(marker_id, marker_size)
        marker_hash = hashlib.md5(marker).hexdigest()  # noqa: S324
        if marker_hash in hashes:
            print(
                f"""Warning: Marker with ID {marker_id} is identical to marker with 
                ID {hashes[marker_hash]}"""
            )
        else:
            hashes[marker_hash] = marker_id
        markers[marker_id] = marker

    for page_number in range(len(pdf_document)):
        page = pdf_document[page_number]

        positions = [
            (padding, padding),  # top-left
            (page.rect.width - (marker_size) - padding, padding),  # top-right
            (padding, page.rect.height - marker_size - padding),  # bottom-left
            (
                page.rect.width - (marker_size) - padding,
                page.rect.height - marker_size - padding,
            ),  # bottom-right
        ]

        for marker_id, position in zip(
            [page_number * 4 + i for i in range(4)], positions, strict=False
        ):
            x, y = position
            page.insert_image(
                pymupdf.Rect(x, y, x + marker_size, y + marker_size),
                stream=markers[marker_id],
                xref=0,
            )

    output_path = os.path.splitext(input_pdf_path)[0] + "_marked.pdf"
    pdf_document.save(output_path)
    print(f"PDF saved at {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Add markers to the corners of each page in a PDF."
    )
    parser.add_argument(
        "input_pdf",
        help="Path to the input PDF file.",
    )
    parser.add_argument(
        "--marker_size",
        type=int,
        default=50,
        help="Size of the markers.",
    )
    parser.add_argument(
        "--padding",
        type=int,
        default=10,
        help="Padding of the markers.",
    )
    parser.add_argument(
        "--white_margin",
        type=int,
        default=15,
        help="White margin of the markers.",
    )
    args = parser.parse_args()

    add_markers_to_pdf(
        args.input_pdf, int(args.marker_size), int(args.padding), int(args.white_margin)
    )