# -*- coding: utf-8 -*-
"""
Semi-automatic pre-annotation tool for engineering drawings.

Pipeline:
    1. Split each raw engineering drawing into front view and left view.
    2. Detect solid and dashed circular primitives in the selected cropped views.
    3. Save cropped views, visualization images, JSON annotations, and optional YOLO labels.

This script merges:
    - the view-segmentation script for front/left view cropping;
    - the solid/dashed circle recognition script for circular primitive extraction.

It is designed as an offline semi-automatic annotation aid. Manual review is still
recommended, especially for concentric circles, dense dimension lines, leader-line
occlusion, and severely broken dashed circles.

Dependencies:
    pip install opencv-python numpy

Example on Windows:
    python semi_auto_view_circle_annotator_merged.py ^
        --input_dir "D:\\drawidfc\\盘类零件03" ^
        --output_dir "D:\\drawidfc\\pre_annotation_output" ^
        --circle_views front ^
        --save_yolo

If you also want to detect circles in the left view:
    python semi_auto_view_circle_annotator_merged.py ^
        --input_dir "D:\\drawidfc\\盘类零件03" ^
        --output_dir "D:\\drawidfc\\pre_annotation_output" ^
        --circle_views front left ^
        --save_yolo

Output structure:
    output_dir/
        crops/                 # cropped front/left view images
        visualizations/         # view boxes and circle-check images
        json_per_image/         # one JSON annotation file per source drawing
        yolo_labels/            # optional YOLO-format labels
        annotations_all.json    # merged annotation file for all drawings
"""

from __future__ import annotations

import argparse
import json
import math
import os
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import cv2
import numpy as np


# ============================================================================
# 1. Global constants and configuration
# ============================================================================

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}

# class_id is used by YOLO labels and JSON annotations.
CLASS_ID = {
    "solid_circle": 0,
    "dashed_circle": 1,
}


@dataclass
class SplitConfig:
    """Parameters for front/left view segmentation."""

    threshold: int = 200
    close_kernel: Tuple[int, int] = (25, 25)
    dilate_kernel: Tuple[int, int] = (5, 5)
    dilate_iterations: int = 2
    min_view_area_ratio: float = 0.0005
    force_split_search_range: Tuple[float, float] = (0.25, 0.75)
    padding: int = 20


@dataclass
class CircleConfig:
    """Parameters for solid/dashed circle detection."""

    threshold: int = 200
    solid_open_kernel: Tuple[int, int] = (7, 7)
    dashed_open_kernel: Tuple[int, int] = (4, 4)
    solid_mask_thickness: int = 9
    dashed_connect_kernel_ratio: float = 0.06
    dashed_connect_kernel_min: int = 21
    min_solid_contour_area: float = 20.0
    min_dashed_blob_area: float = 100.0
    min_radius: float = 5.0
    max_radius_ratio: float = 0.60
    circularity_threshold: float = 0.80
    max_angle_gap_deg: float = 160.0
    concentric_bandwidth_threshold: float = 10.0
    concentric_gap_threshold: float = 3.0
    duplicate_center_ratio: float = 0.02
    duplicate_center_min: float = 10.0
    duplicate_radius_threshold: float = 10.0
    debug_limit: int = 500


# ============================================================================
# 2. Unicode-safe image I/O
# ============================================================================


def imread_unicode(path: Path) -> Optional[np.ndarray]:
    """Read image from a path that may contain Chinese characters."""
    data = np.fromfile(str(path), dtype=np.uint8)
    if data.size == 0:
        return None
    return cv2.imdecode(data, cv2.IMREAD_COLOR)


def imwrite_unicode(path: Path, image: np.ndarray) -> bool:
    """Write image to a path that may contain Chinese characters."""
    path.parent.mkdir(parents=True, exist_ok=True)
    suffix = path.suffix if path.suffix else ".png"
    ok, buffer = cv2.imencode(suffix, image)
    if not ok:
        return False
    buffer.tofile(str(path))
    return True


def binarize_inverse(gray: np.ndarray, threshold: int) -> np.ndarray:
    """Map black drawing lines to white foreground."""
    _, binary = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY_INV)
    return binary


# ============================================================================
# 3. View segmentation: front view and left view
# ============================================================================


