unity2moveit2/unity-project/Assets/Scripts/Communication/ROSTCPBridge.cs

using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using UnityMoveIt2.Core;
using Unity.Robotics.ROSTCPConnector;

namespace UnityMoveIt2.Communication
{
    /// <summary>
    /// ROS TCP通信桥梁 / ROS TCP Communication Bridge
    /// 负责Unity与ROS2之间的TCP通信 / Responsible for TCP communication between Unity and ROS2
    /// </summary>
    public class ROSTCPBridge : MonoBehaviour
    {
        #region Inspector Fields

        [Header("连接配置 / Connection Configuration")]
        [SerializeField]
        [Tooltip("ROS2服务器IP地址 / ROS2 Server IP Address")]
        private string rosIP = "127.0.0.1";

        [SerializeField]
        [Tooltip("ROS2服务器端口 / ROS2 Server Port")]
        private int rosPort = 10000;

        [SerializeField]
        [Tooltip("自动重连 / Auto Reconnect")]
        private bool autoReconnect = true;

        [SerializeField]
        [Tooltip("重连间隔(秒) / Reconnect Interval (seconds)")]
        private float reconnectInterval = 5f;

        [Header("通信设置 / Communication Settings")]
        [SerializeField]
        [Tooltip("发布频率(Hz) / Publish Rate (Hz)")]
        private float publishRate = 50f;

        [SerializeField]
        [Tooltip("消息缓冲区大小 / Message Buffer Size")]
        private int messageBufferSize = 100;

        [SerializeField]
        [Tooltip("超时时间(秒) / Timeout (seconds)")]
        private float timeout = 5f;

        [SerializeField]
        [Tooltip("每帧最大处理消息数 / Max Messages Per Frame")]
        private int maxMessagesPerFrame = 10;

        [Header("调试设置 / Debug Settings")]
        [SerializeField]
        private bool debugMode = false;

        [SerializeField]
        private bool logMessages = false;

        #endregion

        #region Private Fields

        // 连接状态 / Connection state
        private bool isConnected = false;
        private bool isConnecting = false;
        private float lastReconnectAttempt = 0f;
        private int reconnectAttempts = 0;
        private const int MAX_RECONNECT_ATTEMPTS = 10;

        // ROS连接实例 / ROS Connection instance
        private ROSConnection rosConnection;

        // 通信 / Communication
        private float lastPublishTime = 0f;
        private float publishInterval;
        private Queue<ROSMessage> messageQueue;
        private Dictionary<string, Action<object>> messageHandlers;

        // 统计 / Statistics
        private int messagesSent = 0;
        private int messagesReceived = 0;
        private int messagesDropped = 0;

        #endregion

        #region Properties

        /// <summary>
        /// 是否已连接到ROS2 / Is Connected to ROS2
        /// </summary>
        public bool IsConnected => isConnected;

        /// <summary>
        /// 连接状态字符串 / Connection Status String
        /// </summary>
        public string ConnectionStatus
        {
            get
            {
                if (isConnected)
                    return "已连接 / Connected";
                else if (isConnecting)
                    return "连接中... / Connecting...";
                else
                    return "未连接 / Disconnected";
            }
        }

        /// <summary>
        /// 统计信息 / Statistics
        /// </summary>
        public ConnectionStatistics Statistics => new ConnectionStatistics
        {
            messagesSent = this.messagesSent,
            messagesReceived = this.messagesReceived,
            messagesDropped = this.messagesDropped,
            reconnectAttempts = this.reconnectAttempts
        };

        #endregion

        #region Events

        public event Action OnConnected;
        public event Action OnDisconnected;
        public event Action<string> OnConnectionError;
        public event Action<ROSMessage> OnMessageReceived;

        #endregion

        #region Unity Lifecycle

        private void Awake()
        {
            InitializeBridge();
        }

        private void Start()
        {
            StartCoroutine(ConnectToROSAsync());
        }

        private void Update()
        {
            // 处理消息队列 / Process message queue
            ProcessMessageQueue();

            // 发布周期性消息 / Publish periodic messages
            PublishPeriodicMessages();

            // 自动重连 / Auto reconnect
            HandleAutoReconnect();
        }

        private void OnDestroy()
        {
            DisconnectFromROS();
        }

        private void OnApplicationQuit()
        {
            DisconnectFromROS();
        }

