#!/usr/bin/env python3
"""Generate a single AprilTag image with configurable id and physical size."""

import argparse
import os
from pathlib import Path
import sys
from typing import Tuple

try:
    import numpy as np
except ImportError as exc:
    raise SystemExit(
        "Missing dependency 'numpy'. Install it with: pip install numpy"
    ) from exc

try:
    import cv2
except ImportError as exc:
    raise SystemExit(
        "Missing dependency 'opencv-contrib-python'. Install it with: pip install opencv-contrib-python"
    ) from exc

# Force non-GUI backend to avoid Qt/xcb plugin issues in headless environments.
os.environ.setdefault("MPLBACKEND", "Agg")
try:
    import matplotlib

    matplotlib.use("Agg", force=True)
    import matplotlib.pyplot as plt
except ImportError as exc:
    raise SystemExit(
        "Missing dependency 'matplotlib'. Install it with: pip install matplotlib"
    ) from exc


APRILTAG_FAMILIES = {
    "tag16h5": ("DICT_APRILTAG_16h5", 30),
    "tag25h9": ("DICT_APRILTAG_25h9", 35),
    "tag36h10": ("DICT_APRILTAG_36h10", 2320),
    "tag36h11": ("DICT_APRILTAG_36h11", 587),
}

A4_WIDTH_CM = 21.0
A4_HEIGHT_CM = 29.7


def cm_to_px(cm: float, dpi: int) -> int:
    return int(round(cm / 2.54 * dpi))


def parse_ratio_text(value: str) -> float:
    text = value.strip()
    if not text:
        raise ValueError("empty ratio")
    if ":" in text:
        parts = text.split(":")
        if len(parts) != 2:
            raise ValueError(f"invalid ratio '{value}', expected like 1:4")
        left = float(parts[0].strip())
        right = float(parts[1].strip())
        if left < 0 or right <= 0:
            raise ValueError(f"invalid ratio '{value}', left must be >= 0 and right > 0")
        return left / right

    ratio = float(text)
    if ratio < 0:
        raise ValueError(f"invalid ratio '{value}', must be >= 0")
    return ratio


def resolve_aruco_dict(family: str) -> Tuple["cv2.aruco_Dictionary", int]:
    if family not in APRILTAG_FAMILIES:
        supported = ", ".join(sorted(APRILTAG_FAMILIES))
        raise ValueError(f"Unsupported family '{family}'. Supported: {supported}")

    dict_name, max_tags = APRILTAG_FAMILIES[family]
    dict_id = getattr(cv2.aruco, dict_name, None)
    if dict_id is None:
        raise RuntimeError(
            f"OpenCV build does not include {dict_name}. "
            "Please install opencv-contrib-python."
        )

    return cv2.aruco.getPredefinedDictionary(dict_id), max_tags


def make_marker(aruco_dict: "cv2.aruco_Dictionary", marker_id: int, size_px: int) -> np.ndarray:
    if hasattr(cv2.aruco, "generateImageMarker"):
        marker = cv2.aruco.generateImageMarker(aruco_dict, marker_id, size_px)
        if marker is not None:
            return marker

    marker = np.zeros((size_px, size_px), dtype=np.uint8)
    cv2.aruco.drawMarker(aruco_dict, marker_id, size_px, marker, 1)
    return marker


def save_printable_png(output_path: Path, canvas: np.ndarray, width_cm: float, height_cm: float, dpi: int) -> None:
    fig = plt.figure(figsize=(width_cm / 2.54, height_cm / 2.54), dpi=dpi, frameon=False)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.imshow(canvas, cmap="gray", vmin=0, vmax=255, interpolation="nearest")
    ax.axis("off")
    fig.savefig(output_path, dpi=dpi, pad_inches=0)
    plt.close(fig)


