#!/usr/bin/env python3
"""Capture paired camera and white-line-mask screenshots from an existing ROS topic."""

from pathlib import Path
import threading
import time

import cv2
import numpy as np
import rospy
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image


class SnapshotCapture:
    """A subscriber-only capture tool: it never opens /dev/video* itself."""

    def __init__(self):
        rospy.init_node("line_snapshot_capture")

        self.image_topic = rospy.get_param("~image_topic", "/usb_cam/image_raw")
        default_root = Path(__file__).resolve().parent.parent / "data"
        self.save_root = Path(rospy.get_param("~save_root", str(default_root))).expanduser()
        self.lower = self._read_hsv("~hsv_lower", [0, 0, 180])
        self.upper = self._read_hsv("~hsv_upper", [180, 60, 255])

        self.raw_dir = self.save_root / "raw"
        self.mask_dir = self.save_root / "mask"
        self.raw_dir.mkdir(parents=True, exist_ok=True)
        self.mask_dir.mkdir(parents=True, exist_ok=True)

        self.bridge = CvBridge()
        self.lock = threading.Lock()
        self.raw_image = None
        self.mask_image = None

        self.subscriber = rospy.Subscriber(
            self.image_topic, Image, self.image_callback, queue_size=1
        )
        rospy.loginfo(
            "Line snapshot capture subscribes to %s; saving pairs under %s",
            self.image_topic,
            self.save_root,
        )

    @staticmethod
    def _read_hsv(param_name, default):
        values = rospy.get_param(param_name, default)
        if not isinstance(values, (list, tuple)) or len(values) != 3:
            rospy.logwarn("%s must contain three numbers; using %s", param_name, default)
            values = default
        values = [int(max(0, min(255, value))) for value in values]
        values[0] = min(180, values[0])
        return np.array(values, dtype=np.uint8)

    def image_callback(self, message):
        try:
            raw = self.bridge.imgmsg_to_cv2(message, desired_encoding="bgr8")
        except CvBridgeError as error:
            rospy.logwarn_throttle(2.0, "Unable to convert camera image: %s", error)
            return

        hsv = cv2.cvtColor(raw, cv2.COLOR_BGR2HSV)
        mask = cv2.inRange(hsv, self.lower, self.upper)
        with self.lock:
            self.raw_image = raw.copy()
            self.mask_image = mask.copy()

    def save_pair(self):
        with self.lock:
            if self.raw_image is None or self.mask_image is None:
                rospy.logwarn("No camera frame received yet; nothing was saved.")
                return
            raw = self.raw_image.copy()
            mask = self.mask_image.copy()

        timestamp = time.time_ns()
        raw_path = self.raw_dir / ("line_%d_raw.png" % timestamp)
        mask_path = self.mask_dir / ("line_%d_mask.png" % timestamp)
        raw_ok = cv2.imwrite(str(raw_path), raw)
        mask_ok = cv2.imwrite(str(mask_path), mask)
        if raw_ok and mask_ok:
            rospy.loginfo("Saved snapshot pair: raw=%s mask=%s", raw_path, mask_path)
        else:
            rospy.logerr("Failed to save snapshot pair: raw=%s mask=%s", raw_path, mask_path)

    def run(self):
        window_name = "Line follower snapshot capture (S: save, Q/Esc: quit)"
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
        rate = rospy.Rate(30)

        while not rospy.is_shutdown():
            with self.lock:
                raw = None if self.raw_image is None else self.raw_image.copy()
                mask = None if self.mask_image is None else self.mask_image.copy()

            if raw is not None and mask is not None:
                mask_bgr = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
                display = np.hstack((raw, mask_bgr))
                width = raw.shape[1]
                cv2.putText(display, "RAW: %s" % self.image_topic, (12, 28),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                cv2.putText(display, "WHITE MASK HSV %s-%s" % (self.lower.tolist(), self.upper.tolist()),
                            (width + 12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                cv2.putText(display, "S: save pair    Q/Esc: quit", (12, display.shape[0] - 16),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 255, 255), 2)
                cv2.imshow(window_name, display)

            key = cv2.waitKey(1) & 0xFF
            if key in (ord("s"), ord("S")):
                self.save_pair()
            elif key in (ord("q"), ord("Q"), 27):
                rospy.loginfo("Line snapshot capture stopped by keyboard.")
                break
            rate.sleep()

        cv2.destroyAllWindows()


if __name__ == "__main__":
    try:
        SnapshotCapture().run()
    except rospy.ROSInterruptException:
        pass