import rclpy

from rclpy.node import Node
from sensor_msgs.msg import JointState
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint


class TrajectoryComposer(Node):
    def __init__(self):
        super().__init__("trajectory_composer")
        self.currentConfiguration = JointTrajectoryPoint()
        self.jointStateSubscription = self.create_subscription(JointState, 'joint_states',
                                                               self.joint_states_callback,
                                                               10)
        self.goalSubscription = self.create_subscription(JointTrajectory, 'goal_trajectory',
                                                         self.goal_trajectory_callback,
                                                         10)
        self.trajectoryGoalPublisher = self.create_publisher(JointTrajectory,
                                                             'bridge/follow_joint_trajectory/goal', 10)

    def joint_states_callback(self, joint_states):
        self.get_logger().info('JointState message received')
        self.currentConfiguration.positions = joint_states.position

    def goal_trajectory_callback(self, goal_trajectory):
        trajectory = JointTrajectory()
        trajectory.header.frame_id = 'base_link'
        trajectory.header.stamp = self.get_clock().now().to_msg()
        trajectory.joint_names = ['J1', 'J2', 'J3', 'J4', 'J5', 'J6']

        if goal_trajectory.points[0].time_from_start.sec == 0 and goal_trajectory.points[
            0].time_from_start.nanosec == 0:
            self.get_logger().info('Wrong')
            trajectory.points = goal_trajectory.points
        else:
            self.get_logger().info('Right')
            trajectory.points = goal_trajectory.points
            trajectory.points.insert(0, self.currentConfiguration)

        self.trajectoryGoalPublisher.publish(trajectory)
        self.get_logger().info('Published JointTrajectory message')


def main(args=None):
    rclpy.init(args=args)
    node = TrajectoryComposer()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