        #endregion

        #region Initialization

        private void InitializeBridge()
        {
            Debug.Log("[ROSTCPBridge] 初始化ROS TCP桥梁 / Initializing ROS TCP Bridge");

            publishInterval = 1f / publishRate;
            messageQueue = new Queue<ROSMessage>();
            messageHandlers = new Dictionary<string, Action<object>>();

            // 注册消息处理器 / Register message handlers
            RegisterMessageHandlers();
        }

        private void RegisterMessageHandlers()
        {
            // 这里注册各种ROS消息类型的处理器 / Register handlers for various ROS message types
            RegisterMessageHandler("joint_state", HandleJointStateMessage);
            RegisterMessageHandler("ik_response", HandleIKResponseMessage);
            RegisterMessageHandler("plan_response", HandlePlanResponseMessage);
            RegisterMessageHandler("collision_status", HandleCollisionStatusMessage);
        }

        #endregion

        #region Connection Management

        /// <summary>
        /// 连接到ROS2服务器 / Connect to ROS2 Server
        /// </summary>
        private IEnumerator ConnectToROSAsync()
        {
            if (isConnected || isConnecting)
                yield break;

            isConnecting = true;
            Debug.Log($"[ROSTCPBridge] 正在连接到ROS2服务器: {rosIP}:{rosPort} / Connecting to ROS2 server");

            // 模拟连接延迟 / Simulate connection delay
            yield return new WaitForSeconds(0.5f);

            // 实际项目中,这里应该建立TCP连接
            // In actual project, establish TCP connection here
            bool connectionSuccess = AttemptConnection();

            if (connectionSuccess)
            {
                isConnected = true;
                isConnecting = false;
                reconnectAttempts = 0;
                lastReconnectAttempt = 0f;

                Debug.Log("[ROSTCPBridge] 成功连接到ROS2服务器 / Successfully connected to ROS2 server");
                OnConnected?.Invoke();
            }
            else
            {
                isConnecting = false;
                reconnectAttempts++;

                string errorMsg = $"连接失败(尝试 {reconnectAttempts}/{MAX_RECONNECT_ATTEMPTS}) / " +
                                $"Connection failed (attempt {reconnectAttempts}/{MAX_RECONNECT_ATTEMPTS})";
                Debug.LogWarning($"[ROSTCPBridge] {errorMsg}");
                OnConnectionError?.Invoke(errorMsg);
            }
        }

        private bool AttemptConnection()
        {
            // 实现真实的TCP连接逻辑 / Implement real TCP connection logic
            try
            {
                // 获取或创建ROSConnection单例 / Get or create ROSConnection singleton
                rosConnection = ROSConnection.GetOrCreateInstance();

                if (rosConnection == null)
                {
                    Debug.LogError("[ROSTCPBridge] 无法创建ROSConnection实例 / Failed to create ROSConnection instance");
                    return false;
                }

                // 配置连接参数 / Configure connection parameters
                rosConnection.RosIPAddress = rosIP;
                rosConnection.RosPort = rosPort;
                rosConnection.ConnectOnStart = false; // 手动控制连接 / Manual connection control

                Debug.Log($"[ROSTCPBridge] 尝试真实TCP连接到 / Attempting real TCP connection to {rosIP}:{rosPort}");

                // 执行连接 / Execute connection
                rosConnection.Connect();

                // 立即检查连接状态 (不等待) / Check connection status immediately (no wait)
                // ROS2端期望连接后立即收到握手消息 / ROS2 expects handshake immediately after connection
                bool hasError = rosConnection.HasConnectionError;

                if (!hasError)
                {
                    Debug.Log("[ROSTCPBridge] ✓ 真实TCP连接建立成功 / Real TCP connection established successfully");

                    // 发送握手消息 / Send handshake message
                    if (SendHandshake())
                    {
                        Debug.Log("[ROSTCPBridge] ✓ 握手完成 / Handshake completed");
                        return true;
                    }
                    else
                    {
                        Debug.LogWarning("[ROSTCPBridge] ⚠ 握手失败，但保持连接 / Handshake failed, but keeping connection");
                        return true; // 即使握手失败，TCP连接仍然有效
                    }
                }
                else
                {
                    Debug.LogWarning("[ROSTCPBridge] ✗ TCP连接失败 / TCP connection failed");
                    return false;
                }
            }
            catch (Exception ex)
            {
                Debug.LogError($"[ROSTCPBridge] 连接异常 / Connection exception: {ex.Message}");
                Debug.LogError($"[ROSTCPBridge] 堆栈跟踪 / Stack trace: {ex.StackTrace}");
                return false;
            }
        }

