import rclpy
from rclpy.node import Node
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from sensor_msgs.msg import JointState
from osc4py3.as_allthreads import *
from osc4py3 import oscmethod as osm
import xml.etree.ElementTree as ET
import time
import numpy as np
import os

class ScaledJointTrajectoryPublisher(Node):
    """Node to publish joint trajectories based on OSC messages."""
    
    def __init__(self, joint_names, joint_velocity_limits):
        super().__init__('scaled_joint_trajectory_publisher')
        
        # ROS2 Publisher
        self.publisher = self.create_publisher(
            JointTrajectory, 
            '/scaled_joint_trajectory_controller/joint_trajectory', 
            1
        )

        self.subscription = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_states_callback,
            1  # Increased queue size for joint states
        )

        # Store received joint positions
        self.current_joint_positions = [0.0] * len(joint_names)
        self.joint_names = joint_names
        self.joint_velocity_limits = joint_velocity_limits
        self.desired_joint_positions = [0.0] * len(joint_names)
        self.previous_desired = [0.0] * len(joint_names)

        ip = "0.0.0.0"  # Listen on all network interfaces
        port = 8000     # Must match the sender's port in `joint_state_osc.py`

        osc_startup()
        osc_udp_server(ip, port, "osc_server")
        print("Server started on 0.0.0.0:8000 \n ready to receive messages in the following format: /joint_trajectroy [joint_positions] optional: duration as last argument")
        # Register OSC handler
        osc_method("/joint_angles", self.joint_angles_handler, argscheme=osm.OSCARG_DATAUNPACK)

        while True:
            try:
                self.x_limits = [float(i) if i != 'x' else None for i in input("Enter the lower and upper limits for x (space-separated, enter 'x' for no limit): ").split()]
                self.y_limits = [float(i) if i != 'x' else None for i in input("Enter the lower and upper limits for y (space-separated, enter 'x' for no limit): ").split()]
                self.z_limits = [float(i) if i != 'x' else None for i in input("Enter the lower and upper limits for z (space-separated, enter 'x' for no limit): ").split()]
                
                if len(self.x_limits) != 2 or len(self.y_limits) != 2 or len(self.z_limits) != 2:
                    print("Invalid input. Please enter exactly two values (or leave blank) for each limit.")
                    continue
                
                if (self.x_limits[0] is not None and self.x_limits[1] is not None and self.x_limits[0] >= self.x_limits[1]) or \
                (self.y_limits[0] is not None and self.y_limits[1] is not None and self.y_limits[0] >= self.y_limits[1]) or \
                (self.z_limits[0] is not None and self.z_limits[1] is not None and self.z_limits[0] >= self.z_limits[1]):
                    print("Invalid input. Lower limit must be less than upper limit for each axis.")
                    continue
                
                print(f"Current limits:")
                print(f"x: {self.x_limits}")
                print(f"y: {self.y_limits}")
                print(f"z: {self.z_limits}")
                confirm = input("Do you want your robot to move in this range? (y/n): ").strip().lower()
                if confirm == 'y':
                    break
                elif confirm == 'n':
                    print("Please re-enter the limits.")
                else:
                    print("Invalid input. Please enter 'y' or 'n'.")
            except ValueError:
                print("Invalid input. Please enter numeric values only.")
        
        # Ask the user if they want to set new joint limits
        update_limits = input("Do you want to set new joint limits? (y/n): ").strip().lower()
        if update_limits == 'y':
            for joint_name in self.joint_names:
                while True:
                    try:
                        print(f"Current position limits for joint '{joint_name}': {self.robot.links[joint_name]} rad")
                        lower_limit = input(f"Enter the new lower limit for joint '{joint_name}' (or press Enter to keep current): ").strip()
                        upper_limit = input(f"Enter the new upper limit for joint '{joint_name}' (or press Enter to keep current): ").strip()
                        
                        if lower_limit is not None and upper_limit is not None and lower_limit >= upper_limit:
                            print("Invalid input. Lower limit must be less than upper limit.")
                            continue
                        
                        if lower_limit:
                            self.robot.links[joint_name][0] = float(lower_limit)
                        if upper_limit:
                            self.robot.links[joint_name][1] = float(upper_limit)
                        break
                    except ValueError:
                        print("Invalid input. Please enter numeric values or leave blank to keep current limits.")
            

        self.hz = float(input("Enter the desired refresh frequency (Hz): "))
        # Start the OSC server in a separate thread to avoid blocking the ROS 2 event loop
        self.create_timer(1/self.hz, self.update_position)  # Timer to call osc_process periodically

    def joint_angles_handler(self, *args):
    # Ensure the desired joint positions are within the specified limits

        x, y, z, r, p, yaw, *_ = [float(i) for i in list(args)]
        if self.x_limits[0] is not None:
            x = max(self.x_limits[0], x)
        if self.x_limits[1] is not None:
            x = min(self.x_limits[1], x)
        if self.y_limits[0] is not None:
            y = max(self.y_limits[0], y)
        if self.y_limits[1] is not None:
            y = min(self.y_limits[1], y)
        if self.z_limits[0] is not None:
            z = max(self.z_limits[0], z)
        if self.z_limits[1] is not None:
            z = min(self.z_limits[1], z)

        if x != args[0] or y != args[1] or z != args[2]:
            self.get_logger().warn(
            f"Desired joint positions adjusted to fit within limits: "
            f"x={x}, y={y}, z={z} (original: x={args[0]}, y={args[1]}, z={args[2]})"
            )

        self.desired_joint_positions = [x, y, z, r, p, yaw]
        

    def joint_states_callback(self, msg):
        """Callback function to handle incoming joint states."""
        joint_names = msg.name  # List of joint names
        joint_position_dict = dict(zip(joint_names, self.current_joint_positions))
        self.current_joint_positions = [joint_position_dict[name] for name in self.joint_names]

    def update_position(self):
        time1 = time.time()
        if self.desired_joint_positions == self.previous_desired:
            return
        duration = 0
        for p1, p2, max_vel in zip(self.desired_joint_positions, self.current_joint_positions, self.joint_velocity_limits.values()):
            duration = max(max(duration, abs(p1 - p2) / max_vel),1/self.hz)# as minimun
        duration *= 1.8
        #print(f"Duration: {duration}")
        #print(f'vel: {max(np.array(self.desired_joint_positions_joint_positions) - np.array(self.joint_positions)/duration)}')
        msg = JointTrajectory()
        msg.joint_names = self.joint_names
        point = JointTrajectoryPoint()
        point.positions = self.desired_joint_positions  # Updated joint positions
        point.time_from_start.sec = int(duration)
        #point.time_from_start.sec = 10
        point.time_from_start.nanosec = int((duration - int(duration)) * 1e9)
        #point.time_from_start.nanosec = 0
        #point.velocities = [i*0.95 for i in self.joint_velocity_limits.values()]
        msg.points.append(point)
        self.publisher.publish(msg)
        #print(f"desired: {self.desired_joint_positions}")
        #print(f"desired: {[180/3.141*i for i in self.desired_joint_positions]}")
        #print(f"current: {[180/3.141*i for i in self.current_joint_positions]}")
        self.previous_desired = self.desired_joint_positions

