#!/usr/bin/env python3
"""
capture_calibration_images.py — Capture checkerboard calibration images from IMX219 cameras.

Usage (standalone, no ROS2 required — uses OpenCV VideoCapture for CSI):
  python3 capture_calibration_images.py --camera front --n-images 30 --output-dir /tmp/cal

Usage (with ROS2 — subscribes to image topic):
  python3 capture_calibration_images.py --ros --camera front --n-images 30

Controls:
  SPACE  — capture current frame (if checkerboard detected)
  q      — quit / done
  r      — reset captured images for this camera

Outputs: output_dir/cam_{name}/frame_{NNNN}.png
"""

import argparse
import os
import sys
from pathlib import Path

import cv2
import numpy as np


CAMERA_TOPICS = {
    'front': '/camera/front/image_raw',
    'right': '/camera/right/image_raw',
    'rear':  '/camera/rear/image_raw',
    'left':  '/camera/left/image_raw',
}

ALL_CAMERAS = ['front', 'right', 'rear', 'left']


def parse_board_size(s: str) -> tuple:
    """Parse 'WxH' string to (cols, rows) tuple of ints."""
    parts = s.lower().split('x')
    if len(parts) != 2:
        raise argparse.ArgumentTypeError(f'board-size must be WxH, got: {s}')
    return (int(parts[0]), int(parts[1]))


def find_checkerboard(img_bgr: np.ndarray, board_size: tuple):
    """Detect checkerboard corners. Returns (found, corners, gray)."""
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    flags = cv2.CALIB_CB_ADAPTIVE_THRESH | cv2.CALIB_CB_NORMALIZE_IMAGE
    found, corners = cv2.findChessboardCorners(gray, board_size, flags)
    if found:
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        corners = cv2.cornerSubPix(gray, corners, (7, 7), (-1, -1), criteria)
    return found, corners, gray


