#!/usr/bin/env python3
"""Generate single AprilTag artwork in raster or vector formats."""

import argparse
from dataclasses import dataclass
import os
from pathlib import Path
import sys
from typing import List, Optional, Sequence, 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.patches as mpatches
    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
MARKER_BORDER_BITS = 1
SUPPORTED_SUFFIXES = {".png", ".pdf", ".svg", ".ai"}


@dataclass(frozen=True)
class TagGeometry:
    tag_id: int
    family: str
    marker_size_cm: float
    tag_white_ratio: float
    tag_white_border_cm: float
    tag_tile_size_cm: float
    margin_cm: float
    canvas_size_cm: float
    marker_modules: int


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, MARKER_BORDER_BITS)
    return marker


def compute_tag_dimensions(
    size_cm: float,
    size_mode: str,
    tag_white_ratio_text: str,
    tag_white_border_cm: Optional[float],
    margin_cm: float,
) -> Tuple[float, float, float, float, float]:
    if size_cm <= 0:
        raise SystemExit("size-cm/tag-size-cm must be > 0")
    if margin_cm < 0:
        raise SystemExit("margin-cm must be >= 0")

    if tag_white_border_cm is None:
        try:
            tag_white_ratio = parse_ratio_text(tag_white_ratio_text)
        except ValueError as exc:
            raise SystemExit(f"tag-white-ratio parse error: {exc}") from exc

        if size_mode == "marker":
            marker_size_cm = size_cm
            tag_white_border_cm = marker_size_cm * tag_white_ratio / 2.0
            tag_tile_size_cm = marker_size_cm + 2.0 * tag_white_border_cm
            canvas_size_cm = tag_tile_size_cm + 2.0 * margin_cm
        elif size_mode == "canvas":
            canvas_size_cm = size_cm
            tag_tile_size_cm = canvas_size_cm - 2.0 * margin_cm
            if tag_tile_size_cm <= 0:
                raise SystemExit("canvas size must be larger than 2 * margin-cm")
            marker_size_cm = tag_tile_size_cm / (1.0 + tag_white_ratio)
            if marker_size_cm <= 0:
                raise SystemExit("canvas size is too small for the requested white border ratio")
            tag_white_border_cm = marker_size_cm * tag_white_ratio / 2.0
        else:
            raise SystemExit(f"Unsupported size mode '{size_mode}'")
    else:
        if tag_white_border_cm < 0:
            raise SystemExit("tag-white-border-cm must be >= 0")
        if size_mode == "marker":
            marker_size_cm = size_cm
            tag_tile_size_cm = marker_size_cm + 2.0 * tag_white_border_cm
            canvas_size_cm = tag_tile_size_cm + 2.0 * margin_cm
        elif size_mode == "canvas":
            canvas_size_cm = size_cm
            tag_tile_size_cm = canvas_size_cm - 2.0 * margin_cm
            if tag_tile_size_cm <= 0:
                raise SystemExit("canvas size must be larger than 2 * margin-cm")
            marker_size_cm = tag_tile_size_cm - 2.0 * tag_white_border_cm
            if marker_size_cm <= 0:
                raise SystemExit(
                    "canvas size is too small after subtracting margin-cm and tag-white-border-cm"
                )
        else:
            raise SystemExit(f"Unsupported size mode '{size_mode}'")
        tag_white_ratio = 2.0 * tag_white_border_cm / marker_size_cm

    return (
        marker_size_cm,
        tag_white_ratio,
        tag_white_border_cm,
        tag_tile_size_cm,
        canvas_size_cm,
    )


def build_tag_geometry(
    tag_id: int,
    size_cm: float,
    size_mode: str,
    family: str,
    tag_white_ratio_text: str,
    tag_white_border_cm: Optional[float],
    margin_cm: float,
) -> Tuple[TagGeometry, "cv2.aruco_Dictionary"]:
    if tag_id < 0:
        raise SystemExit("id must be >= 0")

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

    (
        marker_size_cm,
        tag_white_ratio,
        resolved_white_border_cm,
        tag_tile_size_cm,
        canvas_size_cm,
    ) = compute_tag_dimensions(
        size_cm=size_cm,
        size_mode=size_mode,
        tag_white_ratio_text=tag_white_ratio_text,
        tag_white_border_cm=tag_white_border_cm,
        margin_cm=margin_cm,
    )

    geometry = TagGeometry(
        tag_id=tag_id,
        family=family,
        marker_size_cm=marker_size_cm,
        tag_white_ratio=tag_white_ratio,
        tag_white_border_cm=resolved_white_border_cm,
        tag_tile_size_cm=tag_tile_size_cm,
        margin_cm=margin_cm,
        canvas_size_cm=canvas_size_cm,
        marker_modules=aruco_dict.markerSize + 2 * MARKER_BORDER_BITS,
    )
    return geometry, aruco_dict