def main():
    """Main function to get joint names and start the ROS 2 & OSC system."""
    while True:
        path_to_urdf = input("Enter the path to the URDF file: ")
        if os.path.isfile(path_to_urdf):
            if not path_to_urdf.endswith('.urdf'):
                print("The file is not a URDF file. Please enter a valid URDF file.")
                continue
            break
        else:
            print("Invalid path. Please enter a valid path to the URDF file.")
    tree = ET.parse(path_to_urdf)
    root = tree.getroot()
    joint_names = [joint.get('name') for joint in root.findall('joint') if joint.get('type') == 'revolute' or joint.get('type') == 'continuous' or joint.get('type') == 'prismatic']
    
    joint_velocity_limits = {}

    # Iterate over all joints in the URDF
    for joint in root.findall('.//joint'):
        joint_name = joint.get('name')  # Get the name of the joint
        
        # Look for the <limit> tag under each joint
        limit = joint.find('limit')
        
        if limit is not None:
            # Extract the velocity limit (if it exists)
            velocity_limit = limit.get('velocity')
            
            if velocity_limit is not None:
                joint_velocity_limits[joint_name] = float(velocity_limit)

    rclpy.init()

    node = ScaledJointTrajectoryPublisher(joint_names, joint_velocity_limits)
    
    # Run ROS 2 spin, and osc_process will be handled by the timer
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        print("")
    finally:
        node.destroy_node()
        rclpy.shutdown()
        osc_terminate()

if __name__ == '__main__':
    main()