        /// <summary>
        /// 断开ROS2连接 / Disconnect from ROS2
        /// </summary>
        public void DisconnectFromROS()
        {
            if (!isConnected)
                return;

            Debug.Log("[ROSTCPBridge] 断开ROS2连接 / Disconnecting from ROS2");

            // 关闭真实的TCP连接 / Close real TCP connection
            if (rosConnection != null)
            {
                try
                {
                    rosConnection.Disconnect();
                    Debug.Log("[ROSTCPBridge] ✓ TCP连接已关闭 / TCP connection closed");
                }
                catch (Exception ex)
                {
                    Debug.LogWarning($"[ROSTCPBridge] 断开连接时出错 / Error during disconnection: {ex.Message}");
                }
            }

            isConnected = false;
            isConnecting = false;

            OnDisconnected?.Invoke();
        }

        private void HandleAutoReconnect()
        {
            if (!autoReconnect || isConnected || isConnecting)
                return;

            if (reconnectAttempts >= MAX_RECONNECT_ATTEMPTS)
            {
                Debug.LogError("[ROSTCPBridge] 已达到最大重连次数,停止尝试 / Max reconnect attempts reached, stopping");
                autoReconnect = false;
                return;
            }

            if (Time.time - lastReconnectAttempt >= reconnectInterval)
            {
                lastReconnectAttempt = Time.time;
                StartCoroutine(ConnectToROSAsync());
            }
        }

        #endregion

        #region Message Handling

        /// <summary>
        /// 注册消息处理器 / Register Message Handler
        /// </summary>
        public void RegisterMessageHandler(string messageType, Action<object> handler)
        {
            if (messageHandlers.ContainsKey(messageType))
            {
                messageHandlers[messageType] += handler;
            }
            else
            {
                messageHandlers[messageType] = handler;
            }

            if (debugMode)
                Debug.Log($"[ROSTCPBridge] 已注册消息处理器: {messageType} / Registered message handler: {messageType}");
        }

        /// <summary>
        /// 发送消息到ROS2 / Send Message to ROS2
        /// </summary>
        public bool SendMessage(ROSMessage message)
        {
            if (!isConnected)
            {
                Debug.LogWarning("[ROSTCPBridge] 未连接到ROS2,无法发送消息 / Not connected to ROS2, cannot send message");
                return false;
            }

            if (messageQueue.Count >= messageBufferSize)
            {
                messagesDropped++;
                Debug.LogWarning("[ROSTCPBridge] 消息缓冲区已满,丢弃消息 / Message buffer full, dropping message");
                return false;
            }

            messageQueue.Enqueue(message);

            if (logMessages)
                Debug.Log($"[ROSTCPBridge] 消息已加入队列: {message.messageType} / Message queued: {message.messageType}");

            return true;
        }

        /// <summary>
        /// 处理消息队列 / Process Message Queue
        /// </summary>
        private void ProcessMessageQueue()
        {
            if (!isConnected || messageQueue.Count == 0)
                return;

            int processedCount = 0;
            while (messageQueue.Count > 0 && processedCount < maxMessagesPerFrame)
            {
                ROSMessage message = messageQueue.Dequeue();
                SendMessageInternal(message);
                processedCount++;
            }
        }

        private void SendMessageInternal(ROSMessage message)
        {
            // 实际项目中通过TCP发送消息 / Send message via TCP in actual project
            // 这里模拟发送成功 / Simulate successful send here

            messagesSent++;

            if (logMessages)
                Debug.Log($"[ROSTCPBridge] 消息已发送: {message.messageType} / Message sent: {message.messageType}");
        }

