import time
import pymongo
from bson import ObjectId

import rclpy
from builtin_interfaces.msg import Time
from rclpy.node import Node
from rosidl_runtime_py import convert
from sensor_msgs.msg import JointState
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from fanuc_web_connector_msgs.msg import TrajectoryDB, TrajectoryDBItem, RecordTrajectory, ReplayTrajectory
import json


class ReplayManager(Node):
    def __init__(self):
        super().__init__('replay_manager')

        self.mongo_client = pymongo.MongoClient('localhost', 27017)
        self.mongo_db = self.mongo_client['ros_web_service_demo']
        self.mongo_collection = self.mongo_db['recorded_trajectories']

        self.fanuc_1_joint_states = JointState()
        self.fanuc_2_joint_states = JointState()

        self.fanuc_1_record_name = None
        self.fanuc_2_record_name = None

        self.fanuc_1_record = False
        self.fanuc_2_record = False

        self.fanuc_1_trajectory = JointTrajectory()
        self.fanuc_2_trajectory = JointTrajectory()

        self.db_trajectories = []

        self.pair_trajectory = None

        self.fanuc_1_waiting_for_fanuc_2 = False
        self.fanuc_2_waiting_for_fanuc_1 = False

        self.timer_pair_replay_trajectory = None
        self.timer_pair_replay_trajectory_counter = 0

        self.timer_db_publisher = self.create_timer(1.0, self.db_publisher_callback)

        self.publisher_fanuc_1_goal_trajectory = self.create_publisher(JointTrajectory, '/fanuc_1/goal_trajectory', 10)
        self.publisher_fanuc_2_goal_trajectory = self.create_publisher(JointTrajectory, '/fanuc_2/goal_trajectory', 10)

        self.subscription_fanuc_1_joint_states = self.create_subscription(JointState, '/fanuc_1/joint_states',
                                                                          self.fanuc_1_joint_states_callback,
                                                                          10)
        self.subscription_fanuc_2_joint_states = self.create_subscription(JointState, '/fanuc_2/joint_states',
                                                                          self.fanuc_2_joint_states_callback,
                                                                          10)

        self.publisher_trajectory_db = self.create_publisher(TrajectoryDB, '/trajectory_db', 10)

        self.subscription_record_trajectory = self.create_subscription(RecordTrajectory, '/record_trajectory',
                                                                       self.record_trajectory_callback,
                                                                       10)

        self.subscription_replay_trajectory = self.create_subscription(ReplayTrajectory, '/replay_trajectory',
                                                                       self.replay_trajectory_callback,
                                                                       10)




        self.get_logger().info('Node initialized')

    def db_publisher_callback(self):
        cursor = self.mongo_collection.find({}, {"_id": 1, "trajectory_name": 1})
        self.db_trajectories = []

        for document in cursor:
            item = TrajectoryDBItem()
            item.trajectory_name = document["trajectory_name"]
            item.trajectory_id = str(document["_id"])
            self.db_trajectories.append(item)

        trajectory_db = TrajectoryDB()
        trajectory_db.header.stamp = self.get_clock().now().to_msg()
        trajectory_db.trajectories = self.db_trajectories
        self.publisher_trajectory_db.publish(trajectory_db)



    def record_trajectory_callback(self, record_trajectory):
        self.get_logger().info('RecordTrajectory message received')
        if self.fanuc_1_record_name is None and self.fanuc_2_record_name is None and record_trajectory.record:
            if record_trajectory.robot_name == "fanuc_1":
                self.fanuc_1_record_name = record_trajectory.trajectory_name
                self.fanuc_1_record = True
                self.fanuc_1_trajectory = JointTrajectory()
                self.fanuc_1_trajectory.joint_names = ['J1', 'J2', 'J3', 'J4', 'J5', 'J6']

            elif record_trajectory.robot_name == "fanuc_2":
                self.fanuc_2_record_name = record_trajectory.trajectory_name
                self.fanuc_2_record = True
                self.fanuc_2_trajectory = JointTrajectory()
                self.fanuc_2_trajectory.joint_names = ['J1', 'J2', 'J3', 'J4', 'J5', 'J6']

        elif record_trajectory.record is False:
            if record_trajectory.robot_name == "fanuc_1":
                self.fanuc_1_record = False
            elif record_trajectory.robot_name == "fanuc_2":
                self.fanuc_2_record = False

    def pair_replay_trajectory_callback(self):
        if self.timer_pair_replay_trajectory_counter < 150:
            self.timer_pair_replay_trajectory_counter += 1
            if self.pair_trajectory is not None:
                if not self.fanuc_1_waiting_for_fanuc_2 and not self.fanuc_2_waiting_for_fanuc_1:
                    time.sleep(0.4)
                    self.publisher_fanuc_1_goal_trajectory.publish(self.pair_trajectory)
                    self.publisher_fanuc_2_goal_trajectory.publish(self.pair_trajectory)
                    self.get_logger().info('Trajectory published for fanuc_1 and fanuc_2')
                    self.pair_trajectory = None
                    self.timer_pair_replay_trajectory_counter = 0
                    self.timer_pair_replay_trajectory.cancel()

        else:
            self.get_logger().info('Trajectory not published for fanuc_1 and fanuc_2')
            self.pair_trajectory = None
            self.fanuc_1_waiting_for_fanuc_2 = False
            self.fanuc_2_waiting_for_fanuc_1 = False
            self.timer_pair_replay_trajectory_counter = 0
            self.timer_pair_replay_trajectory.cancel()

    def replay_trajectory_callback(self, replay_trajectory):
        if len(replay_trajectory.robot_names) == 1:
            if replay_trajectory.robot_names[0] == "fanuc_1":
                self.get_logger().info('ReplayTrajectory message received')
                self.get_logger().info('Replaying trajectory for fanuc_1')

                trajectory = self.load_trajectory_from_db(replay_trajectory.trajectory_id)
                if trajectory is not None:
                    self.publisher_fanuc_1_goal_trajectory.publish(trajectory)
                    self.get_logger().info('Trajectory published for fanuc_1')
                else:
                    self.get_logger().info('Trajectory not found in DB')
            elif replay_trajectory.robot_names[0] == "fanuc_2":
                self.get_logger().info('ReplayTrajectory message received')
                self.get_logger().info('Replaying trajectory for fanuc_2')

                trajectory = self.load_trajectory_from_db(replay_trajectory.trajectory_id)
                if trajectory is not None:
                    self.publisher_fanuc_2_goal_trajectory.publish(trajectory)
                    self.get_logger().info('Trajectory published for fanuc_2')
                else:
                    self.get_logger().info('Trajectory not found in DB')
        else:
            if (replay_trajectory.robot_names[0] == "fanuc_1" and replay_trajectory.robot_names[1] == "fanuc_2") or \
                    (replay_trajectory.robot_names[0] == "fanuc_2" and replay_trajectory.robot_names[1] == "fanuc_1"):
                self.get_logger().info('ReplayTrajectory message received for pair')
                pair_trajectory = self.load_trajectory_from_db(replay_trajectory.trajectory_id)
                if pair_trajectory is not None:
                    temp_trajectory = JointTrajectory()
                    temp_trajectory.joint_names = ['J1', 'J2', 'J3', 'J4', 'J5', 'J6']
                    temp_trajectory.points.append(pair_trajectory.points[0])
                    self.publisher_fanuc_1_goal_trajectory.publish(temp_trajectory)
                    self.publisher_fanuc_2_goal_trajectory.publish(temp_trajectory)
                    self.fanuc_1_waiting_for_fanuc_2 = True
                    self.fanuc_2_waiting_for_fanuc_1 = True
                    self.pair_trajectory = pair_trajectory

                    self.timer_pair_replay_trajectory = self.create_timer(0.1, self.pair_replay_trajectory_callback)

                    self.get_logger().info('Trajectory published for pair')


            else:
                self.get_logger().info('ReplayTrajectory message received')
                self.get_logger().info('Robot names are not correct')

    def fanuc_1_joint_states_callback(self, joint_states):
        self.get_logger().info('JointState message received')
        self.fanuc_1_joint_states = joint_states
        if self.fanuc_1_record_name is not None:
            if self.fanuc_1_record:
                joint_trajectory_point = JointTrajectoryPoint()
                joint_trajectory_point.time_from_start.sec = joint_states.header.stamp.sec
                joint_trajectory_point.time_from_start.nanosec = joint_states.header.stamp.nanosec
                joint_trajectory_point.positions = joint_states.position
                self.fanuc_1_trajectory.points.append(joint_trajectory_point)
            else:
                first_time_stamp = self.fanuc_1_trajectory.points[0].time_from_start
                final_trajectory = JointTrajectory()
                for point in self.fanuc_1_trajectory.points:
                    new_point = JointTrajectoryPoint()
                    new_point.positions = point.positions

                    new_time_stamp = self.calculate_time_difference(first_time_stamp, point.time_from_start)
                    new_point.time_from_start.sec = new_time_stamp.sec + 1
                    new_point.time_from_start.nanosec = new_time_stamp.nanosec
                    final_trajectory.points.append(new_point)
                self.save_joint_trajectory(final_trajectory, (str(self.fanuc_1_record_name)))
                self.fanuc_1_record_name = None
                self.fanuc_1_trajectory = JointTrajectory()
        elif self.pair_trajectory is not None:
            self.get_logger().info('Publishing pair trajectory')
            self.get_logger().info('Fanuc_1 waiting for Fanuc_2: ' + str(self.fanuc_1_waiting_for_fanuc_2))
            self.get_logger().info('Fanuc_2 waiting for Fanuc_1: ' + str(self.fanuc_2_waiting_for_fanuc_1))

            if self.fanuc_2_waiting_for_fanuc_1 and self.confiuration_reached(joint_states, self.pair_trajectory.points[0]):
                self.fanuc_2_waiting_for_fanuc_1 = False

    def fanuc_2_joint_states_callback(self, joint_states):
        self.get_logger().info('JointState message received')
        self.fanuc_2_joint_states = joint_states
        if self.fanuc_2_record_name is not None:
            if self.fanuc_2_record:
                joint_trajectory_point = JointTrajectoryPoint()
                joint_trajectory_point.time_from_start.sec = joint_states.header.stamp.sec
                joint_trajectory_point.time_from_start.nanosec = joint_states.header.stamp.nanosec
                joint_trajectory_point.positions = joint_states.position
                self.fanuc_2_trajectory.points.append(joint_trajectory_point)
            else:
                first_time_stamp = self.fanuc_2_trajectory.points[0].time_from_start
                final_trajectory = JointTrajectory()
                for point in self.fanuc_2_trajectory.points:
                    new_point = JointTrajectoryPoint()
                    new_point.positions = point.positions

                    new_time_stamp = self.calculate_time_difference(first_time_stamp, point.time_from_start)
                    new_point.time_from_start.sec = new_time_stamp.sec + 1
                    new_point.time_from_start.nanosec = new_time_stamp.nanosec
                    final_trajectory.points.append(new_point)
                self.save_joint_trajectory(final_trajectory, (str(self.fanuc_2_record_name)))
                self.fanuc_2_record_name = None
                self.fanuc_2_trajectory = JointTrajectory()
        elif self.pair_trajectory is not None:
            self.get_logger().info('Publishing pair trajectory')

            if self.fanuc_1_waiting_for_fanuc_2 and self.confiuration_reached(joint_states, self.pair_trajectory.points[0]):
                self.fanuc_1_waiting_for_fanuc_2 = False

    def save_joint_trajectory(self, joint_trajectory, trajectory_name):
        try:
            serialized_states = []
            for state in joint_trajectory.points:
                serialized_states.append(convert.message_to_ordereddict(state))

            self.mongo_collection.insert_one({"trajectory_name": trajectory_name, "trajectory": serialized_states})

        except Exception as e:
            self.get_logger().error(f"Failed to save joint states to file: {str(e)}")

    def load_trajectory_from_db(self, trajectory_id):
        try:
            entry = self.mongo_collection.find_one({"_id": ObjectId(trajectory_id)})
            serialized_states = entry["trajectory"]

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

            for serialized_state in serialized_states:
                state = JointTrajectoryPoint()
                state.positions = serialized_state["positions"]
                state.time_from_start.sec = serialized_state["time_from_start"]["sec"]
                state.time_from_start.nanosec = serialized_state["time_from_start"]["nanosec"]
                trajectory.points.append(state)

            return trajectory

        except Exception as e:
            self.get_logger().error(f"Failed to load joint states from file: {str(e)}")
            return None



    def calculate_time_difference(self, zero_time_stamp, current_time):
        # Convert the given times into milliseconds
        time1_ms = (zero_time_stamp.sec * 1000) + (zero_time_stamp.nanosec // 1000000)
        time2_ms = (current_time.sec * 1000) + (current_time.nanosec // 1000000)

        # Calculate the time difference in milliseconds
        time_diff_ms = time2_ms - time1_ms

        # Convert the time difference back to seconds and nanoseconds
        diff_time = Time()

        diff_time.sec = time_diff_ms // 1000
        diff_time.nanosec = (time_diff_ms % 1000) * 1000000

        return diff_time

    def confiuration_reached(self, joint_states, target_joint_states):
        for i in range(len(joint_states.position)):
            if abs(joint_states.position[i] - target_joint_states.positions[i]) > 0.1:
                return False
        return True


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


if __name__ == '__main__':
    main()
