#!/usr/bin/env python3
"""
机械臂客户端
连接到手部检测服务器，接收控制信号并控制机械臂
"""

import asyncio
import json
import logging
import time
from typing import Dict, Any, Optional

import socketio
from ZDT_Servo import Emm42Driver

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)



class RobotArmClient:
    def _start_watchdog(self, timeout=0.2):
        import threading
        self._watchdog_timeout = timeout
        self._watchdog_timer = None
        def stop_all_motors():
            logger.warning("控制信号超时，自动停止所有电机！")
            for m in self.motors.values():
                if m:
                    try:
                        m.set_speed_mode(direction=True, speed_rpm=0, acceleration=5)
                    except Exception as e:
                        logger.error(f"电机停止异常: {e}")
        self._watchdog_stop_all = stop_all_motors
        def reset_timer():
            if self._watchdog_timer:
                self._watchdog_timer.cancel()
            self._watchdog_timer = threading.Timer(self._watchdog_timeout, self._watchdog_stop_all)
            self._watchdog_timer.daemon = True
            self._watchdog_timer.start()
        self._watchdog_reset = reset_timer
        self._watchdog_reset()
    def _init_motor_read_idx(self):
        if not hasattr(self, '_motor_read_idx'):
            self._motor_read_idx = 1
    """机械臂客户端"""
    
    def __init__(self, server_url='http://localhost:5000'):
        self.server_url = server_url
        self.sio = socketio.Client()
        self.is_connected = False
        self.current_position = {
            'x_angle': 90,
            'y_angle': 90,
            'z_angle': 90,
            'grip': 0
        }
        # 初始化3个电机
        self.motors = {}
        # 初始化限位标志位 {电机ID: 是否达到限位}
        self.limit_flags = {1: False, 2: False, 3: False}
        for i in range(1, 4):
            try:
                self.motors[i] = Emm42Driver(port='/dev/ttyAMA5', device_id=i, baudrate=115200)
                self.motors[i].enable_motor(enable=True)
            except Exception as e:
                self.motors[i] = None
                print(f"电机{i}初始化失败: {e}")

        # 启动看门狗定时器
        self._start_watchdog(timeout=0.2)



        
        self.setup_events()
    
    def setup_events(self):
        """设置事件处理器"""
        
        @self.sio.event
        def connect():
            """连接成功"""
            self.is_connected = True
            logger.info("已连接到手部检测服务器")
            
            # 注册为机械臂客户端
            self.sio.emit('register_client', {
                'type': 'robot'
            })
        
        @self.sio.event
        def disconnect():
            """连接断开"""
            self.is_connected = False
            logger.info("与服务器断开连接")
        
        @self.sio.event
        def registration_success(data):
            """注册成功"""
            logger.info(f"客户端注册成功: {data}")
        
        @self.sio.event
        def robot_control(data):
            """接收机械臂控制信号"""
            try:
                self.process_control_signal(data)
            except Exception as e:
                logger.error(f"处理控制信号时出错: {e}")
        
        @self.sio.event
        def error(data):
            """错误处理"""
            logger.error(f"服务器错误: {data}")
    

    
    def process_control_signal(self, control_signal: Dict[str, Any]):
        # 重置看门狗定时器，防止超时自动停机
        if hasattr(self, '_watchdog_reset'):
            self._watchdog_reset()
        """处理控制信号"""
        # logger.info(f"接收到控制信号: {control_signal}")
        
        # 检查信号变化
        x_changed = abs(control_signal['x_angle'] - self.current_position['x_angle']) > 1
        y_changed = abs(control_signal['y_angle'] - self.current_position['y_angle']) > 1
        z_changed = abs(control_signal['z_angle'] - self.current_position['z_angle']) > 1
        grip_changed = control_signal['grip'] != self.current_position['grip']
        
        # 如果有变化，执行相应动作
        if x_changed or y_changed or z_changed:
            self.move_arm(control_signal)
        
        if grip_changed:
            self.control_gripper(control_signal['grip'])
        
        # 更新当前位置
        self.current_position = {
            'x_angle': control_signal['x_angle'],
            'y_angle': control_signal['y_angle'],
            'z_angle': control_signal['z_angle'],
            'grip': control_signal['grip']
        }

    
    def move_arm(self, control_signal: Dict[str, Any]):
        """移动机械臂，支持左右和上下跟随（不再负责电机位置读取和限位）"""
        x_angle = control_signal['x_angle']
        y_angle = control_signal['y_angle']
        z_angle = control_signal['z_angle']
        speed = control_signal['speed']
        logger.info(f"移动机械臂: X={x_angle:.1f}°, Y={y_angle:.1f}°, 速度={speed}")
        
        m1 = self.motors.get(1)
        m3 = self.motors.get(3)
        
        # 1. 底座左右跟随 (电机1)
        if not self.limit_flags[1]:  # 只有未达到限位时才执行
            base_angle = x_angle - 90
            if base_angle > 0:
                dir1 = True
            else:
                dir1 = False
                base_angle = -base_angle
            if m1:
                try:
                    m1.set_speed_mode(direction=dir1, speed_rpm=base_angle * 0.1, acceleration=5)
                except Exception as e:
                    print(f"ID1异常:{e}")
        
        # 2. 小臂上下 (电机3)
        if not self.limit_flags[3]:  # 只有未达到限位时才执行
            y_angle_cmd = y_angle - 90
            if y_angle_cmd > 0:
                dir3 = True
            else:
                dir3 = False
                y_angle_cmd = -y_angle_cmd
            if m3:
                try:
                    m3.set_speed_mode(direction=dir3, speed_rpm=y_angle_cmd * 0.1, acceleration=5)
                except Exception as e:
                    print(f"ID3异常:{e}")

    def poll_motor_and_limit(self):
        """主循环调用：顺序读取电机位置并做限位保护"""
        self._init_motor_read_idx()
        m1 = self.motors.get(1)
        m2 = self.motors.get(2)
        m3 = self.motors.get(3)
        try:
            if self._motor_read_idx == 1:
                pos1 = m1.read_position() if m1 else None
                
                if pos1 is not None:
                    logger.info(f"电机1脉冲: {pos1}")
                    # 限位保护
                    if (pos1 >= 190000 or pos1 <= -190000):
                        logger.warning(f"电机1到达限位，强制停止")
                        self.limit_flags[1] = True
                        try:
                            m1.set_speed_mode(direction=True, speed_rpm=0, acceleration=5)
                        except Exception:
                            pass
                    else:
                        self.limit_flags[1] = False
                self._motor_read_idx = 2
            elif self._motor_read_idx == 2:
                pos2 = m2.read_position() if m2 else None

                if pos2 is not None:
                    logger.info(f"电机2脉冲: {pos2}")  
                    # 电机2限位保护
                    if (pos2 >= 190000 or pos2 <= -190000):
                        logger.warning(f"电机2到达限位，强制停止")
                        self.limit_flags[2] = True
                        try:
                            m2.set_speed_mode(direction=True, speed_rpm=0, acceleration=5)
                        except Exception:
                            pass
                    else:
                        self.limit_flags[2] = False
                # 可加电机2限位保护
                self._motor_read_idx = 3

            elif self._motor_read_idx == 3:
                pos3 = m3.read_position() if m3 else None
                logger.info(f"电机3脉冲: {pos3}")
                # 限位保护示例
                if pos3 is not None:
                    if pos3 >= 1000:
                        self.limit_flags[3] = True
                        logger.warning(f"电机3下限，自动反向脱限")
                        try:
                            # 反向（向上）低速脱限
                            m3.set_speed_mode(direction=True, speed_rpm=2, acceleration=5)
                        except Exception:
                            pass
                    elif pos3 <= -180000:
                        self.limit_flags[3] = True
                        logger.warning(f"电机3上限，自动反向脱限")
                        try:
                            # 反向（向下）低速脱限
                            m3.set_speed_mode(direction=False, speed_rpm=2, acceleration=5)
                        except Exception:
                            pass
                    else:
                        # 安全区间，正常可停 
                        try:
                            self.limit_flags[3] = False
                            m3.set_speed_mode(direction=True, speed_rpm=0, acceleration=5)
                        except Exception:
                            pass
                self._motor_read_idx = 1
        except Exception as e:
            logger.error(f"电机读取异常: {e}")

        




    def control_gripper(self, grip_state: int):
        """控制抓取器"""
        action = "抓取" if grip_state == 1 else "松开"
        logger.info(f"控制抓取器: {action}")
        
        # 这里应该是实际的抓取器控制代码
        # 例如：
        # self.gripper_controller.set_state(grip_state)
        
        # 模拟控制延迟
        # time.sleep(0.05)
    
    def connect(self):
        """连接到服务器"""
        try:
            logger.info(f"正在连接到服务器: {self.server_url}")
            self.sio.connect(self.server_url)
            return True
        except Exception as e:
            logger.error(f"连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开连接"""
        if self.is_connected:
            self.sio.disconnect()
    
    def run(self):
        """运行客户端"""
        if self.connect():
            try:
                # 保持连接
                while self.is_connected:
                    # 主动轮询电机位置和限位保护
                    self.poll_motor_and_limit()
                    time.sleep(0.05)

                    # 发送心跳
                    if self.is_connected:
                        self.sio.emit('ping')
            except KeyboardInterrupt:
                logger.info("收到中断信号，正在退出...")
            finally:
                self.disconnect()
        else:
            logger.error("无法连接到服务器")

class MockRobotArmController:
    """模拟机械臂控制器"""
    
    def __init__(self):
        self.current_angles = [90, 90, 90]  # X, Y, Z轴角度
        self.gripper_state = 0  # 0=松开, 1=抓取
        
    def move_to_angles(self, x_angle: float, y_angle: float, z_angle: float, speed: int):
        """移动到指定角度"""
        logger.info(f"模拟移动: X={x_angle:.1f}°, Y={y_angle:.1f}°, Z={z_angle:.1f}°")
        
        # 模拟平滑移动
        target_angles = [x_angle, y_angle, z_angle]
        
        for i in range(3):
            step = (target_angles[i] - self.current_angles[i]) / 10
            for _ in range(10):
                self.current_angles[i] += step
                time.sleep(0.01)  # 模拟移动时间
        
        self.current_angles = target_angles
        logger.info(f"移动完成: {self.current_angles}")
    
    def set_gripper_state(self, state: int):
        """设置抓取器状态"""
        self.gripper_state = state
        action = "抓取" if state == 1 else "松开"
        logger.info(f"抓取器状态: {action}")

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='机械臂客户端')
    parser.add_argument('--server', default='http://localhost:5000', help='服务器地址')
    parser.add_argument('--mock', action='store_true', help='使用模拟控制器')
    
    args = parser.parse_args()
    
    # 创建客户端
    client = RobotArmClient(server_url=args.server)
    
    # 如果使用模拟控制器
    if args.mock:
        mock_controller = MockRobotArmController()
        
        # 重写控制方法以使用模拟控制器
        original_move_arm = client.move_arm
        original_control_gripper = client.control_gripper
        
        def mock_move_arm(control_signal):
            mock_controller.move_to_angles(
                control_signal['x_angle'],
                control_signal['y_angle'],
                control_signal['z_angle'],
                control_signal['speed']
            )
        
        def mock_control_gripper(grip_state):
            mock_controller.set_gripper_state(grip_state)
        
        client.move_arm = mock_move_arm
        client.control_gripper = mock_control_gripper
    
    # 运行客户端
    client.run()

if __name__ == '__main__':
    main()