import grpc
import numpy as np
from typing import List, Optional, Tuple
from typing import Dict

from .enums import CMVRErrorCode, RobotCartesian, RobotJointIndexDirection
from .models import JointCmd, Pose3D, JointState

class HumanoidRobotClient:
    """人形机器人客户端"""

    def __init__(self, device_id: str, stub):
        """
        初始化机器人客户端
        
        参数:
            device_id: 设备ID
            stub: gRPC stub实例
        """
        self.device_id = device_id
        self.stub = stub
        self.generated = None  # 推迟导入

    def _import_generated(self):
        """推迟导入 generated 模块"""
        if self.generated is None:
            try:
                # 从正确的路径导入生成的模块
                from generated.cmvr.api import common_pb2, humanoid_robot_command_pb2
                self.generated = type('GeneratedModules', (), {
                    'common_pb2': common_pb2,
                    'humanoid_robot_command_pb2': humanoid_robot_command_pb2
                })
            except ImportError as e:
                print(f"导入生成的模块失败: {e}")
                print("请确保已生成 protobuf 代码")
                raise
        return self.generated

    def _create_command_header(self):
        """创建命令头"""
        generated = self._import_generated()
        header = generated.common_pb2.CommandHeader.Request()
        header.device_id = self.device_id
        header.timestamp.GetCurrentTime()
        return header

    # ==================== 基本控制接口 ====================

    def torqueOn(self) -> CMVRErrorCode:
        try:
            generated = self._import_generated()
            request = generated.common_pb2.CommandHeader.Request()
            request.CopyFrom(self._create_command_header())

            response = self.stub.torqueOn(request)

            return CMVRErrorCode.CMVR_SUCCESS if response.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"开启使能失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    def torqueOff(self) -> CMVRErrorCode:
        try:
            generated = self._import_generated()
            request = generated.common_pb2.CommandHeader.Request()
            request.CopyFrom(self._create_command_header())

            response = self.stub.torqueOff(request)

            return CMVRErrorCode.CMVR_SUCCESS if response.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"关闭使能失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    # ==================== 关节空间运动 ====================

    def moveJ(self, joint_commands: List[JointCmd], 
               overall_vel: float = 0.1, overall_acc: float = 0.5) -> Tuple[CMVRErrorCode, Optional[str]]:
        try:
            generated = self._import_generated()
            # 创建MoveJ请求
            request = generated.humanoid_robot_command_pb2.MoveJ.Request()
            request.header.CopyFrom(self._create_command_header())
            request.vel = overall_vel
            request.acc = overall_acc

            # 添加关节命令
            for cmd in joint_commands:
                joint_cmd = request.cmds.add()
                joint_cmd.joint_name = cmd.joint_name
                joint_cmd.rad = cmd.rad
                joint_cmd.vel = cmd.vel

            # 调用RPC接口
            response = self.stub.moveJ(request)

            return CMVRErrorCode.CMVR_SUCCESS if response.header.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            return CMVRErrorCode.CMVR_RPC_FAILED

    # ==================== 直线空间运动 ====================

    def moveL(self, ee_link: str, target_pose: Pose3D, 
               vel: float = 0.1, acc: float = 0.5) -> Tuple[CMVRErrorCode, Optional[str]]:
        """
        直线空间运动 (MoveL)
        
        参数:
            ee_link: 末端执行器连杆名称
            target_pose: 目标位姿 (Pose3D)
            vel: 速度 (m/s)
            acc: 加速度 (m/s²)
            
        返回:
            Tuple[CMVRErrorCode, Optional[str]]: 错误码和错误消息
        """
        try:
            generated = self._import_generated()
            # 创建MoveL请求
            request = generated.humanoid_robot_command_pb2.MoveL.Request()
            request.header.CopyFrom(self._create_command_header())
            request.ee_link = ee_link
            request.vel = vel
            request.acc = acc

            # 设置目标位姿
            request.target_pose.x = target_pose.x
            request.target_pose.y = target_pose.y
            request.target_pose.z = target_pose.z
            request.target_pose.rx = target_pose.rx
            request.target_pose.ry = target_pose.ry
            request.target_pose.rz = target_pose.rz

            # 调用RPC接口
            response = self.stub.moveL(request)
            return CMVRErrorCode.CMVR_SUCCESS if response.header.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"直线空间运动失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED

    # ==================== 关节速度控制 ====================

    def speedJ(self, joint_name: str, velocity: float, 
                direction: RobotJointIndexDirection = RobotJointIndexDirection.FORWARD,
                acc: float = 0.5) -> Tuple[CMVRErrorCode, Optional[str]]:
        """
        关节速度控制 (SpeedJ)
        
        参数:
            joint_name: 关节名称
            velocity: 速度 (rad/s)
            direction: 运动方向
            acc: 加速度 (rad/s²)
            
        返回:
            Tuple[CMVRErrorCode, Optional[str]]: 错误码和错误消息
        """
        try:
            generated = self._import_generated()
            # 创建SpeedJ请求
            request = generated.humanoid_robot_command_pb2.SpeedJ.Request()
            request.header.CopyFrom(self._create_command_header())
            request.joint_name = joint_name
            request.vel = abs(velocity)  # 确保速度为正值
            request.acc = acc
            request.direction = direction.value

            # 调用RPC接口
            response = self.stub.speedJ(request)
            return CMVRErrorCode.CMVR_SUCCESS if response.header.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"关节速度控制失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED, str(e)
        
    # ==================== 笛卡尔速度控制 ====================

    def speedL(self, ee_link: str, velocity: float, 
                cartesian: RobotCartesian = RobotCartesian.X,
                direction: RobotJointIndexDirection = RobotJointIndexDirection.FORWARD,
                acc: float = 0.5) -> Tuple[CMVRErrorCode, Optional[str]]:
        """
        笛卡尔速度控制 (SpeedL)
        
        参数:
            ee_link: 末端执行器连杆名称
            velocity: 速度 (m/s 或 rad/s，取决于cartesian类型)
            cartesian: 笛卡尔坐标轴
            direction: 运动方向
            acc: 加速度 (m/s² 或 rad/s²)
            
        返回:
            Tuple[CMVRErrorCode, Optional[str]]: 错误码和错误消息
        """
        try:
            generated = self._import_generated()
            # 创建SpeedL请求
            request = generated.humanoid_robot_command_pb2.SpeedL.Request()
            request.header.CopyFrom(self._create_command_header())
            request.ee_link = ee_link
            request.vel = abs(velocity)
            request.acc = acc
            request.cartesian = cartesian.value
            request.direction = direction.value

            # 调用RPC接口
            response = self.stub.speedL(request)
            return CMVRErrorCode.CMVR_SUCCESS if response.header.success else CMVRErrorCode.CMVR_RPC_FAILED

        except grpc.RpcError as e:
            print(f"笛卡尔速度控制失败: {e}")
            return CMVRErrorCode.CMVR_RPC_FAILED, str(e)

    # ==================== 状态获取 ====================

    def get_joint_states(self) -> Tuple[CMVRErrorCode, List[JointState]]:
        """
        获取机器人所有关节状态
        
        返回:
            Tuple[CMVRErrorCode, List[JointState]]: 错误码和关节状态对象列表
        """
        try:
            generated = self._import_generated()
            # 创建关节状态请求
            request = generated.humanoid_robot_command_pb2.JointRequest()
            request.header.CopyFrom(self._create_command_header())

            # 调用RPC接口
            response = self.stub.getJointState(request)

            if not response.header.success:
                print(f"RPC调用失败: {response.header}")
                return CMVRErrorCode.CMVR_RPC_FAILED, []

            # 检查是否有状态数据
            if not response.state:
                print("响应中没有关节状态数据")
                return CMVRErrorCode.CMVR_NO_DATA, []

            # 解析所有JointState
            joint_states = []
            for proto_joint_state in response.state:
                # 解析响应数据到JointState对象
                # 注意proto字段名和Python dataclass字段名的映射
                joint_state = JointState(
                    name=list(proto_joint_state.name),           # proto: name -> Python: name
                    position=list(proto_joint_state.position),   # proto: position -> Python: position
                    velocity=list(proto_joint_state.velocity),   # proto: velocity -> Python: velocity  
                    effort=list(proto_joint_state.effort),       # proto: effort -> Python: effort
                    timestamp=proto_joint_state.timestamp        # proto: timestamp -> Python: timestamp
                )
                joint_states.append(joint_state)

            if not joint_states:
                print("解析后没有有效的关节状态数据")
                return CMVRErrorCode.CMVR_NO_DATA, []

            return CMVRErrorCode.CMVR_SUCCESS, joint_states

        except grpc.RpcError as e:
            print(f"gRPC错误 - 获取关节状态失败: {e}")
            if hasattr(e, 'code'):
                if e.code() == grpc.StatusCode.NOT_FOUND:
                    return CMVRErrorCode.CMVR_NOT_FOUND, []
                elif e.code() == grpc.StatusCode.UNAVAILABLE:
                    return CMVRErrorCode.CMVR_CONNECTION_FAILED, []
            return CMVRErrorCode.CMVR_RPC_FAILED, []
        
        except Exception as e:
            print(f"未知错误 - 获取关节状态失败: {e}")
            import traceback
            traceback.print_exc()
            return CMVRErrorCode.CMVR_UNKNOWN_ERROR, []