#!/usr/bin/env python3

import rospy
import actionlib
from trajectory_msgs.msg import JointTrajectory
from control_msgs.msg import FollowJointTrajectoryAction, FollowJointTrajectoryGoal
from fanuc_msg_translator_msgs.msg import FanucTrajectoryConverterQueue, FanucTrajectoryConverterQueueItem, \
    FanucTrajectoryConverterCancel

# Namespace of the node
ns = rospy.get_namespace()

# Queue for all trajectories
queue = []

# Queue for all finished trajectories
cemetery = []

# Conntection to the action server of the fanuc_interface_node
client = actionlib.SimpleActionClient(ns + 'arm_controller/follow_joint_trajectory', FollowJointTrajectoryAction)

# Publisher for the queue (for ROS-2)
queue_publisher = rospy.Publisher(ns + 'bridge/follow_joint_trajectory/queue', FanucTrajectoryConverterQueue,
                                  queue_size=10)

# Marker for 500 ms delay in the queue handling, activated when a trajectory is canceled
canceled_trajectory = False


def goal_callback(trajectory):
    """
    This function is called when a new trajectory is received. It adds the trajectory to the queue.

    :param trajectory: The received trajectory message
    """
    global queue
    rospy.loginfo("[Trajectory Converter] Trajectory received!")
    trajectory_wrapper = TrajectoryWrapper(trajectory)
    queue.append(trajectory_wrapper)


def cancel_callback(cancel):
    """
    This function is called when a cancel command is received. It deletes the trajectory from the queue if it is not
    sent to the action server yet. If it is sent to the action server, the goal is canceled and the trajectory is
    moved to the cemetery.

    :param cancel: The received cancel message
    """
    global client
    global queue
    global cemetery
    global canceled_trajectory

    rospy.loginfo("[Trajectory Converter] Cancel command received!")

    for trajectory_wrapper in queue:
        if trajectory_wrapper.get_trajectory().header.stamp == cancel.received_time_stamp:
            if trajectory_wrapper.get_status() == 0:
                trajectory_wrapper.set_status(100)
                cemetery.append(queue.pop(queue.index(trajectory_wrapper)))
            elif trajectory_wrapper.get_status() == 1:
                client.cancel_all_goals()
                canceled_trajectory = True
            break


def done_callback(state, result):
    """
    This function is called when the action server is done with the trajectory. It sets the status of the trajectory
    in the queue to 100, moves it to the cemetery and sets the action status of the trajectory to the state of
    the action server.

    :param state: The state of the action server
    :param result: The result of the action server
    """
    global queue

    rospy.loginfo("[Trajectory Converter] Action server done! State: " + str(state))

    queue[0].set_action_status(state)


class TrajectoryWrapper:
    """
    This class is used to wrap the trajectory message and add two variables to it: status and action_status.

    status: The status of the trajectory in the queue.
    action_status: The status of the action server of the fanuc_interface_node.
    """

    def __init__(self, trajectory=None):
        """
        The constructor of the TrajectoryWrapper class.

        :param trajectory: The trajectory message to be wrapped
        """
        self.trajectory = trajectory
        self.status = 0
        self.action_status = -1

    def set_action_status(self, status):
        """
        This function sets the action status of the trajectory.

        :param status: The action status to be set
        """
        self.action_status = status

    def get_action_status(self):
        """
        This function returns the action status of the trajectory.

        :return: The action status of the trajectory
        """
        return self.action_status

    def set_status(self, status):
        """
        This function sets the status of the trajectory.

        :param status: The status to be set
        """
        self.status = status

    def get_status(self):
        """
        This function returns the status of the trajectory.

        :return: The status of the trajectory
        """
        return self.status

    def get_trajectory(self):
        """
        This function returns the trajectory message.

        :return: The trajectory message
        """
        return self.trajectory


def main():
    """
    This is the main function of the node. It initializes the node, subscribes to the topics and runs the queue handler.

    The queue hanlder is called every 0.1 seconds (while loop). It works like a state machine for the queue.
    Trajectories can have the following states:
    0: Not started
    1: In progress
    2: Finished
    100: Finished and in the cemetery
    101 - 129: Finished and in the cemetery (timer)
    130: Finished and removed from the cemetery (not published to /msg_translator/follow_joint_trajectory/queue anymore)

    The action status of the trajectory is saved in the variable action_status. The action status is the status of the
    action server of the fanuc_interface_node. The action status is -1 until the action server reaches a final state for
    this goal. Then it changes to the resulting state of the action server. The action status can be found here:
    http://docs.ros.org/en/api/actionlib_msgs/html/msg/GoalStatus.html

    This function also publishes the queue and the cemetery to the topic /msg_translator/follow_joint_trajectory/queue.
    """

    global client
    global queue_publisher
    global queue
    global cemetery
    global canceled_trajectory

    rospy.init_node('fanuc_trajectory_converter_node')

    rospy.Subscriber(ns + 'bridge/follow_joint_trajectory/goal', JointTrajectory, goal_callback, queue_size=10)
    rospy.Subscriber(ns + 'bridge/follow_joint_trajectory/cancel', FanucTrajectoryConverterCancel, cancel_callback,
                     queue_size=10)

    client.wait_for_server()

    rospy.loginfo("[Trajectory Converter] Node initialized!")

    while not rospy.is_shutdown():
        if len(queue) > 0:
            if queue[0].get_status() == 0 \
                    and client.get_state() != actionlib.GoalStatus.ACTIVE \
                    and client.get_state() != actionlib.GoalStatus.PENDING:
                goal = FollowJointTrajectoryGoal()
                goal.trajectory = queue[0].get_trajectory()
                queue[0].set_status(1)
                if canceled_trajectory:
                    rospy.sleep(0.5)
                    canceled_trajectory = False
                client.send_goal(goal, done_cb=done_callback)
            elif queue[0].get_status() == 1:
                if 2 <= client.get_state() <= 8:
                    queue[0].set_status(2)
                    queue[0].set_action_status(client.get_state())
            elif queue[0].get_status() == 2:
                queue[0].set_status(100)
                cemetery.append(queue.pop(0))

        if len(cemetery) > 0:
            for trajectory_wrapper in cemetery:
                if trajectory_wrapper.get_status() == 130:
                    cemetery.remove(trajectory_wrapper)
                else:
                    trajectory_wrapper.set_status(trajectory_wrapper.get_status() + 1)

        msg = FanucTrajectoryConverterQueue()
        msg.header.stamp = rospy.Time.now()

        for trajectory_wrapper in cemetery:
            goal = FanucTrajectoryConverterQueueItem()
            goal.received_time_stamp = trajectory_wrapper.get_trajectory().header.stamp
            goal.status_in_queue = trajectory_wrapper.get_status()
            goal.status_on_server = trajectory_wrapper.get_action_status()
            msg.goals.append(goal)

        for trajectory_wrapper in queue:
            goal = FanucTrajectoryConverterQueueItem()
            goal.received_time_stamp = trajectory_wrapper.get_trajectory().header.stamp
            goal.status_in_queue = trajectory_wrapper.get_status()
            goal.status_on_server = trajectory_wrapper.get_action_status()
            msg.goals.append(goal)

        queue_publisher.publish(msg)

        rospy.sleep(0.1)


if __name__ == '__main__':
    main()
