#!/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, RobotState
from moveit_msgs.srv import GetCartesianPath, GetPositionFK
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. 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"""
        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"] 
        
        robot_state = RobotState()
        robot_state.joint_state = self.current_joint_state
        req.robot_state = robot_state

        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: 
            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 freely, settling horizontally at the end destination"""
        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"
        
        # Define the target position and orientation
        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)
        
        # Enforces the horizontal alignment strictly at the final destination pose
        target_pose.pose.orientation = self.get_horizontal_camera_quaternion()

        # Add target pose to the path request constraints
        goal_constraints = Constraints()
        goal_msg.request.goal_constraints.append(goal_constraints)
        
        # MoveIt handles position and end-point orientation without path constraints
        # By avoiding constraints during movement, the planner should avoid RRTConnect failures.
        goal_msg.request.workspace_parameters.header.frame_id = "base_link"
        
        # Pack the target pose directly into the goal message requests
        from moveit_msgs.msg import PositionConstraint, OrientationConstraint
        
        # Create a simple end-state constraint matching our target pose
        # Position Constraint
        pos_con = PositionConstraint()
        pos_con.header.frame_id = "base_link"
        pos_con.link_name = "tool_link"
        from moveit_msgs.msg import BoundingVolume
        from shape_msgs.msg import SolidPrimitive
        box = SolidPrimitive()
        box.type = SolidPrimitive.BOX
        box.dimensions = [0.01, 0.01, 0.01] # Tight convergence tolerance
        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)
        pos_con.constraint_region = volume
        pos_con.weight = 1.0
        
        # Orientation Constraint for the destination frame
        ori_con = OrientationConstraint()
        ori_con.header.frame_id = "base_link"
        ori_con.link_name = "tool_link"
        ori_con.orientation = target_pose.pose.orientation
        ori_con.absolute_x_axis_tolerance = 0.05
        ori_con.absolute_y_axis_tolerance = 0.05
        ori_con.absolute_z_axis_tolerance = 0.05
        ori_con.weight = 1.0
        
        goal_constraints.position_constraints.append(pos_con)
        goal_constraints.orientation_constraints.append(ori_con)

        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):
        """Fetches live position from the simulation and advances the camera straight forward."""
        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}")
        
        req = GetCartesianPath.Request()
        req.group_name = "arm"
        req.header.frame_id = "base_link"
        req.start_state.is_diff = True 
        
        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     

        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()

    # Coordinates are now evaluated relative to the arm's base link.
    # Reaches 40cm forward, 30cm down toward crops. It will rotate freely during
    # transit and level out perfectly flat once it locks into the target destination.
    success = node.send_pt_to_pt_goal(.2, 0.3, 0.3)
    
    if success:
        time.sleep(2.0) 
        # Execute the forward crawl sweep
        node.execute_straight_camera_sweep(distance=0.2)

    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()