#!/usr/bin/env python3
"""Subscribe /feedback/joint_states and preview camera."""

import math

import cv2
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState

FEEDBACK_TOPIC = "/feedback/joint_states"
CONTROL_TOPIC = "/control/joint_states"
CAMERA_INDEX = 4
TARGET_NAMES = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "gripper"]
GRIPPER_MIN_M = 0.0
GRIPPER_MAX_M = 0.1
# Motion amplitudes (radians / meters). Increase these to make movement larger.
JOINT2_DELTA_RAD = 0.12
JOINT3_DELTA_RAD = 0.12
JOINT4_DELTA_RAD = 0.06
GRIPPER_DELTA_M = 0.05


class JointStatePrinter(Node):
    def __init__(self) -> None:
        super().__init__("joint_state_printer")
        self.latest_msg: JointState | None = None
        self.latest_frame_shape: tuple[int, int, int] | None = None
        self.show_window = False
        self.window_name = "camera_preview"

        self.cap = cv2.VideoCapture(CAMERA_INDEX, cv2.CAP_V4L2)
        if self.cap.isOpened():
            self.get_logger().info(f"OpenCV camera opened: /dev/video{CAMERA_INDEX}")
            try:
                cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
                cv2.resizeWindow(self.window_name, 1280, 720)
                self.show_window = True
            except cv2.error:
                self.get_logger().warn("No GUI backend, run without preview window.")
        else:
            self.get_logger().warn(f"OpenCV camera open failed: /dev/video{CAMERA_INDEX}")

        self.ctrl_pub = self.create_publisher(JointState, CONTROL_TOPIC, 10)
        self.traj_msgs: list[JointState] = []
        self.traj_idx = 0
        self.traj_cycle = 0
        self.traj_started = False

        self.create_subscription(JointState, FEEDBACK_TOPIC, self._callback, 10)
        self.create_timer(0.05, self._read_camera_frame)
        self.create_timer(0.05, self._publish_traj_step)  # 20Hz
        self.create_timer(1.0, self._print_latest)
        self.get_logger().info(
            f"Subscribed to {FEEDBACK_TOPIC}; publish planned stream to {CONTROL_TOPIC} at 20Hz"
        )

    def _callback(self, msg: JointState) -> None:
        self.latest_msg = msg
        if self.traj_started:
            return
        base = self._extract_target_positions(msg)
        if base is None:
            self.get_logger().warn("Feedback JointState missing required joints, waiting next message.")
            return
        self.traj_msgs = self._build_trajectory(base, steps=50)
        self.traj_idx = 0
        self.traj_cycle = 0
        self.traj_started = True
        self.get_logger().info("Generated 50-step motion stream from current feedback position (loop mode).")

    def _extract_target_positions(self, msg: JointState) -> list[float] | None:
        name_to_pos = {n: p for n, p in zip(msg.name, msg.position)}
        out: list[float] = []
        for joint in TARGET_NAMES[:6]:
            if joint not in name_to_pos:
                return None
            out.append(float(name_to_pos[joint]))
        if "gripper" in name_to_pos:
            out.append(float(name_to_pos["gripper"]))
        elif "gripper_joint1" in name_to_pos and "gripper_joint2" in name_to_pos:
            out.append(abs(float(name_to_pos["gripper_joint1"])) + abs(float(name_to_pos["gripper_joint2"])))
        else:
            out.append(0.05)
        out[6] = max(GRIPPER_MIN_M, min(GRIPPER_MAX_M, out[6]))
        return out

    def _build_trajectory(self, base: list[float], steps: int) -> list[JointState]:
        traj: list[JointState] = []
        for i in range(steps):
            t = i / max(steps - 1, 1)
            s = math.sin(2.0 * math.pi * t)
            s2 = math.sin(4.0 * math.pi * t)
            p = list(base)
            # Larger adjustments around current pose.
            p[1] += JOINT2_DELTA_RAD * s
            p[2] += -JOINT3_DELTA_RAD * s
            p[3] += JOINT4_DELTA_RAD * s2
            # Gripper opening width in meters, clamped to [0, 0.1].
            p[6] = max(GRIPPER_MIN_M, min(GRIPPER_MAX_M, p[6] + GRIPPER_DELTA_M * s))

            m = JointState()
            m.name = TARGET_NAMES
            m.position = p
            m.velocity = []
            m.effort = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5]
            traj.append(m)
        return traj

    def _publish_traj_step(self) -> None:
        if not self.traj_started or not self.traj_msgs:
            return
        if self.traj_idx >= len(self.traj_msgs):
            self.traj_idx = 0
            self.traj_cycle += 1
            self.get_logger().info(f"Looping motion stream, cycle={self.traj_cycle}")
        msg = self.traj_msgs[self.traj_idx]
        msg.header.stamp = self.get_clock().now().to_msg()
        self.ctrl_pub.publish(msg)
        self.traj_idx += 1

    def _read_camera_frame(self) -> None:
        if not self.cap.isOpened():
            return
        ok, frame = self.cap.read()
        if not ok:
            return
        self.latest_frame_shape = frame.shape
        if self.show_window:
            cv2.imshow(self.window_name, frame)
            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                self.get_logger().info("Pressed q, shutting down...")
                if rclpy.ok():
                    rclpy.shutdown()

    def _print_latest(self) -> None:
        if self.latest_msg is None:
            return
        msg = self.latest_msg
        stamp = f"{msg.header.stamp.sec}.{msg.header.stamp.nanosec:09d}"
        print("==== /feedback/joint_states ====")
        print(f"stamp: {stamp}")
        print(f"name: {list(msg.name)}")
        print(f"position: {list(msg.position)}")
        print(f"velocity: {list(msg.velocity)}")
        print(f"effort: {list(msg.effort)}")
        if self.traj_started:
            print(f"stream_progress: {self.traj_idx}/{len(self.traj_msgs)} cycle={self.traj_cycle}")
        if self.latest_frame_shape is not None:
            h, w, c = self.latest_frame_shape
            print(f"camera_frame: {w}x{h} channels={c}")
        else:
            print("camera_frame: N/A")


def main() -> None:
    rclpy.init()
    node = JointStatePrinter()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        if node.cap.isOpened():
            node.cap.release()
            if node.show_window:
                cv2.destroyAllWindows()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    main()