def make_marker_module_grid(
    aruco_dict: "cv2.aruco_Dictionary",
    marker_id: int,
    marker_modules: int,
) -> np.ndarray:
    pixels_per_module = 16
    marker = make_marker(aruco_dict, marker_id, marker_modules * pixels_per_module)
    return (
        marker.reshape(marker_modules, pixels_per_module, marker_modules, pixels_per_module)
        .mean(axis=(1, 3))
        < 128
    )


def build_vector_rects(geometry: TagGeometry, module_grid: np.ndarray) -> List[Tuple[float, float, float, float]]:
    rects = []
    module_size_cm = geometry.marker_size_cm / float(geometry.marker_modules)
    origin_cm = geometry.margin_cm + geometry.tag_white_border_cm

    for row_index in range(geometry.marker_modules):
        start_col = None
        for col_index in range(geometry.marker_modules + 1):
            is_black = col_index < geometry.marker_modules and bool(module_grid[row_index, col_index])
            if is_black and start_col is None:
                start_col = col_index
            elif not is_black and start_col is not None:
                rects.append(
                    (
                        origin_cm + start_col * module_size_cm,
                        origin_cm + row_index * module_size_cm,
                        (col_index - start_col) * module_size_cm,
                        module_size_cm,
                    )
                )
                start_col = None

    return rects


