#! /usr/bin/env python3

import rclpy
from rclpy.node import Node
from tf2_msgs.msg import TFMessage
from geometry_msgs.msg import TransformStamped
from rbs_utils_interfaces.msg import AssemblyConfig
from tf2_ros.transform_broadcaster import TransformBroadcaster


class TF2PublisherNode(Node):

    def __init__(self):
        super().__init__('tf2_publisher_node')
        self.tfb_ = TransformBroadcaster(self)
        self.subscription = self.create_subscription(
            AssemblyConfig, '/assembly_config', self.config_callback, 10)

    def config_callback(self, msg):
        # Populate transforms for absolute_part
        for part in msg.absolute_part:
            ts = TransformStamped()
            ts.header.stamp = self.get_clock().now().to_msg()
            ts.header.frame_id = "world"
            ts.child_frame_id = part.name
            ts.transform.translation.x = part.pose.position.x
            ts.transform.translation.y = part.pose.position.y
            ts.transform.translation.z = part.pose.position.z
            ts.transform.rotation.x = part.pose.orientation.x
            ts.transform.rotation.y = part.pose.orientation.y
            ts.transform.rotation.z = part.pose.orientation.z
            ts.transform.rotation.w = part.pose.orientation.w
            self.tfb_.sendTransform(ts)

        # Populate transforms for relative_part
        for part in msg.relative_part:
            ts = TransformStamped()
            ts.header.stamp = self.get_clock().now().to_msg()
            ts.header.frame_id = part.relative_at
            ts.child_frame_id = part.name
            ts.transform.translation.x = part.pose.position.x
            ts.transform.translation.y = part.pose.position.y
            ts.transform.translation.z = part.pose.position.z
            ts.transform.rotation.x = part.pose.orientation.x
            ts.transform.rotation.y = part.pose.orientation.y
            ts.transform.rotation.z = part.pose.orientation.z
            ts.transform.rotation.w = part.pose.orientation.w
            self.tfb_.sendTransform(ts)

        # Populate transforms for grasp_pose
        for part in msg.grasp_pose:
            ts = TransformStamped()
            ts.header.stamp = self.get_clock().now().to_msg()
            ts.header.frame_id = part.relative_at
            ts.child_frame_id = part.name
            ts.transform.translation.x = part.pose.position.x
            ts.transform.translation.y = part.pose.position.y
            ts.transform.translation.z = part.pose.position.z
            ts.transform.rotation.x = part.pose.orientation.x
            ts.transform.rotation.y = part.pose.orientation.y
            ts.transform.rotation.z = part.pose.orientation.z
            ts.transform.rotation.w = part.pose.orientation.w
            self.tfb_.sendTransform(ts)


def main(args=None):
    rclpy.init(args=args)
    tf2_publisher_node = TF2PublisherNode()
    rclpy.spin(tf2_publisher_node)
    tf2_publisher_node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()