        /// <summary>
        /// 接收并处理ROS消息 / Receive and Handle ROS Message
        /// 在实际实现中,这会由接收线程调用 / In actual implementation, called by receive thread
        /// </summary>
        private void ReceiveMessage(ROSMessage message)
        {
            messagesReceived++;

            if (logMessages)
                Debug.Log($"[ROSTCPBridge] 收到消息: {message.messageType} / Message received: {message.messageType}");

            // 调用注册的处理器 / Invoke registered handlers
            if (messageHandlers.ContainsKey(message.messageType))
            {
                try
                {
                    messageHandlers[message.messageType]?.Invoke(message.data);
                }
                catch (Exception ex)
                {
                    Debug.LogError($"[ROSTCPBridge] 处理消息时发生错误 / Error handling message: {ex.Message}");
                }
            }

            OnMessageReceived?.Invoke(message);
        }

        #endregion

        #region Periodic Publishing

        private void PublishPeriodicMessages()
        {
            if (!isConnected)
                return;

            if (Time.time - lastPublishTime < publishInterval)
                return;

            // 这里可以发布周期性的状态消息 / Publish periodic status messages here
            // 例如:关节状态、末端执行器位姿等 / e.g. joint states, end effector pose, etc.

            lastPublishTime = Time.time;
        }

        #endregion

        #region Message Handlers

        private void HandleJointStateMessage(object data)
        {
            if (debugMode)
                Debug.Log("[ROSTCPBridge] 处理关节状态消息 / Handling joint state message");

            // 处理关节状态更新 / Handle joint state updates
        }

        private void HandleIKResponseMessage(object data)
        {
            if (debugMode)
                Debug.Log("[ROSTCPBridge] 处理IK响应消息 / Handling IK response message");

            // 处理IK解算响应 / Handle IK solution response
        }

        private void HandlePlanResponseMessage(object data)
        {
            if (debugMode)
                Debug.Log("[ROSTCPBridge] 处理规划响应消息 / Handling plan response message");

            // 处理路径规划响应 / Handle path planning response
        }

        private void HandleCollisionStatusMessage(object data)
        {
            if (debugMode)
                Debug.Log("[ROSTCPBridge] 处理碰撞状态消息 / Handling collision status message");

            // 处理碰撞检测状态 / Handle collision detection status
        }

        #endregion

        #region Public Service Methods

        /// <summary>
        /// 请求IK解算 / Request IK Solution
        /// </summary>
        public void RequestIKSolution(Vector3 targetPosition, Quaternion targetRotation, Action<float[], bool> callback)
        {
            if (!isConnected)
            {
                Debug.LogWarning("[ROSTCPBridge] 未连接,IK请求失败 / Not connected, IK request failed");
                callback?.Invoke(null, false);
                return;
            }

            // 创建IK请求消息 / Create IK request message
            var request = new ROSMessage
            {
                messageType = "ik_request",
                data = new IKRequest
                {
                    targetPosition = targetPosition,
                    targetRotation = targetRotation,
                    timestamp = Time.time
                }
            };

            SendMessage(request);

            // 在实际实现中,这里会等待响应并调用callback / In actual implementation, wait for response and invoke callback
            // 目前模拟成功响应 / Currently simulate successful response
            StartCoroutine(SimulateIKResponse(callback));
        }

        private IEnumerator SimulateIKResponse(Action<float[], bool> callback)
        {
            yield return new WaitForSeconds(0.1f);
            // 模拟返回一组关节角度 / Simulate returning joint angles
            float[] jointAngles = new float[6] { 0f, 0f, 0f, 0f, 0f, 0f };
            callback?.Invoke(jointAngles, true);
        }

        /// <summary>
        /// 请求路径规划 / Request Path Planning
        /// </summary>
        public void RequestPathPlanning(Vector3 targetPosition, Quaternion targetRotation, Action<bool> callback)
        {
            if (!isConnected)
            {
                Debug.LogWarning("[ROSTCPBridge] 未连接,规划请求失败 / Not connected, planning request failed");
                callback?.Invoke(false);
                return;
            }

            var request = new ROSMessage
            {
                messageType = "plan_request",
                data = new PlanRequest
                {
                    targetPosition = targetPosition,
                    targetRotation = targetRotation,
                    timestamp = Time.time
                }
            };

            SendMessage(request);
            StartCoroutine(SimulatePlanResponse(callback));
        }

        private IEnumerator SimulatePlanResponse(Action<bool> callback)
        {
            yield return new WaitForSeconds(0.2f);
            callback?.Invoke(true);
        }

        #endregion

        #region Handshake Protocol

