import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Joy, JointState
from geometry_msgs.msg import TwistStamped
from control_msgs.msg import JointJog
from std_srvs.srv import Trigger
from robotiq_2f_msgs.msg import ForwardCommand
import yaml
import os

def load_yaml(file_path):
    with open(file_path, 'r') as file:
        return yaml.safe_load(file)

class PS5ControllerNode(Node):
    def __init__(self):
        super().__init__('ps5_controller_node')

        # Load the YAML file path from parameters
        self.yaml_file_path = self.declare_parameter('yaml_file_path', '').get_parameter_value().string_value

        # Validate and load the YAML file
        if not self.yaml_file_path or not os.path.exists(self.yaml_file_path):
            self.get_logger().error(f"Invalid YAML file path: {self.yaml_file_path}")
            rclpy.shutdown()
            return

        self.mappings = load_yaml(self.yaml_file_path)
        self.get_logger().info(f"Loaded mappings: {self.mappings}")


        # States
        self.mode = 'twist'  # Initialize mode to 'twist'. Alternatives: 'twist', 'joint'
        self.last_button_state = 0  # Track the last state of the toggle button to detect presses
        self.last_gripper_gap = 0.0
        self.gripper_cmd = 0.0

        # Parameters
        self.linear_speed_multiplier = self.declare_parameter('linear_speed_multiplier', 1.0)
        self.linear_speed_multiplier = self.get_parameter('linear_speed_multiplier').get_parameter_value().double_value
        self.get_logger().info(f"Linear speed multiplier: {self.linear_speed_multiplier}")

        self.use_fake_hardware = self.declare_parameter('use_fake_hardware', False)
        self.use_fake_hardware = self.get_parameter('use_fake_hardware').get_parameter_value().bool_value
        self.get_logger().info(f"Use fake hardware: {self.use_fake_hardware}")

        self.gripper_position_multiplier = self.declare_parameter('gripper_position_multiplier', 1.0)
        self.gripper_position_multiplier = (self.get_parameter('gripper_position_multiplier')
                                         .get_parameter_value().double_value)
        self.get_logger().info(f"Gripper position multiplier: {self.gripper_position_multiplier}")

        self.gripper_cycle_time = self.declare_parameter('gripper_cycle_time', 1.0)
        self.gripper_cycle_time = (self.get_parameter('gripper_cycle_time')
                                         .get_parameter_value().double_value)
        self.get_logger().info(f"Gripper cycle time: {self.gripper_cycle_time}")

        self.gripper_lower_limit = self.declare_parameter('gripper_lower_limit', 0.0)
        self.gripper_lower_limit = (self.get_parameter('gripper_lower_limit')
                                         .get_parameter_value().double_value)
        self.get_logger().info(f"Gripper lower limit: {self.gripper_lower_limit}")

        self.gripper_upper_limit = self.declare_parameter('gripper_upper_limit', 0.14)
        self.gripper_upper_limit = (self.get_parameter('gripper_upper_limit')
                                         .get_parameter_value().double_value)
        self.get_logger().info(f"Gripper upper limit: {self.gripper_upper_limit}")

        # Subscriber and Publisher
        self.joint_state_sub = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_state_callback,
            10)

        self.joy_sub = self.create_subscription(
            Joy,
            '/joy',
            self.joy_callback,
            10)

        self.twist_pub = self.create_publisher(
            TwistStamped,
            '/servo_node/delta_twist_cmds',
            10)

        self.joint_pub = self.create_publisher(
            JointJog,
            '/servo_node/delta_joint_cmds',
            10)

        self.gripper_cmd_pub = self.create_publisher(
            ForwardCommand,
            '/forward_gripper_position_controller/command',
            10)

        # Services
        self.servo_client = self.create_client(Trigger, '/servo_node/start_servo')

        while not self.servo_client.wait_for_service(timeout_sec=5.0):
            self.get_logger().info('/servo_node/start_servo service not available, waiting again...')

        self.call_start_servo()

        # Timer for publishing gripper command
        self.gripper_timer = self.create_timer(self.gripper_cycle_time, self.publish_gripper_command)

        self.get_logger().info('ps5_servo_node started!')

    def call_start_servo(self):
        request = Trigger.Request()
        future = self.servo_client.call_async(request)
        rclpy.spin_until_future_complete(self, future)
        response = future.result()
        if response.success:
            self.get_logger().info(f'Successfully called start_servo: {response.message}')
        else:
            self.get_logger().info('Failed to call start_servo')

    def joint_state_callback(self, msg):
        try:
            index = msg.name.index('gripper_gap')
            self.last_gripper_gap = msg.position[index]
        except ValueError:
            self.get_logger().error('gripper_gap not found in /joint_states msg')

    def publish_gripper_command(self):
        if self.gripper_cmd is not None:
            self.get_logger().debug(f'Publishing gripper command: {self.gripper_cmd}')
            new_gripper_position = self.gripper_cmd + self.last_gripper_gap
            if new_gripper_position >= self.gripper_upper_limit:
                new_gripper_position = self.gripper_upper_limit
            elif new_gripper_position <= self.gripper_lower_limit:
                new_gripper_position = self.gripper_lower_limit


            self.gripper_cmd = 0.0
            gripper_msg = ForwardCommand()
            gripper_msg.position = new_gripper_position
            gripper_msg.max_velocity = 0.01
            gripper_msg.max_effort = 235.0
            self.gripper_cmd_pub.publish(gripper_msg)
        else:
            self.get_logger().warn('Gripper command not available!')

    def joy_callback(self, msg):
        # Check for button press to toggle mode
        toggle_button_index = self.mappings['buttons']['Options']  # Update to the correct button index for toggling
        current_button_state = msg.buttons[toggle_button_index]
        if current_button_state == 1 and self.last_button_state == 0:
            self.mode = 'joint' if self.mode == 'twist' else 'twist'
            self.get_logger().info(f'Mode switched to: {self.mode}')
        self.last_button_state = current_button_state

        left_trigger = (msg.axes[self.mappings['axes']['L2']] - 1) / -2.0
        right_trigger = (msg.axes[self.mappings['axes']['R2']] - 1) / -2.0

        # Process control input based on current mode
        if self.mode == 'twist':
            twist_msg = TwistStamped()
            twist_msg.header.frame_id = 'tool0'
            twist_msg.header.stamp = self.get_clock().now().to_msg()
            twist_msg.twist.linear.x = msg.axes[self.mappings['axes']['LeftStickHorizontal']] * self.linear_speed_multiplier
            twist_msg.twist.linear.y = msg.axes[self.mappings['axes']['LeftStickVertical']] * self.linear_speed_multiplier
            twist_msg.twist.linear.z = (left_trigger - right_trigger) * self.linear_speed_multiplier
            twist_msg.twist.angular.x = msg.axes[self.mappings['axes']['RightStickHorizontal']]
            twist_msg.twist.angular.y = msg.axes[self.mappings['axes']['RightStickVertical']]
            twist_msg.twist.angular.z = (msg.buttons[self.mappings['buttons']['L1']] - msg.buttons[self.mappings['buttons']['R1']]) * 1.0
            self.twist_pub.publish(twist_msg)
        elif self.mode == 'joint':
            joint_msg = JointJog()
            joint_msg.header.frame_id = 'tool0'
            joint_msg.header.stamp = self.get_clock().now().to_msg()
            joint_msg.joint_names = [
                "shoulder_pan_joint",
                "shoulder_lift_joint",
                "elbow_joint",
                "wrist_1_joint",
                "wrist_2_joint",
                "wrist_3_joint"
            ]
            joint_msg.velocities = [
                msg.axes[self.mappings['axes']['LeftStickHorizontal']],
                msg.axes[self.mappings['axes']['LeftStickVertical']],
                msg.axes[self.mappings['axes']['RightStickHorizontal']],
                msg.axes[self.mappings['axes']['RightStickVertical']],
                msg.axes[self.mappings['axes']['DPadHorizontal']],
                msg.axes[self.mappings['axes']['DPadVertical']]
            ]
            self.joint_pub.publish(joint_msg)

        # Gripper controller
        gripper_open_button = self.mappings['buttons']['Circle']
        gripper_close_button = self.mappings['buttons']['X']
        self.gripper_cmd  = self.gripper_cmd + int((msg.buttons[gripper_open_button] - msg.buttons[gripper_close_button])) * self.gripper_position_multiplier

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


if __name__ == '__main__':
    main()
