server.py

"""
远程操作服务器模块
"""

import os
import socket
from flask import Flask, render_template, jsonify
from flask_socketio import SocketIO, emit
from typing import Optional, Dict, Any

from .data_handler import IMUDataHandler


class TeleoperationServer:
    """远程操作服务器类"""
    
    def __init__(self, host: str = '0.0.0.0', port: int = 4443, debug: bool = False):
        """初始化服务器
        
        Args:
            host: 服务器主机地址
            port: 服务器端口
            debug: 是否开启调试模式
        """
        self.host = host
        self.port = port
        self.debug = debug
        
        # 初始化Flask应用
        self.app = Flask(__name__, 
                        template_folder=os.path.join(os.path.dirname(__file__), 'templates'))
        self.app.config['SECRET_KEY'] = 'your-secret-key'
        
        # 关闭Werkzeug的HTTP请求日志
        if not debug:
            import logging
            log = logging.getLogger('werkzeug')
            log.setLevel(logging.ERROR)
        
        # 初始化SocketIO
        self.socketio = SocketIO(self.app, cors_allowed_origins="*")
        
        # 初始化数据处理器
        self.data_handler = IMUDataHandler()
        
        # 注册路由和事件处理器
        self._register_routes()
        self._register_socket_events()
    
    def _register_routes(self) -> None:
        """注册Flask路由"""
        
        @self.app.route('/')
        def index():
            """主页路由"""
            return render_template('index.html')
        
        @self.app.route('/api/imu_data')
        def get_imu_data():
            """REST API获取当前IMU数据"""
            return jsonify(self.data_handler.get_combined_data())
        
        @self.app.route('/api/control_data')
        def get_control_data():
            """REST API获取当前控制数据"""
            return jsonify(self.data_handler.get_control_data())
    
    def _register_socket_events(self) -> None:
        """注册SocketIO事件处理器"""
        
        @self.socketio.on('connect')
        def handle_connect():
            """客户端连接事件"""
            count = self.data_handler.update_device_count(1)
            emit('device_count_update', {'count': count}, broadcast=True)
            print(f"Client connected, current connections: {count}")
        
        @self.socketio.on('disconnect')
        def handle_disconnect():
            """客户端断开连接事件"""
            count = self.data_handler.update_device_count(-1)
            emit('device_count_update', {'count': count}, broadcast=True)
            print(f"Client disconnected, current connections: {count}")
        
        @self.socketio.on('imu_data')
        def handle_imu_data(data):
            """处理接收到的IMU数据"""
            try:
                if data['type'] == 'motion':
                    self.data_handler.process_motion_data(data)
                elif data['type'] == 'orientation':
                    self.data_handler.process_orientation_data(data)
                
                # 广播更新的数据给所有连接的客户端
                emit('imu_data_update', self.data_handler.get_current_data(), broadcast=True)
                
            except Exception as e:
                print(f"Error processing IMU data: {e}")
        
        @self.socketio.on('control_data')
        def handle_control_data(data):
            """处理接收到的控制数据"""
            try:
                if data['type'] == 'control':
                    self.data_handler.process_control_data(data)
                    # 可以选择是否广播控制数据更新
                    # emit('control_data_update', data, broadcast=True)
                
            except Exception as e:
                print(f"Error processing control data: {e}")
        
        @self.socketio.on('set_reference_pose')
        def handle_set_reference_pose():
            """处理设置参考姿态请求"""
            try:
                self.data_handler.set_reference_pose()
                emit('reference_pose_set', {'status': 'success'})
            except Exception as e:
                print(f"Error setting reference pose: {e}")
                emit('reference_pose_set', {'status': 'error', 'message': str(e)})
        
        @self.socketio.on('request_device_count')
        def handle_device_count_request():
            """处理设备计数请求"""
            emit('device_count_update', {'count': self.data_handler.get_device_count()})
        
        @self.socketio.on('switch_active_arm')
        def handle_switch_active_arm():
            """处理切换活跃手臂请求"""
            # 这里可以添加逻辑来通知连接的设备切换手臂
            emit('arm_switched', {'status': 'success'}, broadcast=True)
            print("Active arm switch requested")
        
        @self.socketio.on('switch_active_robot')
        def handle_switch_active_robot():
            """处理切换活跃机器人请求"""
            # 这里可以添加逻辑来通知连接的设备切换机器人
            emit('robot_switched', {'status': 'success'}, broadcast=True)
            print("Active robot switch requested")
    
    def get_local_ip(self) -> str:
        """获取本机局域网IP地址"""
        try:
            # 连接到外部地址来获取本机IP
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except Exception:
            return "127.0.0.1"
    
    def _print_startup_info(self) -> None:
        """打印启动信息"""
        local_ip = self.get_local_ip()
        
        print("=" * 60)
        print("🚀 Starting Mobile IMU Data Acquisition Server...")
        print("=" * 60)
        print(f"📱 Access from mobile: https://{local_ip}:{self.port}")
        print(f"💻 Access from computer: https://localhost:{self.port}")
        print("=" * 60)
        print("📋 Instructions:")
        print("1. Ensure your phone and computer are on the same local network")
        print("2. Open the address above in your mobile browser")
        print("3. Click the 'Start Collection' button")
        print("4. Allow the browser to access sensor permissions")
        print("5. Move your phone to see real-time IMU data")
        print("6. Use virtual buttons to control the robot")
        print("=" * 60)
    
    def run(self) -> None:
        """启动服务器"""
        self._print_startup_info()
        
        try:
            self.socketio.run(self.app, host=self.host, port=self.port, 
                            debug=self.debug, ssl_context='adhoc')
        except KeyboardInterrupt:
            print("\n👋 Server stopped")
        except Exception as e:
            print(f"❌ Server failed to start: {e}")
            print("💡 Please check if the port is in use, or try another port")