#!/usr/bin/env python3
"""视觉抓取 ROS 节点

基于相机坐标系到基坐标系的变换，实现自动抓取和释放功能。

功能：
1. 抓取服务：输入相机坐标 (x, y, z)，自动转换并执行抓取
2. 释放服务：输入基坐标系位置，移动并释放物体
"""

import math
from typing import Tuple

import numpy as np
import rclpy
from rclpy.node import Node

from arm_control_msgs.srv import MovePose, GetPose, SetGripper
from std_srvs.srv import Trigger
from geometry_msgs.msg import Point


# ============================================================================
# 坐标变换函数（从 camera_to_base.py 复制）
# ============================================================================

def euler_to_rotation_matrix(roll_deg: float, pitch_deg: float, yaw_deg: float) -> np.ndarray:
    """欧拉角转旋转矩阵（ZYX顺序）"""
    roll = math.radians(roll_deg)
    pitch = math.radians(pitch_deg)
    yaw = math.radians(yaw_deg)

    Rx = np.array([
        [1, 0, 0],
        [0, math.cos(roll), -math.sin(roll)],
        [0, math.sin(roll), math.cos(roll)]
    ])

    Ry = np.array([
        [math.cos(pitch), 0, math.sin(pitch)],
        [0, 1, 0],
        [-math.sin(pitch), 0, math.cos(pitch)]
    ])

    Rz = np.array([
        [math.cos(yaw), -math.sin(yaw), 0],
        [math.sin(yaw), math.cos(yaw), 0],
        [0, 0, 1]
    ])

    return Rz @ Ry @ Rx


def camera_to_tcp(
    xc: float, yc: float, zc: float,
    tx: float = 0.0, ty: float = 0.0, tz: float = 0.0,
    roll: float = 0.0, pitch: float = 0.0, yaw: float = 0.0
) -> Tuple[float, float, float]:
    """相机坐标系 → TCP 坐标系"""
    R = euler_to_rotation_matrix(roll, pitch, yaw)
    T = np.array([tx, ty, tz])
    P_cam = np.array([xc, yc, zc])
    P_tcp = R @ P_cam + T
    return float(P_tcp[0]), float(P_tcp[1]), float(P_tcp[2])


def tcp_to_base(
    xt: float, yt: float, zt: float,
    tcp_x: float, tcp_y: float, tcp_z: float, tcp_phi_deg: float
) -> Tuple[float, float, float]:
    """TCP 坐标系 → 机械臂基坐标系（水平相机版本）"""
    phi = math.radians(tcp_phi_deg)

    # 旋转矩阵：TCP → 基坐标系
    # TCP X → 基坐标 -Y
    # TCP Y → 基坐标 -Z
    # TCP Z → 基坐标 X
    R_tcp_to_base = np.array([
        [-math.sin(phi), 0,  math.cos(phi)],   # X_base
        [-math.cos(phi), 0, -math.sin(phi)],   # Y_base (修正：添加负号)
        [0,             -1,  0]                # Z_base
    ])

    P_tcp = np.array([xt, yt, zt])
    P_base_relative = R_tcp_to_base @ P_tcp
    P_base = P_base_relative + np.array([tcp_x, tcp_y, tcp_z])

    return float(P_base[0]), float(P_base[1]), float(P_base[2])


def camera_to_base(
    xc: float, yc: float, zc: float,
    tcp_x: float, tcp_y: float, tcp_z: float, tcp_phi_deg: float,
    cam_tx: float = 0.0, cam_ty: float = 0.0, cam_tz: float = 0.0,
    cam_roll: float = 0.0, cam_pitch: float = 0.0, cam_yaw: float = 0.0
) -> Tuple[float, float, float]:
    """完整变换：相机坐标系 → 基坐标系"""
    xt, yt, zt = camera_to_tcp(xc, yc, zc, cam_tx, cam_ty, cam_tz,
                                cam_roll, cam_pitch, cam_yaw)
    xb, yb, zb = tcp_to_base(xt, yt, zt, tcp_x, tcp_y, tcp_z, tcp_phi_deg)
    return xb, yb, zb


