#!/usr/bin/env python3
"""机械臂控制客户端示例

演示如何调用机械臂控制服务
"""

import rclpy
from rclpy.node import Node
from arm_control_msgs.srv import MoveJoints, MovePose, GetPose, SetGripper


class ArmControlClient(Node):
    """机械臂控制客户端"""

    def __init__(self):
        super().__init__('arm_control_client')

        # 创建服务客户端
        self.cli_move_joints = self.create_client(MoveJoints, 'arm_control/move_joints')
        self.cli_move_pose = self.create_client(MovePose, 'arm_control/move_pose')
        self.cli_get_pose = self.create_client(GetPose, 'arm_control/get_pose')
        self.cli_set_gripper = self.create_client(SetGripper, 'arm_control/set_gripper')

        # 等待服务可用
        self.get_logger().info('等待服务...')
        self.cli_move_joints.wait_for_service()
        self.cli_move_pose.wait_for_service()
        self.cli_get_pose.wait_for_service()
        self.cli_set_gripper.wait_for_service()
        self.get_logger().info('服务已连接')

    def get_current_pose(self):
        """查询当前位姿"""
        req = GetPose.Request()
        future = self.cli_get_pose.call_async(req)
        rclpy.spin_until_future_complete(self, future)

        if future.result().success:
            result = future.result()
            self.get_logger().info(
                f'当前位姿: x={result.x:.1f}, y={result.y:.1f}, '
                f'z={result.z:.1f}, phi={result.phi:.1f}°'
            )
            self.get_logger().info(
                f'关节角度: height={result.height}, j2={result.j2}, '
                f'j3={result.j3}, j4={result.j4}, j5={result.j5}, j6={result.j6}'
            )
            return result
        else:
            self.get_logger().error(f'查询失败: {future.result().message}')
            return None

    def move_to_joints(self, height, j2, j3, j4, j5=81, j6=30, duration=2.0):
        """关节空间运动"""
        req = MoveJoints.Request()
        req.height = height
        req.j2 = j2
        req.j3 = j3
        req.j4 = j4
        req.j5 = j5
        req.j6 = j6
        req.duration = duration

        self.get_logger().info(f'关节运动: height={height}, j2={j2}, j3={j3}, j4={j4}')
        future = self.cli_move_joints.call_async(req)
        rclpy.spin_until_future_complete(self, future)

        if future.result().success:
            self.get_logger().info('运动完成')
            return True
        else:
            self.get_logger().error(f'运动失败: {future.result().message}')
            return False

    def move_to_pose(self, x, y, z, phi, duration=2.0, grip=False, release=False):
        """笛卡尔空间运动"""
        req = MovePose.Request()
        req.x = x
        req.y = y
        req.z = z
        req.phi = phi
        req.duration = duration
        req.grip = grip
        req.release = release
        req.gripper_state = MovePose.Request.GRIPPER_KEEP
        req.elbow_up = False

        self.get_logger().info(f'位姿运动: ({x:.1f}, {y:.1f}, {z:.1f}), phi={phi:.1f}°')
        future = self.cli_move_pose.call_async(req)
        rclpy.spin_until_future_complete(self, future)

        if future.result().success:
            self.get_logger().info('运动完成')
            return True
        else:
            self.get_logger().error(f'运动失败: {future.result().message}')
            return False

    def set_gripper(self, grip=False, release=False):
        """夹爪控制"""
        req = SetGripper.Request()
        req.grip = grip
        req.release = release
        req.gripper_state = SetGripper.Request.GRIPPER_KEEP

        action = "抓取" if grip else ("释放" if release else "保持")
        self.get_logger().info(f'夹爪控制: {action}')
        future = self.cli_set_gripper.call_async(req)
        rclpy.spin_until_future_complete(self, future)

        if future.result().success:
            self.get_logger().info('夹爪控制完成')
            return True
        else:
            self.get_logger().error(f'夹爪控制失败: {future.result().message}')
            return False


def demo_sequence(client):
    """演示抓取流程"""
    print("\n" + "="*60)
    print("演示序列：查询 → 移动 → 抓取")
    print("="*60 + "\n")

    # 1. 查询当前位姿
    print("1. 查询当前位姿...")
    client.get_current_pose()

    # 2. 移动到上方
    print("\n2. 移动到物体上方...")
    client.move_to_pose(x=200.0, y=100.0, z=-50.0, phi=45.0, duration=2.0, release=True)

    # 3. 下降到抓取位置
    print("\n3. 下降到抓取位置...")
    client.move_to_pose(x=200.0, y=100.0, z=-150.0, phi=45.0, duration=1.0, release=True)

    # 4. 抓取
    print("\n4. 执行抓取...")
    client.set_gripper(grip=True)

    # 5. 提升
    print("\n5. 提升物体...")
    client.move_to_pose(x=200.0, y=100.0, z=-50.0, phi=45.0, duration=1.0, grip=True)

    # 6. 移动到目标位置
    print("\n6. 移动到目标位置...")
    client.move_to_pose(x=100.0, y=200.0, z=-50.0, phi=90.0, duration=2.0, grip=True)

    # 7. 下降
    print("\n7. 下降...")
    client.move_to_pose(x=100.0, y=200.0, z=-150.0, phi=90.0, duration=1.0, grip=True)

    # 8. 释放
    print("\n8. 释放物体...")
    client.set_gripper(release=True)

    # 9. 提升
    print("\n9. 提升...")
    client.move_to_pose(x=100.0, y=200.0, z=-50.0, phi=90.0, duration=1.0, release=True)

    print("\n" + "="*60)
    print("演示完成！")
    print("="*60 + "\n")


def main():
    rclpy.init()
    client = ArmControlClient()

    try:
        # 运行演示序列
        demo_sequence(client)

    except KeyboardInterrupt:
        pass
    finally:
        client.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