def save_a4_pdf(
    output_path: Path,
    canvas: np.ndarray,
    width_cm: float,
    height_cm: float,
    dpi: int,
    landscape: bool,
) -> None:
    if landscape:
        page_w_cm = A4_HEIGHT_CM
        page_h_cm = A4_WIDTH_CM
    else:
        page_w_cm = A4_WIDTH_CM
        page_h_cm = A4_HEIGHT_CM

    if width_cm > page_w_cm + 1e-9 or height_cm > page_h_cm + 1e-9:
        raise SystemExit(
            f"Tag size {width_cm}x{height_cm} cm does not fit A4 "
            f"({'landscape' if landscape else 'portrait'}) {page_w_cm}x{page_h_cm} cm."
        )

    fig = plt.figure(figsize=(page_w_cm / 2.54, page_h_cm / 2.54), dpi=dpi, frameon=False)
    x0 = (page_w_cm - width_cm) / 2.0 / page_w_cm
    y0 = (page_h_cm - height_cm) / 2.0 / page_h_cm
    w = width_cm / page_w_cm
    h = height_cm / page_h_cm
    ax = fig.add_axes([x0, y0, w, h])
    ax.imshow(canvas, cmap="gray", vmin=0, vmax=255, interpolation="nearest")
    ax.axis("off")
    fig.savefig(output_path, format="pdf", dpi=dpi, pad_inches=0)
    plt.close(fig)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate single AprilTag image")
    parser.add_argument("--id", type=int, default=0, help="AprilTag ID (default: 0)")
    parser.add_argument(
        "--size-cm",
        "--tag-size-cm",
        dest="tag_size_cm",
        type=float,
        default=15,
        help="AprilTag marker side length in cm (excluding per-tag white border)",
    )
    parser.add_argument(
        "--family",
        type=str,
        default="tag36h11",
        choices=sorted(APRILTAG_FAMILIES.keys()),
        help="AprilTag family",
    )
    parser.add_argument(
        "--tag-white-ratio",
        type=str,
        default="1:4",
        help="Total white border to marker ratio (e.g. 1:4 means total white width is 1/4 of marker size)",
    )
    parser.add_argument(
        "--tag-white-border-cm",
        type=float,
        default=None,
        help="Optional override: white margin on each side in cm; if omitted, auto-computed from --tag-white-ratio",
    )
    parser.add_argument(
        "--margin-cm",
        type=float,
        default=0.0,
        help="Extra white margin around the full tag tile (cm)",
    )
    parser.add_argument("--dpi", type=int, default=300, help="Output DPI")
    parser.add_argument(
        "--a4-landscape",
        action="store_true",
        help="When output is PDF, use A4 landscape instead of portrait",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("apriltag_id0_A4.pdf"),
        help="Output path (.pdf for centered A4 page, .png for exact-size image)",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    if args.id < 0:
        raise SystemExit("id must be >= 0")
    if args.tag_size_cm <= 0:
        raise SystemExit("size-cm/tag-size-cm must be > 0")
    if args.margin_cm < 0:
        raise SystemExit("margin-cm must be >= 0")

    if args.tag_white_border_cm is None:
        try:
            tag_white_ratio = parse_ratio_text(args.tag_white_ratio)
        except ValueError as exc:
            raise SystemExit(f"tag-white-ratio parse error: {exc}") from exc
        tag_white_border_cm = args.tag_size_cm * tag_white_ratio / 2.0
    else:
        if args.tag_white_border_cm < 0:
            raise SystemExit("tag-white-border-cm must be >= 0")
        tag_white_border_cm = args.tag_white_border_cm
        tag_white_ratio = 2.0 * tag_white_border_cm / args.tag_size_cm

    tag_tile_size_cm = args.tag_size_cm + 2.0 * tag_white_border_cm
    canvas_size_cm = tag_tile_size_cm + 2.0 * args.margin_cm

    aruco_dict, max_tags = resolve_aruco_dict(args.family)
    if args.id >= max_tags:
        raise SystemExit(
            f"Tag ID {args.id} exceeds family max ({max_tags - 1}) for family {args.family}."
        )

    canvas_px = cm_to_px(canvas_size_cm, args.dpi)
    tag_tile_px = cm_to_px(tag_tile_size_cm, args.dpi)
    marker_px = cm_to_px(args.tag_size_cm, args.dpi)

    if canvas_px <= 0 or tag_tile_px <= 0 or marker_px <= 0:
        raise SystemExit("At current DPI, tag size is too small to draw")
    if canvas_px < tag_tile_px:
        raise SystemExit("Canvas size is smaller than the tag tile")

    canvas = np.full((canvas_px, canvas_px), 255, dtype=np.uint8)

    tile_x_px = cm_to_px(args.margin_cm, args.dpi)
    tile_y_px = cm_to_px(args.margin_cm, args.dpi)
    tile_x_px = min(max(tile_x_px, 0), canvas_px - tag_tile_px)
    tile_y_px = min(max(tile_y_px, 0), canvas_px - tag_tile_px)

    marker_x_px = cm_to_px(args.margin_cm + tag_white_border_cm, args.dpi)
    marker_y_px = cm_to_px(args.margin_cm + tag_white_border_cm, args.dpi)
    marker_x_px = min(max(marker_x_px, tile_x_px), tile_x_px + tag_tile_px - marker_px)
    marker_y_px = min(max(marker_y_px, tile_y_px), tile_y_px + tag_tile_px - marker_px)

    marker = make_marker(aruco_dict, args.id, marker_px)
    canvas[marker_y_px : marker_y_px + marker_px, marker_x_px : marker_x_px + marker_px] = marker

    args.output.parent.mkdir(parents=True, exist_ok=True)
    if args.output.suffix.lower() == ".pdf":
        save_a4_pdf(
            args.output,
            canvas,
            canvas_size_cm,
            canvas_size_cm,
            args.dpi,
            landscape=args.a4_landscape,
        )
    else:
        save_printable_png(args.output, canvas, canvas_size_cm, canvas_size_cm, args.dpi)

    print(f"Saved single tag image: {args.output}")
    print(
        f"family={args.family}, id={args.id}, marker={args.tag_size_cm}cm, "
        f"tag_white_ratio={round(tag_white_ratio, 6)}, "
        f"tag_white_border={round(tag_white_border_cm, 6)}cm, "
        f"tag_tile={round(tag_tile_size_cm, 6)}cm, "
        f"margin={args.margin_cm}cm, canvas={round(canvas_size_cm, 6)}cm"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())