# ============================================================================
# 自定义服务定义
# ============================================================================

# 由于没有预定义服务，我们使用简化的接口
# 实际使用时可以创建自定义 .srv 文件


# ============================================================================
# 视觉抓取节点
# ============================================================================

class VisionGraspNode(Node):
    """视觉抓取 ROS 节点"""

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

        # 声明参数
        self.declare_parameters(
            namespace='',
            parameters=[
                # 相机到 TCP 的变换参数
                ('cam_tx', 0.0),
                ('cam_ty', 0.0),
                ('cam_tz', 0.0),
                ('cam_roll', 0.0),
                ('cam_pitch', 0.0),
                ('cam_yaw', 0.0),
                # 抓取参数
                ('approach_height_offset', 50.0),  # 接近高度偏移 (mm)
                ('retract_position_x', 200.0),     # 回收位置 X
                ('retract_position_y', 0.0),       # 回收位置 Y
                ('grasp_duration', 3.0),           # 抓取移动时长 (秒)
                ('release_duration', 2.0),         # 释放移动时长 (秒)
                ('gripper_duration', 1.0),         # 夹爪动作时长 (秒)
            ]
        )

        # 获取参数
        self.cam_tx = self.get_parameter('cam_tx').value
        self.cam_ty = self.get_parameter('cam_ty').value
        self.cam_tz = self.get_parameter('cam_tz').value
        self.cam_roll = self.get_parameter('cam_roll').value
        self.cam_pitch = self.get_parameter('cam_pitch').value
        self.cam_yaw = self.get_parameter('cam_yaw').value

        self.approach_offset = self.get_parameter('approach_height_offset').value
        self.retract_x = self.get_parameter('retract_position_x').value
        self.retract_y = self.get_parameter('retract_position_y').value
        self.grasp_duration = self.get_parameter('grasp_duration').value
        self.release_duration = self.get_parameter('release_duration').value
        self.gripper_duration = self.get_parameter('gripper_duration').value

        # 创建服务客户端（连接到 arm_control 节点）
        self.move_cli = self.create_client(MovePose, 'arm_control/move_pose')
        self.get_pose_cli = self.create_client(GetPose, 'arm_control/get_pose')
        self.gripper_cli = self.create_client(SetGripper, 'arm_control/set_gripper')

        # 等待服务可用
        self.get_logger().info('等待 arm_control 服务...')
        self.move_cli.wait_for_service(timeout_sec=5.0)
        self.get_pose_cli.wait_for_service(timeout_sec=5.0)
        self.gripper_cli.wait_for_service(timeout_sec=5.0)
        self.get_logger().info('arm_control 服务已连接')

        # 创建订阅者：接收检测结果
        self.grasp_sub = self.create_subscription(
            Point,
            'vision_grasp/grasp_target',
            self.handle_grasp_target,
            10
        )

        self.release_sub = self.create_subscription(
            Point,
            'vision_grasp/release_target',
            self.handle_release_target,
            10
        )

        self.get_logger().info('视觉抓取节点已启动')
        self.get_logger().info('订阅话题:')
        self.get_logger().info('  - /vision_grasp/grasp_target (geometry_msgs/Point)')
        self.get_logger().info('  - /vision_grasp/release_target (geometry_msgs/Point)')

    def get_current_tcp_pose(self) -> Tuple[float, float, float, float, int]:
        """查询当前 TCP 位姿（包括 J5 状态）"""
        req = GetPose.Request()
        future = self.get_pose_cli.call_async(req)

        # 等待结果（不阻塞其他回调）
        import time
        start_time = time.time()
        timeout = 5.0

        while not future.done():
            if time.time() - start_time > timeout:
                raise RuntimeError("获取 TCP 位姿超时")
            time.sleep(0.01)

        result = future.result()
        if not result.success:
            raise RuntimeError(f"获取 TCP 位姿失败: {result.message}")

        return result.x, result.y, result.z, result.phi, result.j5

    def move_to(self, x: float, y: float, z: float, phi: float,
                duration: float, grip: bool = False, release: bool = False) -> bool:
        """移动到指定位置"""
        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

        future = self.move_cli.call_async(req)

        # 等待结果
        import time
        start_time = time.time()
        timeout = duration + 5.0

        while not future.done():
            if time.time() - start_time > timeout:
                self.get_logger().error("移动超时")
                return False
            time.sleep(0.01)

        result = future.result()
        if result is None:
            self.get_logger().error("移动失败：无响应")
            return False

        return result.success

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

        future = self.gripper_cli.call_async(req)

        # 等待结果
        import time
        start_time = time.time()
        timeout = 3.0

        while not future.done():
            if time.time() - start_time > timeout:
                self.get_logger().error("夹爪控制超时")
                return False
            time.sleep(0.01)

        result = future.result()
        if result is None:
            self.get_logger().error("夹爪控制失败：无响应")
            return False

        return result.success

    def handle_grasp_target(self, msg: Point):
        """处理抓取目标"""
        self.get_logger().info(f'收到抓取目标: 相机坐标 ({msg.x:.1f}, {msg.y:.1f}, {msg.z:.1f})')

        # 在单独的线程中执行抓取，避免阻塞回调
        import threading
        thread = threading.Thread(
            target=self._execute_grasp_thread,
            args=(msg.x, msg.y, msg.z)
        )
        thread.start()

    def _execute_grasp_thread(self, x: float, y: float, z: float):
        """在独立线程中执行抓取流程"""
        try:
            # 获取当前 TCP 位姿（包括 J5 状态）
            tcp_x, tcp_y, tcp_z, tcp_phi, j5 = self.get_current_tcp_pose()
            self.get_logger().info(f'当前 TCP: ({tcp_x:.1f}, {tcp_y:.1f}, {tcp_z:.1f}), phi={tcp_phi:.1f}°, j5={j5}°')

            # 图像坐标到相机坐标系转换
            # 相机水平安装，但会随 J5 状态旋转 180°
            # J5 = -100 (闭合) → 夹爪朝上 (UP) → 相机正向
            # J5 = 81 (张开) → 夹爪朝下 (DOWN) → 相机旋转 180°

            if j5 < 0:
                # J5 闭合（UP）：相机正向
                # 图像 Y 向下 → 相机 -Y
                xc = x
                yc = -y
                zc = z
                self.get_logger().info(f'J5={j5}° (UP)，相机正向，相机坐标: ({xc:.1f}, {yc:.1f}, {zc:.1f})')
            else:
                # J5 张开（DOWN）：相机旋转 180°
                # 图像 X,Y 都翻转（相对于 UP 状态）
                xc = -x
                yc = y
                zc = z
                self.get_logger().info(f'J5={j5}° (DOWN)，相机旋转 180°，相机坐标: ({xc:.1f}, {yc:.1f}, {zc:.1f})')

            # 坐标变换：相机 → 基坐标系
            target_x, target_y, target_z = camera_to_base(
                xc, yc, zc,
                tcp_x, tcp_y, tcp_z, tcp_phi,
                self.cam_tx, self.cam_ty, self.cam_tz,
                self.cam_roll, self.cam_pitch, self.cam_yaw
            )

            self.get_logger().info(f'目标基坐标: ({target_x:.1f}, {target_y:.1f}, {target_z:.1f})')

            # 执行抓取流程
            self.execute_grasp(target_x, target_y, target_z, tcp_phi)

        except Exception as e:
            self.get_logger().error(f'抓取失败: {e}')

    def handle_release_target(self, msg: Point):
        """处理释放目标"""
        self.get_logger().info(f'收到释放目标: 基坐标 ({msg.x:.1f}, {msg.y:.1f}, {msg.z:.1f})')

        # 在单独的线程中执行释放，避免阻塞回调
        import threading
        thread = threading.Thread(
            target=self._execute_release_thread,
            args=(msg.x, msg.y, msg.z)
        )
        thread.start()

    def _execute_release_thread(self, x: float, y: float, z: float):
        """在独立线程中执行释放流程"""
        try:
            # 获取当前 TCP 位姿（用于获取 phi）
            _, _, _, tcp_phi, _ = self.get_current_tcp_pose()

            # 执行释放流程
            self.execute_release(x, y, z, tcp_phi)

        except Exception as e:
            self.get_logger().error(f'释放失败: {e}')

    def execute_grasp(self, x: float, y: float, z: float, phi: float):
        """执行抓取流程

        1. release 夹爪 (duration=0)
        2. 移动到目标位置 (duration=3)
        3. grip 夹爪 (duration=1)
        4. 移动到回收位置 (200, 0, 当前z)
        """
        self.get_logger().info('=' * 60)
        self.get_logger().info('开始抓取流程')
        self.get_logger().info('=' * 60)

        # 步骤 1: 释放夹爪
        self.get_logger().info('1. 释放夹爪')
        if not self.set_gripper(release=True):
            self.get_logger().error('释放夹爪失败')
            return

        # 步骤 2: 移动到目标位置
        self.get_logger().info(f'2. 移动到目标位置: ({x:.1f}, {y:.1f}, {z:.1f})')
        if not self.move_to(x, y, z, phi, self.grasp_duration, release=True):
            self.get_logger().error('移动到目标位置失败')
            return

        # 步骤 3: 抓取
        self.get_logger().info('3. 抓取物体')
        if not self.move_to(x, y, z, phi, self.gripper_duration, grip=True):
            self.get_logger().error('抓取失败')
            return

        # 步骤 4: 移动到回收位置
        self.get_logger().info(f'4. 移动到回收位置: ({self.retract_x:.1f}, {self.retract_y:.1f}, {z:.1f})')
        if not self.move_to(self.retract_x, self.retract_y, z, phi, self.grasp_duration, grip=True):
            self.get_logger().error('移动到回收位置失败')
            return

        self.get_logger().info('=' * 60)
        self.get_logger().info('✓ 抓取完成！')
        self.get_logger().info('=' * 60)

    def execute_release(self, x: float, y: float, z: float, phi: float):
        """执行释放流程

        1. 移动到指定位置
        2. release 夹爪 (duration=0)
        3. 回收到 (200, 0, 当前z)
        """
        self.get_logger().info('=' * 60)
        self.get_logger().info('开始释放流程')
        self.get_logger().info('=' * 60)

        # 步骤 1: 移动到释放位置
        self.get_logger().info(f'1. 移动到释放位置: ({x:.1f}, {y:.1f}, {z:.1f})')
        if not self.move_to(x, y, z, phi, self.release_duration, grip=True):
            self.get_logger().error('移动到释放位置失败')
            return

        # 步骤 2: 释放夹爪
        self.get_logger().info('2. 释放物体')
        if not self.set_gripper(release=True):
            self.get_logger().error('释放夹爪失败')
            return

        # 步骤 3: 回收到初始位置
        self.get_logger().info(f'3. 回收到初始位置: ({self.retract_x:.1f}, {self.retract_y:.1f}, {z:.1f})')
        if not self.move_to(self.retract_x, self.retract_y, z, phi, self.release_duration, release=True):
            self.get_logger().error('回收失败')
            return

        self.get_logger().info('=' * 60)
        self.get_logger().info('✓ 释放完成！')
        self.get_logger().info('=' * 60)


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

    # 使用多线程执行器以支持在回调中调用服务
    from rclpy.executors import MultiThreadedExecutor

    node = VisionGraspNode()
    executor = MultiThreadedExecutor()
    executor.add_node(node)

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


if __name__ == '__main__':
    main()
