import threading
import time
from typing import Any, Dict
from Robotic_Arm.rm_robot_interface import *
from lerobot.motors import Motor, MotorCalibration, MotorsBus
from .robot_mock import RealmanMock
from scipy.spatial.transform import Rotation as R

import numpy
#动作执行成功
ACTION_SUCCESS = 0


class RealmanMotorsBus(MotorsBus):
    model_number_table={}
    model_ctrl_table={
        "realman":{}
    }
    ip: str
    port: int
    """
        对Realman SDK的二次封装
    """
    def __init__(
        self,
        port: str,
        motors:dict[str, Motor],
        gripper_range:list[int],
        joint: list,
        udp_ip: str,
        udp_port: int,
        mock: bool = False,
        calibration: dict[str, MotorCalibration] | None = None,
    ):
        super().__init__(port,motors, calibration)
        address  = port.split(':')
        self.ip = address[0]
        self.port = int(address[1])
        self.motors = motors
        self.gripper_range = gripper_range
        self.init_joint_position = joint
        self.safe_disable_position = joint
        self.mock = mock
        if mock:
            #self.rmarm = RealmanMock()
            self.rmarm = RoboticArm(rm_thread_mode_e.RM_DUAL_MODE_E)
        else:
            self.rmarm = RoboticArm(rm_thread_mode_e.RM_DUAL_MODE_E)
        self.handle = None
        self.init_pose = numpy.zeros(8)
        self.init_pose[-2] = 1
        self.running = False
        custom = rm_udp_custom_config_t()
        custom.joint_speed = 0
        # custom.lift_state = 1
        custom.hand_state = 1
        custom.plus_state = 1
        custom.plus_base=1
        custom.expand_state = 1
        custom.arm_current_status = 1
        custom.expand_state = 1
        self.udp_config = rm_realtime_push_config_t(1, True, udp_port, 0, udp_ip, custom)
        self.lock = threading.Lock()
        self.temp_pose = numpy.zeros(8)

    @property
    def motor_names(self) -> list[str]:
        return list(self.motors.keys())

    @property
    def motor_models(self) -> list[str]:
        return [model for _, model in self.motors.values()]

    @property
    def motor_indices(self) -> list[int]:
        return [idx for idx, _ in self.motors.values()]
    
    def update_arm_state(self, state:rm_realtime_arm_joint_state_t):
        joint = list(state.joint_status.joint_position)
        gripper = state.expandState.pos
        pose = state.waypoint.position
        euler = state.waypoint.euler
        pose = [pose.x, pose.y, pose.z, euler.rx, euler.ry, euler.rz, 0]
        self.current_arm_state = {"joint": joint, "gripper": gripper, "pose": pose}


    def connect(self, enable=True) -> bool:
        '''
        使能机械臂并检测使能状态,尝试5s,如果使能超时则退出程序
        '''
        self.handle = self.rmarm.rm_create_robot_arm(self.ip, self.port)

        if self.mock:
            self.rmarm.rm_set_arm_run_mode(0)#仿真模式
        else:
            self.rmarm.rm_set_arm_run_mode(1)
        enable_flag = False
        loop_flag = False
        # 设置超时时间（秒）
        timeout = 5
        # 记录进入循环前的时间
        start_time = time.time()
        elapsed_time_flag = False
        
        while not loop_flag:
            elapsed_time = time.time() - start_time
            print("--------------------")
            
            if enable:
                # 获取机械臂状态
                ret = self.rmarm.rm_get_current_arm_state()
                if ret[0] == 0:  # 成功获取状态
                    enable_flag = True
            else:
                enable_flag = False
                # 断开所有连接，销毁线程
                RoboticArm.rm_destory()
            print("使能状态:", enable_flag)
            print("--------------------")
            if(enable_flag == enable):
                loop_flag = True
                enable_flag = True
                break
            else: 
                loop_flag = False
                enable_flag = False
            # 检查是否超过超时时间
            if elapsed_time > timeout:
                print("超时....")
                elapsed_time_flag = True
                enable_flag = True
                break
            time.sleep(1)
        # self.rmarm.rm_set_realtime_push(self.udp_config)
        # self.arm_state_callback = rm_realtime_arm_state_callback_ptr(self.update_arm_state)
        # self.rmarm.rm_realtime_arm_state_call_back(self.arm_state_callback)
        self.rmarm.rm_movej(self.init_joint_position[:-1], 5, 0, 0, 1)
        self.init_pose = self.getQuatPose()
        resp = enable_flag
        self.running = True
        threading.Thread(target=self.update_pose).start()
        threading.Thread(target=self.update_gripper).start()
        print(f"Returning response: {resp}")
        return resp
    
    def getQuatPose(self)->tuple[int,rm_current_arm_state_t]:
        state = rm_current_arm_state_t()
        ret = rm_get_current_arm_state(self.handle, byref(state))
        return numpy.array([
            state.pose.position.x, 
            state.pose.position.y, 
            state.pose.position.z,
            state.pose.quaternion.w,
            state.pose.quaternion.x, 
            state.pose.quaternion.y, 
            state.pose.quaternion.z, 
            self.rmarm.rm_get_gripper_state()[1]['actpos']
        ])
    def read_pose(self):
        return self.init_pose
    
    def motor_names(self):
        return
    
    def read_calibration(self):
        return
    def set_calibration(self):
        return
    
    def revert_calibration(self):
        return

    def apply_calibration(self):
        """
            移动到初始位置
        """
        self.write_arm(target_joint=self.init_joint_position)

    def write(self, data_name: str, motor: str, value, *, normalize: bool = True, num_retry: int = 0):
        return False
    
    def write_arm(self, target_joint: list):
        ret = self.rmarm.rm_movej(self.init_joint_position[:-1], 5, 0, 0, 1)
        if ret != 0: 
            print("movej error:", ret)
            return ret
        
        ret = self.rmarm.rm_set_gripper_position(target_joint[-1], block=False, timeout=2)
        return ret

    def read(self) -> Dict:
        """
            - 机械臂关节消息,单位1度;[-1, 1]
            - 机械臂夹爪消息,[-1, 1]
        """
        joint_state = self.read_pose()
        
        return {
            "joint_1": joint_state[0]/180,
            "joint_2": joint_state[1]/180,
            "joint_3": joint_state[2]/180,
            "joint_4": joint_state[3]/180,
            "joint_5": joint_state[4]/180,
            "joint_6": joint_state[5]/180,
            "gripper": (joint_state[6]-500)/500
        }
    5
    def disconnect(self, disable_torque = True):
        self.safe_disconnect()

    def safe_disconnect(self):
        """ 
            Move to safe disconnect position
        """
        self.write_arm(target_joint=self.safe_disable_position)
        self.running = False
        # 断开所有连接，销毁线程

    """
    向机械臂写入动作
    """
    def sync_write(self,name,actionData: dict[str, Any]):
        if name =="Goal_Position":
            self.write_arm(target_joint=actionData)
        elif name == "Goal_Velocity":#速度透传模式
            values = numpy.array(list(actionData.values()))
            with self.lock:
                self.temp_pose += numpy.array([*values[:6],0,values[-1]])
    def update_pose(self):
        while self.running:
            with self.lock:
                self.temp_pose[-1] = self.temp_pose[-1]*1000
                temp = self.init_pose + self.temp_pose
                newr = R.from_rotvec(self.temp_pose[3:6].tolist())
                oldr = R.from_quat(self.init_pose[3:7].tolist())
                temp[3:7] = (oldr * newr).as_quat()
                self.temp_pose.fill(0)
            ret = self.rmarm.rm_movep_canfd(temp[:-1].tolist(), False)
            if ret == 0:
                self.init_pose = temp
            else:
                print(f"Error: {ret}")
            time.sleep(0.01)
    def clamp(self,value, min_value, max_value):
        return max(min(value, max_value), min_value)
    def update_gripper(self):
        while self.running:
            if self.init_pose[-1]>0:
                self.rmarm.rm_set_gripper_release(500,False,0)
                self.init_pose[-1] = 0
            elif self.init_pose[-1]<0:
                self.rmarm.rm_set_gripper_pick(500,300,False,0)
                self.init_pose[-1] = 0
    
    def sync_read(self, data_name, motors = None, *, normalize = True, num_retry = 0):
        if data_name == "Present_Position":
            return self.read()

    
    @property
    def is_connected(self) -> bool:
        return self.handle is not None
    @property
    def is_calibrated(self) -> bool:
        return True

    def write_calibration(self, calibration_dict, cache = True):
        pass

    def _get_half_turn_homings(self, positions):
        pass
    
    def _find_single_motor(self, motor: str, initial_baudrate: int | None = None) -> tuple[int, int]:
        return 0,self.motors[motor].id

    def _encode_sign(self, data_name, ids_values):
        pass

    def _decode_sign(self, data_name, ids_values):
        pass

    def _split_into_byte_chunks(self, value, length):
        pass

    def broadcast_ping(self, num_retry = 0, raise_on_error = False):
        pass
    def _handshake(self) -> None:
        pass
    def _assert_protocol_is_compatible(self, instruction_name):
        return True


    """
    配置电机
    """
    def configure_motors(self, return_delay_time=0, maximum_acceleration=254, acceleration=254) -> None:
        pass

    def disable_torque(self, motors: int | str | list[str] | None = None, num_retry: int = 0):
        pass

    def _disable_torque(self, motor: int, model: str, num_retry: int = 0):
        pass

    def enable_torque(self, motors = None, num_retry = 0):
        pass