        /// <summary>
        /// 发送ROS握手消息 / Send ROS Handshake Message
        /// 尝试多种握手格式以提高兼容性 / Try multiple formats for compatibility
        /// </summary>
        private bool SendHandshake()
        {
            try
            {
                Debug.Log("[ROSTCPBridge] ======================================");
                Debug.Log("[ROSTCPBridge] 开始ROS握手流程 / Starting ROS handshake");
                Debug.Log("[ROSTCPBridge] ======================================");

                // 只使用标准ROS-TCP协议握手 (带长度前缀的JSON)
                // Only use standard ROS-TCP protocol handshake (length-prefixed JSON)
                Debug.Log("[ROSTCPBridge] 发送标准ROS-TCP握手 (带长度前缀的JSON)...");
                byte[] handshakeData = ROSHandshakeProtocol.CreateLengthPrefixedHandshake();

                if (handshakeData != null && SendHandshakeBytes(handshakeData))
                {
                    Debug.Log("[ROSTCPBridge] ✓ 握手数据发送成功 / Handshake data sent successfully");
                    Debug.Log("[ROSTCPBridge] ======================================");
                    return true;
                }
                else
                {
                    Debug.LogWarning("[ROSTCPBridge] ✗ 握手数据发送失败 / Handshake data send failed");
                    Debug.Log("[ROSTCPBridge] ======================================");
                    return false;
                }
            }
            catch (Exception ex)
            {
                Debug.LogError($"[ROSTCPBridge] 握手异常 / Handshake exception: {ex.Message}");
                return false;
            }
        }

        /// <summary>
        /// 发送握手字节数据 / Send Handshake Bytes
        /// </summary>
        private bool SendHandshakeBytes(byte[] handshakeData)
        {
            if (handshakeData == null || handshakeData.Length == 0)
            {
                Debug.LogWarning("[ROSTCPBridge] 握手数据为空 / Handshake data is empty");
                return false;
            }

            try
            {
                if (rosConnection == null || rosConnection.HasConnectionError)
                {
                    Debug.LogWarning("[ROSTCPBridge] ROS连接不可用 / ROS connection not available");
                    return false;
                }

                // 发送握手数据
                rosConnection.SendMessage(handshakeData);
                Debug.Log($"[ROSTCPBridge] → 已发送握手数据 / Sent handshake data: {handshakeData.Length} bytes");

                return true;
            }
            catch (Exception ex)
            {
                Debug.LogError($"[ROSTCPBridge] 发送握手数据失败 / Failed to send handshake: {ex.Message}");
                return false;
            }
        }

        #endregion

        #region Utility Methods

        /// <summary>
        /// 重置统计信息 / Reset Statistics
        /// </summary>
        public void ResetStatistics()
        {
            messagesSent = 0;
            messagesReceived = 0;
            messagesDropped = 0;
            reconnectAttempts = 0;
        }

        /// <summary>
        /// 获取连接信息字符串 / Get Connection Info String
        /// </summary>
        public string GetConnectionInfo()
        {
            return $"ROS TCP桥梁 / ROS TCP Bridge\n" +
                   $"服务器 / Server: {rosIP}:{rosPort}\n" +
                   $"状态 / Status: {ConnectionStatus}\n" +
                   $"发送 / Sent: {messagesSent}\n" +
                   $"接收 / Received: {messagesReceived}\n" +
                   $"丢弃 / Dropped: {messagesDropped}";
        }

        #endregion
    }

    #region Data Structures

    /// <summary>
    /// ROS消息 / ROS Message
    /// </summary>
    [Serializable]
    public class ROSMessage
    {
        public string messageType;
        public object data;
        public float timestamp;
    }

    /// <summary>
    /// IK请求 / IK Request
    /// </summary>
    [Serializable]
    public class IKRequest
    {
        public Vector3 targetPosition;
        public Quaternion targetRotation;
        public float timestamp;
    }

    /// <summary>
    /// 规划请求 / Plan Request
    /// </summary>
    [Serializable]
    public class PlanRequest
    {
        public Vector3 targetPosition;
        public Quaternion targetRotation;
        public float timestamp;
    }

    /// <summary>
    /// 连接统计信息 / Connection Statistics
    /// </summary>
    [Serializable]
    public struct ConnectionStatistics
    {
        public int messagesSent;
        public int messagesReceived;
        public int messagesDropped;
        public int reconnectAttempts;
    }

    #endregion
}