def render_raster_canvas(
    geometry: TagGeometry,
    aruco_dict: "cv2.aruco_Dictionary",
    dpi: int,
) -> np.ndarray:
    canvas_px = cm_to_px(geometry.canvas_size_cm, dpi)
    tag_tile_px = cm_to_px(geometry.tag_tile_size_cm, dpi)
    marker_px = cm_to_px(geometry.marker_size_cm, 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(geometry.margin_cm, dpi)
    tile_y_px = cm_to_px(geometry.margin_cm, 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(geometry.margin_cm + geometry.tag_white_border_cm, dpi)
    marker_y_px = cm_to_px(geometry.margin_cm + geometry.tag_white_border_cm, 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, geometry.tag_id, marker_px)
    canvas[marker_y_px : marker_y_px + marker_px, marker_x_px : marker_x_px + marker_px] = marker
    return canvas


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 resolve_page_dimensions(
    canvas_width_cm: float,
    canvas_height_cm: float,
    center_on_a4: bool,
    landscape: bool,
) -> Tuple[float, float, float, float]:
    if not center_on_a4:
        return canvas_width_cm, canvas_height_cm, 0.0, 0.0

    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 canvas_width_cm > page_w_cm + 1e-9 or canvas_height_cm > page_h_cm + 1e-9:
        raise SystemExit(
            f"Tag size {canvas_width_cm}x{canvas_height_cm} cm does not fit A4 "
            f"({'landscape' if landscape else 'portrait'}) {page_w_cm}x{page_h_cm} cm."
        )

    return (
        page_w_cm,
        page_h_cm,
        (page_w_cm - canvas_width_cm) / 2.0,
        (page_h_cm - canvas_height_cm) / 2.0,
    )


def save_vector_art(
    output_path: Path,
    geometry: TagGeometry,
    rects_cm: Sequence[Tuple[float, float, float, float]],
    output_format: str,
    dpi: int,
    center_on_a4: bool,
    landscape: bool,
) -> None:
    page_w_cm, page_h_cm, offset_x_cm, offset_y_cm = resolve_page_dimensions(
        geometry.canvas_size_cm,
        geometry.canvas_size_cm,
        center_on_a4=center_on_a4,
        landscape=landscape,
    )

    fig = plt.figure(figsize=(page_w_cm / 2.54, page_h_cm / 2.54), dpi=dpi, frameon=False)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_xlim(0, page_w_cm)
    ax.set_ylim(page_h_cm, 0)
    ax.set_aspect("equal")
    ax.axis("off")

    ax.add_patch(
        mpatches.Rectangle((0.0, 0.0), page_w_cm, page_h_cm, facecolor="white", edgecolor="none")
    )

    for x_cm, y_cm, width_cm, height_cm in rects_cm:
        ax.add_patch(
            mpatches.Rectangle(
                (offset_x_cm + x_cm, offset_y_cm + y_cm),
                width_cm,
                height_cm,
                facecolor="black",
                edgecolor="none",
            )
        )

    fig.savefig(output_path, format=output_format, dpi=dpi, pad_inches=0)
    plt.close(fig)


def normalize_output_suffix(output_path: Path) -> str:
    suffix = output_path.suffix.lower()
    if suffix == ".cdr":
        raise SystemExit(
            "Native .cdr export is not supported in this script. "
            "Use .ai or .svg output and open/import it in CorelDRAW."
        )
    if suffix not in SUPPORTED_SUFFIXES:
        supported = ", ".join(sorted(SUPPORTED_SUFFIXES | {".cdr"}))
        raise SystemExit(f"Unsupported output suffix '{suffix}'. Supported: {supported}")
    return suffix


def save_tag_artwork(
    output_path: Path,
    geometry: TagGeometry,
    aruco_dict: "cv2.aruco_Dictionary",
    dpi: int,
    a4_landscape: bool,
    exact_size_pdf: bool,
) -> None:
    suffix = normalize_output_suffix(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if suffix == ".png":
        canvas = render_raster_canvas(geometry, aruco_dict, dpi)
        save_printable_png(output_path, canvas, geometry.canvas_size_cm, geometry.canvas_size_cm, dpi)
        return

    module_grid = make_marker_module_grid(aruco_dict, geometry.tag_id, geometry.marker_modules)
    rects_cm = build_vector_rects(geometry, module_grid)

    if suffix == ".svg":
        save_vector_art(
            output_path,
            geometry,
            rects_cm,
            output_format="svg",
            dpi=dpi,
            center_on_a4=False,
            landscape=False,
        )
        return

    if suffix == ".ai":
        save_vector_art(
            output_path,
            geometry,
            rects_cm,
            output_format="pdf",
            dpi=dpi,
            center_on_a4=False,
            landscape=False,
        )
        return

    save_vector_art(
        output_path,
        geometry,
        rects_cm,
        output_format="pdf",
        dpi=dpi,
        center_on_a4=not exact_size_pdf,
        landscape=a4_landscape,
    )


def describe_geometry(geometry: TagGeometry) -> str:
    return (
        f"family={geometry.family}, id={geometry.tag_id}, marker={round(geometry.marker_size_cm, 6)}cm, "
        f"tag_white_ratio={round(geometry.tag_white_ratio, 6)}, "
        f"tag_white_border={round(geometry.tag_white_border_cm, 6)}cm, "
        f"tag_tile={round(geometry.tag_tile_size_cm, 6)}cm, "
        f"margin={round(geometry.margin_cm, 6)}cm, "
        f"canvas={round(geometry.canvas_size_cm, 6)}cm"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate single AprilTag artwork")
    parser.add_argument("--id", type=int, default=0, help="AprilTag ID (default: 0)")
    parser.add_argument(
        "--size-cm",
        "--tag-size-cm",
        dest="size_cm",
        type=float,
        default=15.0,
        help=(
            "Base size in cm. Use --size-mode marker to treat it as the black tag size, "
            "or --size-mode canvas to treat it as the final exported square size."
        ),
    )
    parser.add_argument(
        "--size-mode",
        type=str,
        default="marker",
        choices=("marker", "canvas"),
        help="Interpret --size-cm as marker size or final canvas size.",
    )
    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 and --exact-size-pdf is not set, use A4 landscape instead of portrait",
    )
    parser.add_argument(
        "--exact-size-pdf",
        action="store_true",
        help="For .pdf output, save an exact-size vector page instead of centering on A4",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("apriltag_id0_A4.pdf"),
        help=(
            "Output path. Supported suffixes: .png, .pdf, .svg, .ai. "
            ".ai is written as a PDF-based Illustrator-compatible vector file."
        ),
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    geometry, aruco_dict = build_tag_geometry(
        tag_id=args.id,
        size_cm=args.size_cm,
        size_mode=args.size_mode,
        family=args.family,
        tag_white_ratio_text=args.tag_white_ratio,
        tag_white_border_cm=args.tag_white_border_cm,
        margin_cm=args.margin_cm,
    )

    save_tag_artwork(
        output_path=args.output,
        geometry=geometry,
        aruco_dict=aruco_dict,
        dpi=args.dpi,
        a4_landscape=args.a4_landscape,
        exact_size_pdf=args.exact_size_pdf,
    )

    print(f"Saved single tag image: {args.output}")
    print(describe_geometry(geometry))
    return 0


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