#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Joy
import yaml
import os
import time
from ament_index_python.packages import get_package_share_directory
import threading

class PS5ControllerMappingNode(Node):
    def __init__(self):
        super().__init__('ps5_controller_mapping_node')
        self.subscription = self.create_subscription(
            Joy,
            'joy',
            self.joy_callback,
            10)
        self.subscription  # prevent unused variable warning

        self.button_mapping = {}
        self.axis_mapping = {}
        self.joy_msg = None
        self.previous_buttons = []
        self.l2_axis_index = None
        self.r2_axis_index = None
        self.prompt_thread = threading.Thread(target=self.prompt_user)
        self.prompt_thread.start()

        self.get_logger().info("Waiting for joystick messages...")

    def joy_callback(self, msg):
        # Normalize L2 and R2 axis values if indices are known
        if self.l2_axis_index is not None and self.r2_axis_index is not None:
            msg.axes[self.l2_axis_index] -= 1.0
            msg.axes[self.r2_axis_index] -= 1.0

        self.joy_msg = msg

        if not self.previous_buttons:
            self.previous_buttons = list(msg.buttons)

        for index, (prev, curr) in enumerate(zip(self.previous_buttons, msg.buttons)):
            if prev != curr and curr == 1:
                self.get_logger().info(f"Button with index {index} was pressed.")

        self.previous_buttons = list(msg.buttons)  # Update the previous_buttons list

    def prompt_user(self):
        while self.joy_msg is None:
            self.get_logger().info("Waiting for joystick messages...")
            time.sleep(1)  # Allow other ROS 2 operations to run

        self.get_logger().info("Joystick messages received. Please press each button and move each axis one by one.")

        self.prompt_button('X')
        self.prompt_button('Circle')
        self.prompt_button('Square')
        self.prompt_button('Triangle')
        self.prompt_button('L1')
        self.prompt_button('R1')
        self.prompt_button('L2')
        self.prompt_button('R2')
        self.prompt_button('Share')
        self.prompt_button('Options')

        self.prompt_axis('LeftStickHorizontal')
        self.prompt_axis('LeftStickVertical')
        self.prompt_trigger_axis('L2')

        self.prompt_axis('RightStickHorizontal')
        self.prompt_axis('RightStickVertical')
        self.prompt_trigger_axis('R2')

        self.prompt_axis('DPadHorizontal')
        self.prompt_axis('DPadVertical')

        self.save_mapping_to_yaml()

    def prompt_button(self, button_name):
        input(f"Press the {button_name} button and then press Enter.")
        index = self.get_button_index()
        self.button_mapping[button_name] = index
        if button_name == 'L2':
            self.l2_axis_index = self.get_trigger_axis_index()
        elif button_name == 'R2':
            self.r2_axis_index = self.get_trigger_axis_index()
        time.sleep(1)  # Allow other ROS 2 operations to run

    def prompt_axis(self, axis_name):
        input(f"Move the {axis_name} and then press Enter.\n")
        self.get_logger().info(f"Current axis: {self.get_axis_index()} was moved.")
        self.axis_mapping[axis_name] = self.get_axis_index()
        time.sleep(1)  # Allow other ROS 2 operations to run

    def prompt_trigger_axis(self, axis_name):
        input(f"Move the {axis_name} and then press Enter.\n")
        self.axis_mapping[axis_name] = self.get_trigger_axis_index()
        self.get_logger().info(f"Current axis: {self.get_axis_index()} was moved.")
        time.sleep(1)  # Allow other ROS 2 operations to run

    def get_button_index(self):
        if self.joy_msg:
            return self.joy_msg.buttons.index(1)
        else:
            self.get_logger().warn("Joy message not received yet.")
            return -1

    def get_axis_index(self):
        if self.joy_msg:
            max_val = max(self.joy_msg.axes, key=abs)
            return self.joy_msg.axes.index(max_val)
        else:
            self.get_logger().warn("Joy message not received yet.")
            return -1

    def get_trigger_axis_index(self):
        if self.joy_msg:
            min_val = min(self.joy_msg.axes)
            return self.joy_msg.axes.index(min_val)
        else:
            self.get_logger().warn("Joy message not received yet.")
            return -1

    def save_mapping_to_yaml(self):
        package_name = 'ur_robotiq_servo'  # Replace with your package name
        share_path = get_package_share_directory(package_name)

        # Navigate to the source config directory
        config_path = os.path.join(share_path, '..', '..', '..', '..', 'src', package_name, 'config')
        config_path = os.path.realpath(config_path)

        if not os.path.exists(config_path):
            os.makedirs(config_path)

        file_path = os.path.join(config_path, 'ps5-controller-mapping.yaml')

        mapping = {
            'buttons': self.button_mapping,
            'axes': self.axis_mapping
        }

        with open(file_path, 'w') as file:
            yaml.dump(mapping, file)

        self.get_logger().info(f"Mapping saved to {file_path}")

def main(args=None):
    rclpy.init(args=args)

    node = PS5ControllerMappingNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