def split_front_left_views(
    image: np.ndarray,
    split_cfg: SplitConfig,
) -> Optional[Dict[str, Tuple[np.ndarray, Tuple[int, int, int, int]]]]:
    """
    Segment the raw drawing into front and left views.

    Returns:
        {
            "front": (front_crop, (x1, y1, x2, y2)),
            "left":  (left_crop,  (x1, y1, x2, y2)),
        }
    The bbox coordinates are in the original image coordinate system.
    """

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    binary = binarize_inverse(gray, split_cfg.threshold)

    # Morphological topological reconstruction:
    # close dashed fragments and lightly dilate thin line structures.
    close_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, split_cfg.close_kernel)
    morph = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, close_kernel)

    dilate_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, split_cfg.dilate_kernel)
    morph = cv2.dilate(morph, dilate_kernel, iterations=split_cfg.dilate_iterations)

    contours, _ = cv2.findContours(morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    img_h, img_w = image.shape[:2]
    img_area = img_h * img_w
    candidates: List[Tuple[int, int, int, int, int]] = []

    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        area = w * h
        if area > img_area * split_cfg.min_view_area_ratio:
            candidates.append((x, y, w, h, area))

    candidates.sort(key=lambda item: item[4], reverse=True)

    rect_front: Optional[Tuple[int, int, int, int]] = None
    rect_left: Optional[Tuple[int, int, int, int]] = None

    if len(candidates) >= 2:
        # Take the two largest connected view regions and sort them by x.
        two_views = [candidates[0], candidates[1]]
        two_views.sort(key=lambda item: item[0])
        rect_front = tuple(int(v) for v in two_views[0][:4])
        rect_left = tuple(int(v) for v in two_views[1][:4])

    elif len(candidates) == 1:
        # If front and left views are connected, use vertical projection and search
        # for the minimum foreground column in the central 25%-75% range.
        x, y, w, h, _ = candidates[0]
        roi = binary[y : y + h, x : x + w]
        vertical_projection = np.sum(roi, axis=0)

        start_search = int(w * split_cfg.force_split_search_range[0])
        end_search = int(w * split_cfg.force_split_search_range[1])

        if end_search > start_search:
            min_idx = int(np.argmin(vertical_projection[start_search:end_search]))
            split_x_rel = start_search + min_idx
            split_x_abs = x + split_x_rel

            rect_front = (x, y, split_x_rel, h)
            rect_left = (split_x_abs, y, w - split_x_rel, h)

    if rect_front is None or rect_left is None:
        return None

    def crop_with_padding(rect: Tuple[int, int, int, int]) -> Tuple[np.ndarray, Tuple[int, int, int, int]]:
        x, y, w, h = rect
        x1 = max(0, x - split_cfg.padding)
        y1 = max(0, y - split_cfg.padding)
        x2 = min(img_w, x + w + split_cfg.padding)
        y2 = min(img_h, y + h + split_cfg.padding)
        crop = image[y1:y2, x1:x2].copy()
        return crop, (x1, y1, x2, y2)

    front_crop, front_bbox = crop_with_padding(rect_front)
    left_crop, left_bbox = crop_with_padding(rect_left)

    return {
        "front": (front_crop, front_bbox),
        "left": (left_crop, left_bbox),
    }


# ============================================================================
# 4. Circle recognition: solid circles and dashed circles
# ============================================================================


def check_angle_coverage(points: np.ndarray, cx: float, cy: float, max_gap_thres: float = 160.0) -> bool:
    """Validate whether contour points cover a sufficiently closed circular arc."""
    if points is None or len(points) < 10:
        return False

    pts = points.reshape(-1, 2)
    angles = np.degrees(np.arctan2(pts[:, 1] - cy, pts[:, 0] - cx))
    angles = np.mod(angles, 360.0)
    angles.sort()

    gaps = np.diff(angles)
    last_gap = angles[0] + 360.0 - angles[-1]
    all_gaps = np.append(gaps, last_gap)
    return bool(np.max(all_gaps) <= max_gap_thres)


def calculate_max_angle_gap(points: np.ndarray, cx: float, cy: float) -> float:
    """Return the maximum angular gap in degrees for debugging/metadata."""
    if points is None or len(points) < 3:
        return 360.0
    pts = points.reshape(-1, 2)
    angles = np.degrees(np.arctan2(pts[:, 1] - cy, pts[:, 0] - cx))
    angles = np.mod(angles, 360.0)
    angles.sort()
    gaps = np.diff(angles)
    last_gap = angles[0] + 360.0 - angles[-1]
    return float(max(np.max(gaps), last_gap))


def fit_circle_from_points(
    points: np.ndarray,
    min_radius: float,
    max_radius: float,
    max_gap_thres: float,
) -> Optional[Tuple[int, int, int, float]]:
    """Fit a circle using ellipse fitting and validate radius + angular coverage."""
    if points is None or len(points) < 5:
        return None

    try:
        ellipse = cv2.fitEllipse(points)
        (xc, yc), (w, h), _ = ellipse
        radius = (w + h) / 4.0
    except cv2.error:
        return None

    if not (min_radius <= radius <= max_radius):
        return None

    max_gap = calculate_max_angle_gap(points, xc, yc)
    if max_gap > max_gap_thres:
        return None

    return int(round(xc)), int(round(yc)), int(round(radius)), float(max_gap)


def split_and_fit_concentric_circles(
    points: np.ndarray,
    min_radius: float,
    max_radius: float,
    circle_cfg: CircleConfig,
) -> List[Tuple[int, int, int, float]]:
    """
    Split near-concentric dashed-circle points by radial distance, then fit circles.
    This is mainly used for dashed-circle blobs after morphological connection.
    """
    if points is None or len(points) < 20:
        return []

    try:
        ellipse = cv2.fitEllipse(points)
        (cx, cy), _, _ = ellipse
    except cv2.error:
        return []

    pts_flat = points.reshape(-1, 2)
    distances = np.sqrt((pts_flat[:, 0] - cx) ** 2 + (pts_flat[:, 1] - cy) ** 2)
    p05, p95 = np.percentile(distances, 5), np.percentile(distances, 95)
    bandwidth = p95 - p05

    results: List[Tuple[int, int, int, float]] = []

    if bandwidth > circle_cfg.concentric_bandwidth_threshold:
        sorted_indices = np.argsort(distances)
        sorted_distances = distances[sorted_indices]
        diffs = np.diff(sorted_distances)
        split_idx = int(np.argmax(diffs))
        max_gap = float(diffs[split_idx]) if len(diffs) else 0.0

        if max_gap > circle_cfg.concentric_gap_threshold:
            inner_indices = sorted_indices[: split_idx + 1]
            outer_indices = sorted_indices[split_idx + 1 :]

            for group_indices in (inner_indices, outer_indices):
                if len(group_indices) > 10:
                    fitted = fit_circle_from_points(
                        points[group_indices],
                        min_radius=min_radius,
                        max_radius=max_radius,
                        max_gap_thres=circle_cfg.max_angle_gap_deg,
                    )
                    if fitted is not None:
                        results.append(fitted)
            return results

    fitted = fit_circle_from_points(
        points,
        min_radius=min_radius,
        max_radius=max_radius,
        max_gap_thres=circle_cfg.max_angle_gap_deg,
    )
    if fitted is not None:
        results.append(fitted)
    return results


def find_solid_circles(binary_img: np.ndarray, min_radius: float, max_radius: float, circle_cfg: CircleConfig) -> List[List[float]]:
    """Detect solid circles using contour circularity and enclosing circle."""
    contours, _ = cv2.findContours(binary_img, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    candidates: List[List[float]] = []

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < circle_cfg.min_solid_contour_area:
            continue

        perimeter = cv2.arcLength(cnt, True)
        if perimeter <= 0:
            continue

        circularity = 4.0 * np.pi * area / (perimeter * perimeter)
        if circularity < circle_cfg.circularity_threshold:
            continue

        (x, y), radius = cv2.minEnclosingCircle(cnt)
        if not (min_radius <= radius <= max_radius):
            continue

        candidates.append([float(x), float(y), float(radius), float(circularity)])

    return candidates


def clean_circles(circles: List[Dict[str, object]], min_img_dim: int, circle_cfg: CircleConfig) -> List[Dict[str, object]]:
    """Remove repeated circles; when duplicates are close, keep the larger one."""
    if not circles:
        return []

    center_threshold = max(circle_cfg.duplicate_center_min, min_img_dim * circle_cfg.duplicate_center_ratio)
    radius_threshold = circle_cfg.duplicate_radius_threshold

    circles_sorted = sorted(circles, key=lambda c: float(c["r"]), reverse=True)
    kept: List[Dict[str, object]] = []

    for c in circles_sorted:
        cx, cy, r = float(c["cx"]), float(c["cy"]), float(c["r"])
        duplicate = False
        for kc in kept:
            dist = math.hypot(cx - float(kc["cx"]), cy - float(kc["cy"]))
            r_diff = abs(r - float(kc["r"]))
            if dist < center_threshold and r_diff <= radius_threshold:
                duplicate = True
                break
        if not duplicate:
            kept.append(c)

    kept.sort(key=lambda c: (float(c["cy"]), float(c["cx"]), float(c["r"])))
    return kept


def make_circle_record(
    cx: float,
    cy: float,
    r: float,
    label: str,
    view_name: str,
    crop_bbox_original_xyxy: Tuple[int, int, int, int],
    view_w: int,
    view_h: int,
    score: Optional[float] = None,
    max_angle_gap_deg: Optional[float] = None,
) -> Dict[str, object]:
    """Create JSON/YOLO-friendly circle annotation record."""
    x1 = max(0.0, cx - r)
    y1 = max(0.0, cy - r)
    x2 = min(float(view_w - 1), cx + r)
    y2 = min(float(view_h - 1), cy + r)

    bbox_w = max(0.0, x2 - x1)
    bbox_h = max(0.0, y2 - y1)
    norm_x = (x1 + x2) / 2.0 / view_w
    norm_y = (y1 + y2) / 2.0 / view_h
    norm_w = bbox_w / view_w
    norm_h = bbox_h / view_h

    offset_x, offset_y = crop_bbox_original_xyxy[0], crop_bbox_original_xyxy[1]

    return {
        "label": label,
        "class_id": CLASS_ID[label],
        "view": view_name,
        "cx": float(cx),
        "cy": float(cy),
        "r": float(r),
        "bbox_xyxy": [float(x1), float(y1), float(x2), float(y2)],
        "bbox_yolo_xywh_norm": [float(norm_x), float(norm_y), float(norm_w), float(norm_h)],
        "global_cx": float(cx + offset_x),
        "global_cy": float(cy + offset_y),
        "score_or_circularity": None if score is None else float(score),
        "max_angle_gap_deg": None if max_angle_gap_deg is None else float(max_angle_gap_deg),
    }


def detect_circles_in_view(
    view_img: np.ndarray,
    view_name: str,
    crop_bbox_original_xyxy: Tuple[int, int, int, int],
    circle_cfg: CircleConfig,
) -> List[Dict[str, object]]:
    """
    Hierarchical extraction of solid and dashed circular primitives.

    Phase 1: detect solid circles.
    Phase 2: erase solid-circle regions, connect dashed fragments, and fit dashed circles.
    Phase 3: remove duplicate candidates.
    """
    view_h, view_w = view_img.shape[:2]
    min_dim = min(view_h, view_w)
    min_radius = max(circle_cfg.min_radius, min_dim * 0.005)
    max_radius = min_dim * circle_cfg.max_radius_ratio

    gray = cv2.cvtColor(view_img, cv2.COLOR_BGR2GRAY)
    binary = binarize_inverse(gray, circle_cfg.threshold)

    # Phase 1: solid circles. The 7x7 opening follows the paper description;
    # if it is too aggressive for your drawings, set --disable_solid_open.
    solid_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, circle_cfg.solid_open_kernel)
    solid_binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, solid_kernel)
    solid_candidates = find_solid_circles(solid_binary, min_radius, max_radius, circle_cfg)

    records: List[Dict[str, object]] = []
    solid_mask = np.zeros_like(binary)

    for cx, cy, r, circularity in solid_candidates:
        records.append(
            make_circle_record(
                cx=cx,
                cy=cy,
                r=r,
                label="solid_circle",
                view_name=view_name,
                crop_bbox_original_xyxy=crop_bbox_original_xyxy,
                view_w=view_w,
                view_h=view_h,
                score=circularity,
            )
        )
        cv2.circle(
            solid_mask,
            (int(round(cx)), int(round(cy))),
            int(round(r)),
            255,
            thickness=circle_cfg.solid_mask_thickness,
        )

    # Phase 2: dashed circles after removing solid-circle regions.
    dashed_raw = cv2.subtract(binary, solid_mask)

    dashed_kernel_size = max(circle_cfg.dashed_connect_kernel_min, int(min_dim * circle_cfg.dashed_connect_kernel_ratio))
    if dashed_kernel_size % 2 == 0:
        dashed_kernel_size += 1

    dashed_connect_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dashed_kernel_size, dashed_kernel_size))
    dashed_connected = cv2.dilate(dashed_raw, dashed_connect_kernel, iterations=1)

    contours, _ = cv2.findContours(dashed_connected, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < circle_cfg.min_dashed_blob_area:
            continue

        mask_blob = np.zeros_like(binary)
        cv2.drawContours(mask_blob, [cnt], -1, 255, -1)
        raw_pixels = cv2.bitwise_and(dashed_raw, dashed_raw, mask=mask_blob)
        points = cv2.findNonZero(raw_pixels)

        if points is None:
            continue

        fitted_circles = split_and_fit_concentric_circles(
            points=points,
            min_radius=min_radius,
            max_radius=max_radius,
            circle_cfg=circle_cfg,
        )

        for cx, cy, r, max_gap in fitted_circles:
            records.append(
                make_circle_record(
                    cx=cx,
                    cy=cy,
                    r=r,
                    label="dashed_circle",
                    view_name=view_name,
                    crop_bbox_original_xyxy=crop_bbox_original_xyxy,
                    view_w=view_w,
                    view_h=view_h,
                    max_angle_gap_deg=max_gap,
                )
            )

    # Phase 3: deduplication.
    return clean_circles(records, min_dim, circle_cfg)


# ============================================================================
# 5. Visualization and label export
# ============================================================================


def draw_view_boxes(image: np.ndarray, views: Dict[str, Tuple[np.ndarray, Tuple[int, int, int, int]]]) -> np.ndarray:
    """Draw front/left view boxes on the original drawing."""
    canvas = image.copy()
    colors = {"front": (0, 180, 0), "left": (255, 0, 0)}

    for view_name, (_, bbox) in views.items():
        x1, y1, x2, y2 = bbox
        color = colors.get(view_name, (0, 0, 255))
        cv2.rectangle(canvas, (x1, y1), (x2, y2), color, 3)
        cv2.putText(canvas, view_name, (x1, max(25, y1 - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.85, color, 2, cv2.LINE_AA)

    return canvas


def draw_circle_annotations(view_img: np.ndarray, circles: Sequence[Dict[str, object]]) -> np.ndarray:
    """Draw detected circles on a cropped view."""
    canvas = view_img.copy()
    colors = {"solid_circle": (0, 180, 0), "dashed_circle": (0, 0, 255)}

    for idx, circle in enumerate(circles, start=1):
        label = str(circle["label"])
        color = colors.get(label, (255, 0, 0))
        cx, cy, r = int(round(float(circle["cx"]))), int(round(float(circle["cy"]))), int(round(float(circle["r"])))
        cv2.circle(canvas, (cx, cy), r, color, 2)
        cv2.circle(canvas, (cx, cy), 2, color, -1)
        tag = "S" if label == "solid_circle" else "D"
        cv2.putText(canvas, f"{idx}:{tag}", (cx + 4, cy - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 1, cv2.LINE_AA)

    return canvas


def save_yolo_labels(label_path: Path, circles: Sequence[Dict[str, object]]) -> None:
    """Save one YOLO-format .txt file for one cropped view image."""
    label_path.parent.mkdir(parents=True, exist_ok=True)
    lines: List[str] = []

    for circle in circles:
        class_id = int(circle["class_id"])
        x, y, w, h = [float(v) for v in circle["bbox_yolo_xywh_norm"]]
        lines.append(f"{class_id} {x:.6f} {y:.6f} {w:.6f} {h:.6f}")

    label_path.write_text("\n".join(lines), encoding="utf-8")


# ============================================================================
# 6. Main processing logic
# ============================================================================


def process_one_drawing(
    image_path: Path,
    output_dir: Path,
    split_cfg: SplitConfig,
    circle_cfg: CircleConfig,
    circle_views: Sequence[str],
    save_yolo: bool,
    image_index: int,
) -> Optional[Dict[str, object]]:
    """Process one raw drawing and return its annotation dictionary."""
    image = imread_unicode(image_path)
    if image is None:
        print(f"[FAIL] Cannot read image: {image_path}")
        return None

    views = split_front_left_views(image, split_cfg)
    if views is None:
        print(f"[FAIL] Cannot locate two views: {image_path.name}")
        return None

    stem = image_path.stem
    crop_dir = output_dir / "crops"
    vis_dir = output_dir / "visualizations"
    json_dir = output_dir / "json_per_image"
    yolo_dir = output_dir / "yolo_labels"

    crop_dir.mkdir(parents=True, exist_ok=True)
    vis_dir.mkdir(parents=True, exist_ok=True)
    json_dir.mkdir(parents=True, exist_ok=True)

    view_box_vis = draw_view_boxes(image, views)
    imwrite_unicode(vis_dir / f"{stem}_view_split.png", view_box_vis)

    annotation: Dict[str, object] = {
        "source_image": str(image_path),
        "image_name": image_path.name,
        "image_size": {"width": int(image.shape[1]), "height": int(image.shape[0])},
        "classes": CLASS_ID,
        "views": {},
        "review_note": (
            "This annotation was generated by a semi-automatic morphological and geometric fitting tool. "
            "Manual correction is recommended for concentric circles, leader-line occlusion, dense dimension areas, "
            "and severely broken dashed circles."
        ),
    }

    for view_name, (crop_img, bbox_original) in views.items():
        crop_path = crop_dir / f"{stem}_{view_name}.png"
        imwrite_unicode(crop_path, crop_img)

        circles: List[Dict[str, object]] = []
        if view_name in circle_views:
            circles = detect_circles_in_view(crop_img, view_name, bbox_original, circle_cfg)
            circle_vis = draw_circle_annotations(crop_img, circles)
            if image_index < circle_cfg.debug_limit:
                imwrite_unicode(vis_dir / f"{stem}_{view_name}_circles.png", circle_vis)

            if save_yolo:
                save_yolo_labels(yolo_dir / view_name / f"{stem}_{view_name}.txt", circles)

        annotation["views"][view_name] = {
            "crop_path": str(crop_path),
            "bbox_original_xyxy": [int(v) for v in bbox_original],
            "circles": circles,
        }

    per_image_json = json_dir / f"{stem}.json"
    per_image_json.write_text(json.dumps(annotation, ensure_ascii=False, indent=2), encoding="utf-8")

    front_count = len(annotation["views"].get("front", {}).get("circles", []))
    left_count = len(annotation["views"].get("left", {}).get("circles", []))
    print(f"[OK] {image_path.name}: front circles={front_count}, left circles={left_count}")

    return annotation


def iter_images(input_dir: Path) -> Iterable[Path]:
    """Yield all supported image files in the input folder."""
    for path in sorted(input_dir.iterdir()):
        if path.is_file() and path.suffix.lower() in IMAGE_EXTS:
            yield path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Split front/left views and detect solid/dashed circles for semi-automatic annotation."
    )

    parser.add_argument("--input_dir", required=True, help="Folder containing raw engineering drawing images.")
    parser.add_argument("--output_dir", required=True, help="Folder to save crops, JSON annotations, and visualizations.")
    parser.add_argument(
        "--circle_views",
        nargs="+",
        default=["front"],
        choices=["front", "left"],
        help="Which cropped views should be processed for circle detection. Default: front.",
    )
    parser.add_argument("--save_yolo", action="store_true", help="Save YOLO-format label files for detected circles.")

    # Commonly tuned parameters.
    parser.add_argument("--threshold", type=int, default=200, help="Binarization threshold. Default: 200.")
    parser.add_argument("--padding", type=int, default=20, help="Padding around cropped views. Default: 20.")
    parser.add_argument("--min_radius", type=float, default=5.0, help="Minimum valid circle radius. Default: 5.")
    parser.add_argument("--debug_limit", type=int, default=500, help="Maximum number of images for circle visualization output.")

    return parser.parse_args()


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

    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if not input_dir.exists():
        raise FileNotFoundError(f"Input folder does not exist: {input_dir}")

    split_cfg = SplitConfig(
        threshold=args.threshold,
        padding=args.padding,
    )
    circle_cfg = CircleConfig(
        threshold=args.threshold,
        min_radius=args.min_radius,
        debug_limit=args.debug_limit,
    )

    all_annotations: Dict[str, object] = {
        "tool_name": "semi_auto_view_circle_annotator_merged",
        "class_id": CLASS_ID,
        "split_config": asdict(split_cfg),
        "circle_config": asdict(circle_cfg),
        "circle_views": list(args.circle_views),
        "items": {},
    }

    total = 0
    success = 0

    for idx, image_path in enumerate(iter_images(input_dir)):
        total += 1
        annotation = process_one_drawing(
            image_path=image_path,
            output_dir=output_dir,
            split_cfg=split_cfg,
            circle_cfg=circle_cfg,
            circle_views=args.circle_views,
            save_yolo=args.save_yolo,
            image_index=idx,
        )
        if annotation is not None:
            success += 1
            all_annotations["items"][image_path.name] = annotation

    merged_json_path = output_dir / "annotations_all.json"
    merged_json_path.write_text(json.dumps(all_annotations, ensure_ascii=False, indent=2), encoding="utf-8")

    print("\nFinished.")
    print(f"Success: {success}/{total}")
    print(f"Merged JSON: {merged_json_path}")
    print(f"Output folder: {output_dir}")


if __name__ == "__main__":
    main()