def draw_overlay(img: np.ndarray, board_size: tuple, found: bool,
                 corners, captured: int, n_images: int) -> np.ndarray:
    """Draw corner detections and status overlay onto a copy of img."""
    display = img.copy()
    if found and corners is not None:
        cv2.drawChessboardCorners(display, board_size, corners, found)

    # Status bar background
    cv2.rectangle(display, (0, 0), (display.shape[1], 36), (0, 0, 0), -1)

    status_color = (0, 255, 0) if found else (0, 80, 220)
    status_text = 'BOARD DETECTED' if found else 'searching for board...'
    cv2.putText(display, status_text, (10, 24),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, status_color, 2)

    count_text = f'Captured: {captured}/{n_images}  [SPACE=capture  r=reset  q=quit]'
    cv2.putText(display, count_text, (display.shape[1] // 2 - 10, 24),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (200, 200, 200), 1)

    if captured >= n_images:
        cv2.rectangle(display, (0, 0), (display.shape[1], display.shape[0]),
                      (0, 200, 0), 8)
        cv2.putText(display, 'DONE! Press q to finish.', (60, display.shape[0] // 2),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 3)

    return display


def capture_loop_opencv(cam_name: str, device: str, board_size: tuple,
                        n_images: int, output_dir: Path) -> int:
    """Capture calibration images using cv2.VideoCapture (non-ROS mode).

    Returns number of images captured.
    """
    out_dir = output_dir / f'cam_{cam_name}'
    out_dir.mkdir(parents=True, exist_ok=True)

    cap = cv2.VideoCapture(device)
    if not cap.isOpened():
        print(f'[ERROR] Cannot open device: {device}')
        return 0

    # Try to set resolution
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    captured = 0
    frame_idx = 0
    window_title = f'Calibration Capture — {cam_name} | SPACE=capture q=quit r=reset'
    cv2.namedWindow(window_title, cv2.WINDOW_NORMAL)

    print(f'\n[{cam_name.upper()}] Capturing {n_images} images. Output: {out_dir}')
    print('  Point the checkerboard at the camera, press SPACE when board is detected.')

    while captured < n_images:
        ret, frame = cap.read()
        if not ret:
            print('[WARN] Failed to read frame')
            continue

        found, corners, _ = find_checkerboard(frame, board_size)
        display = draw_overlay(frame, board_size, found, corners, captured, n_images)
        cv2.imshow(window_title, display)

        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            print(f'  Quit requested. Captured {captured}/{n_images} images.')
            break
        elif key == ord('r'):
            # Remove all saved images and reset counter
            for f in out_dir.glob('frame_*.png'):
                f.unlink()
            captured = 0
            frame_idx = 0
            print(f'  [RESET] Cleared captured images for {cam_name}.')
        elif key == ord(' '):
            if found:
                save_path = out_dir / f'frame_{frame_idx:04d}.png'
                cv2.imwrite(str(save_path), frame)
                captured += 1
                frame_idx += 1
                print(f'  Saved ({captured}/{n_images}): {save_path.name}')
                # Brief flash to confirm capture
                flash = frame.copy()
                cv2.rectangle(flash, (0, 0), (flash.shape[1], flash.shape[0]),
                              (0, 255, 0), 12)
                cv2.imshow(window_title, flash)
                cv2.waitKey(120)
            else:
                print('  [SKIP] No checkerboard detected — move board into view.')

    cap.release()
    cv2.destroyWindow(window_title)
    print(f'  Done: {captured} images saved to {out_dir}')
    return captured


def capture_loop_ros(cam_name: str, board_size: tuple,
                     n_images: int, output_dir: Path) -> int:
    """Capture calibration images via ROS2 topic subscription.

    Returns number of images captured.
    """
    try:
        import rclpy
        from rclpy.node import Node
        from sensor_msgs.msg import Image
        from cv_bridge import CvBridge
    except ImportError as e:
        print(f'[ERROR] ROS2 not available: {e}')
        print('  Run without --ros for standalone OpenCV mode.')
        return 0

    out_dir = output_dir / f'cam_{cam_name}'
    out_dir.mkdir(parents=True, exist_ok=True)
    topic = CAMERA_TOPICS[cam_name]

    print(f'\n[{cam_name.upper()}] Subscribing to: {topic}')
    print(f'  Capturing {n_images} images. Output: {out_dir}')
    print('  Press SPACE in the preview window when the board is detected.')

    rclpy.init()
    node = Node('calibration_capture')
    bridge = CvBridge()

    captured = [0]
    frame_idx = [0]
    latest_frame = [None]

    def image_callback(msg):
        try:
            latest_frame[0] = bridge.imgmsg_to_cv2(msg, 'bgr8')
        except Exception as e:
            node.get_logger().warn(f'cv_bridge error: {e}')

    sub = node.create_subscription(Image, topic, image_callback, 1)  # noqa: F841
    window_title = f'Calibration Capture — {cam_name} (ROS2)'
    cv2.namedWindow(window_title, cv2.WINDOW_NORMAL)

    import threading
    spin_thread = threading.Thread(target=rclpy.spin, args=(node,), daemon=True)
    spin_thread.start()

    try:
        while captured[0] < n_images:
            frame = latest_frame[0]
            if frame is None:
                cv2.waitKey(50)
                continue

            found, corners, _ = find_checkerboard(frame, board_size)
            display = draw_overlay(frame, board_size, found, corners,
                                   captured[0], n_images)
            cv2.imshow(window_title, display)

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                print(f'  Quit. Captured {captured[0]}/{n_images}.')
                break
            elif key == ord('r'):
                for f in out_dir.glob('frame_*.png'):
                    f.unlink()
                captured[0] = 0
                frame_idx[0] = 0
                print(f'  [RESET] Cleared images for {cam_name}.')
            elif key == ord(' '):
                if found:
                    save_path = out_dir / f'frame_{frame_idx[0]:04d}.png'
                    cv2.imwrite(str(save_path), frame)
                    captured[0] += 1
                    frame_idx[0] += 1
                    print(f'  Saved ({captured[0]}/{n_images}): {save_path.name}')
                else:
                    print('  [SKIP] No board detected.')
    finally:
        cv2.destroyWindow(window_title)
        node.destroy_node()
        rclpy.shutdown()

    print(f'  Done: {captured[0]} images saved to {out_dir}')
    return captured[0]


def main():
    parser = argparse.ArgumentParser(
        description='Capture checkerboard images for IMX219 fisheye calibration.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        '--camera', choices=['front', 'right', 'rear', 'left', 'all'],
        default='front',
        help='Which camera to capture from (default: front). Use "all" to iterate all 4.',
    )
    parser.add_argument(
        '--n-images', type=int, default=30,
        help='Number of images to capture per camera (default: 30).',
    )
    parser.add_argument(
        '--output-dir', type=str, default='/tmp/saltybot_cal',
        help='Root output directory. Images saved to {output_dir}/cam_{name}/ (default: /tmp/saltybot_cal).',
    )
    parser.add_argument(
        '--board-size', type=str, default='8x6',
        help='Checkerboard inner corners as WxH (default: 8x6).',
    )
    parser.add_argument(
        '--square-size', type=float, default=0.025,
        help='Physical square size in metres (default: 0.025 = 25mm).',
    )
    parser.add_argument(
        '--ros', action='store_true',
        help='Use ROS2 topic subscription instead of direct VideoCapture.',
    )
    parser.add_argument(
        '--device', type=str, default='/dev/video0',
        help='Video device (non-ROS mode only, default: /dev/video0).',
    )
    args = parser.parse_args()

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

    cameras = ALL_CAMERAS if args.camera == 'all' else [args.camera]

    print(f'Checkerboard: {board_size[0]}x{board_size[1]} inner corners, '
          f'{args.square_size * 1000:.0f}mm squares')
    print(f'Target: {args.n_images} images per camera')
    print(f'Output: {output_dir}')
    print(f'Mode: {"ROS2 topics" if args.ros else f"VideoCapture ({args.device})"}')

    for i, cam_name in enumerate(cameras):
        if len(cameras) > 1:
            print(f'\n--- Camera {i + 1}/{len(cameras)}: {cam_name.upper()} ---')
            if i > 0:
                input('  Position the board in view of this camera, then press ENTER...')

        if args.ros:
            n_captured = capture_loop_ros(cam_name, board_size, args.n_images, output_dir)
        else:
            n_captured = capture_loop_opencv(cam_name, args.device, board_size,
                                             args.n_images, output_dir)

        if n_captured < 5:
            print(f'[WARN] Only {n_captured} images captured for {cam_name} — '
                  f'calibration may fail (need >= 5, recommend >= 20).')

    print('\nCapture complete. Run calibrate_camera.py for each camera to compute intrinsics.')
    print('Example:')
    for cam in cameras:
        print(f'  python3 calibrate_camera.py \\')
        print(f'    --images-dir {output_dir}/cam_{cam} \\')
        print(f'    --output /mnt/nvme/saltybot/calibration/{cam}/camera_info.yaml \\')
        print(f'    --camera-name camera_{cam} --board-size {args.board_size} \\')
        print(f'    --square-size {args.square_size} --image-size 640x480')


if __name__ == '__main__':
    main()
