#!/usr/bin/env python3
import time
from pathlib import Path

import cv2
import rclpy
from builtin_interfaces.msg import Duration
from cv_bridge import CvBridge
from rclpy.node import Node
from rclpy.qos import qos_profile_sensor_data
from sensor_msgs.msg import Image
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from ultralytics import YOLO


HERE = Path(__file__).resolve().parent
CAMERA = "/RMC1/arm95/svcam/right/image"
GRIPPER = "/RMC1/arm95/gripper_trajectory_controller/joint_trajectory"


class Demo(Node):
    def __init__(self):
        super().__init__("simple_module_b")
        self.model = YOLO(str(HERE / "models/latest.pt"))
        self.bridge = CvBridge()
        self.gripper = self.create_publisher(JointTrajectory, GRIPPER, 10)
        self.create_subscription(Image, CAMERA, self.on_image, qos_profile_sensor_data)
        self.started = None
        self.last_log = -1
        self.step = 0

    def move_gripper(self, position, state):
        message = JointTrajectory()
        message.joint_names = ["left_joint", "right_joint"]
        point = JointTrajectoryPoint()
        point.positions = [position, position]
        point.time_from_start = Duration(sec=1)
        message.points = [point]
        self.gripper.publish(message)
        print(f"[GRIPPER] {state}: {position}", flush=True)

    def on_image(self, message):
        frame = self.bridge.imgmsg_to_cv2(message, "bgr8")
        result = self.model(frame, conf=0.25, verbose=False)[0]

        if self.started is None:
            self.started = time.monotonic()
            print(f"[CAMERA] {CAMERA}", flush=True)
        seconds = int(time.monotonic() - self.started)

        if seconds != self.last_log:
            print(f"[YOLO] {result.verbose().strip()}", flush=True)
            self.last_log = seconds

        if seconds >= 5 and self.step == 0:
            self.move_gripper(0.04, "OPEN")
            self.step = 1
        elif seconds >= 7 and self.step == 1:
            self.move_gripper(0.02, "TAKE_CUBE")
            self.step = 2

        cv2.imshow("YOLO", result.plot())
        if cv2.waitKey(1) & 0xFF in (ord("q"), 27):
            rclpy.shutdown()


def main():
    target = input("Предмет (Кисть, Нож или Резаки): ")
    print(f"[START] {target}", flush=True)
    rclpy.init()
    node = Demo()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        cv2.destroyAllWindows()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    main()