#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient
import math
import time

# ROS 2 Message, Action, and Service Imports
from moveit_msgs.action import MoveGroup, ExecuteTrajectory
from moveit_msgs.msg import Constraints, PositionConstraint, OrientationConstraint, BoundingVolume, RobotState
from moveit_msgs.srv import GetCartesianPath, GetPositionFK
from shape_msgs.msg import SolidPrimitive
from geometry_msgs.msg import PoseStamped, Pose
from sensor_msgs.msg import JointState

class CameraArmController(Node):
    def __init__(self):
        super().__init__('camera_arm_controller')
        
        # 1. Action Client for normal point-to-point planning
        self._action_client = ActionClient(self, MoveGroup, '/move_action')
        self.get_logger().info("Waiting for MoveIt action server...")
        self._action_client.wait_for_server()
        
        # 2. Action Client for running raw pre-calculated paths
        self._execute_client = ActionClient(self, ExecuteTrajectory, '/execute_trajectory')
        self.get_logger().info("Waiting for execution action server...")
        self._execute_client.wait_for_server()
        
        # 3. Service Client for calculating straight horizontal camera paths
        self._cartesian_client = self.create_client(GetCartesianPath, '/compute_cartesian_path')
        self.get_logger().info("Connecting to Cartesian path service...")
        self._cartesian_client.wait_for_service()

        # 4. NEW: Service Client for fetching live end-effector coordinates (Forward Kinematics)
        self._fk_client = self.create_client(GetPositionFK, '/compute_fk')
        self.get_logger().info("Connecting to Forward Kinematics service...")
        self._fk_client.wait_for_service()
        
        # 5. Subscriber to listen to current joint states (Required to compute current position)
        self.current_joint_state = None
        self._joint_sub = self.create_subscription(
            JointState,
            '/joint_states',
            self.joint_state_callback,
            10
        )
        
        self.get_logger().info("🚀 Automated Camera Arm Pipeline Fully Operational!")

    def joint_state_callback(self, msg):
        """Saves current joint angles to calculate position on the fly"""
        self.current_joint_state = msg

    def get_horizontal_camera_quaternion(self):
        """Calculates a 90-degree pitch offset to keep the camera level with the horizon"""
        pitch_angle = math.radians(-90.0)
        q = Pose().orientation
        q.x = 0.0
        q.y = math.sin(pitch_angle / 2.0)
        q.z = 0.0
        q.w = math.cos(pitch_angle / 2.0)
        return q

    def get_current_end_effector_pose(self):
        """Queries the FK service to find exactly where the camera is right now"""
        # Wait until we receive at least one joint state message from the robot
        while self.current_joint_state is None and rclpy.ok():
            self.get_logger().info("Waiting for initial joint states...")
            rclpy.spin_once(self, timeout_sec=0.1)

        req = GetPositionFK.Request()
        req.header.frame_id = "base_link"
        req.fk_link_names = ["tool_link"] # Name of the tip frame to track
        
        # Pack current joint configuration into the request
        robot_state = RobotState()
        robot_state.joint_state = self.current_joint_state
        req.robot_state = robot_state

        # Call the service synchronously
        future = self._fk_client.call_async(req)
        rclpy.spin_until_future_complete(self, future)
        
        res = future.result()
        if res and res.error_code.val == 1: # 1 = SUCCESS
            return res.pose_stamped[0].pose
        else:
            self.get_logger().error("Failed to calculate Forward Kinematics!")
            return None

    def send_pt_to_pt_goal(self, x, y, z):
        """Moves the camera to a specific 3D coordinate while forcing it to be horizontal"""
        self.get_logger().info(f"Planning Point-to-Point Move to: X={x}, Y={y}, Z={z}")
        
        goal_msg = MoveGroup.Goal()
        goal_msg.request.group_name = "arm"
        goal_msg.request.allowed_planning_time = 10.0
        goal_msg.request.num_planning_attempts = 15
        goal_msg.request.pipeline_id = "ompl"
        
        target_pose = PoseStamped()
        target_pose.header.frame_id = "base_link"
        target_pose.pose.position.x = float(x)
        target_pose.pose.position.y = float(y)
        target_pose.pose.position.z = float(z)
        target_pose.pose.orientation = self.get_horizontal_camera_quaternion()

        position_constraint = PositionConstraint()
        position_constraint.header.frame_id = "base_link"
        position_constraint.link_name = "tool_link"
        
        box = SolidPrimitive()
        box.type = SolidPrimitive.BOX
        box.dimensions = [0.03, 0.03, 0.03]
        
        volume = BoundingVolume()
        volume.primitives.append(box)
        
        box_pose = Pose()
        box_pose.position = target_pose.pose.position
        box_pose.orientation.w = 1.0
        volume.primitive_poses.append(box_pose)
        
        position_constraint.constraint_region = volume
        position_constraint.weight = 1.0

        orient_constraint = OrientationConstraint()
        orient_constraint.header.frame_id = "base_link"
        orient_constraint.link_name = "tool_link"
        orient_constraint.orientation = target_pose.pose.orientation
        
        orient_constraint.absolute_x_axis_tolerance = 0.15  
        orient_constraint.absolute_y_axis_tolerance = 0.15  
        orient_constraint.absolute_z_axis_tolerance = 3.14  
        orient_constraint.weight = 1.0

        goal_constraints = Constraints()
        goal_constraints.position_constraints.append(position_constraint)
        goal_constraints.orientation_constraints.append(orient_constraint)
        goal_msg.request.goal_constraints.append(goal_constraints)

        send_goal_future = self._action_client.send_goal_async(goal_msg)
        rclpy.spin_until_future_complete(self, send_goal_future)
        
        goal_handle = send_goal_future.result()
        if not goal_handle.accepted:
            self.get_logger().error("Point-to-Point goal rejected by MoveIt!")
            return False

        self.get_logger().info("Goal accepted! Executing path...")
        get_result_future = goal_handle.get_result_async()
        rclpy.spin_until_future_complete(self, get_result_future)
        return True

    def execute_straight_camera_sweep(self, distance=0.2):
        """
        AUTOMATED NO-XYZ METHOD: Fetches live position from the simulation 
        and advances the camera straight forward parallel to the ground.
        """
        # 1. Fetch live position automatically
        current_pose = self.get_current_end_effector_pose()
        if current_pose is None:
            self.get_logger().error("Cannot track sweep because current pose is unknown.")
            return False

        current_x = current_pose.position.x
        current_y = current_pose.position.y
        current_z = current_pose.position.z

        self.get_logger().info(f"Detected current position: X={current_x:.3f}, Y={current_y:.3f}, Z={current_z:.3f}")
        self.get_logger().info(f"Calculating straight crawl: advancing +{distance}m forward along X...")
        
        req = GetCartesianPath.Request()
        req.group_name = "arm"
        req.header.frame_id = "base_link"
        req.start_state.is_diff = True 
        
        # Target waypoint calculated using live variables
        target_waypoint = Pose()
        target_waypoint.position.x = current_x + distance
        target_waypoint.position.y = current_y
        target_waypoint.position.z = current_z
        target_waypoint.orientation = self.get_horizontal_camera_quaternion() 

        req.waypoints = [target_waypoint]
        req.max_step = 0.01          
        req.jump_threshold = 0.0     

        # Call Cartesian planner service
        future = self._cartesian_client.call_async(req)
        rclpy.spin_until_future_complete(self, future)
        
        res = future.result()
        if res.fraction < 1.0:
            self.get_logger().warn(f" Only calculated {res.fraction*100:.1f}% of path due to link extension limits.")
            if res.fraction < 0.5:
                return False

        self.get_logger().info("Executing calculated straight line Cartesian trajectory...")
        execute_goal = ExecuteTrajectory.Goal()
        execute_goal.trajectory = res.solution 

        send_goal_future = self._execute_client.send_goal_async(execute_goal)
        rclpy.spin_until_future_complete(self, send_goal_future)
        
        goal_handle = send_goal_future.result()
        if goal_handle and goal_handle.accepted:
            get_result_future = goal_handle.get_result_async()
            rclpy.spin_until_future_complete(self, get_result_future)
            self.get_logger().info(" Straight camera sweep complete!")
            return True
        else:
            self.get_logger().error("Trajectory execution rejected.")
            return False

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

    # 1. Send to starting exploration zone
    success = node.send_pt_to_pt_goal(-0.4, -0.3, 0.5)
    
    if success:
        time.sleep(2.0) 
        
        # 2. Call the sweep without passing any XYZ parameters! 
        # It handles the math completely internally.
        #node.execute_straight_camera_sweep(distance=0.2)

    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()