unity2moveit2/unity-project/Assets/Scripts/Editor/UR5SceneSetup.cs

using UnityEngine;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityMoveIt2.Core;
using UnityMoveIt2.Communication;

namespace UnityMoveIt2.Editor
{
    /// <summary>
    /// UR5场景自动配置工具 / UR5 Scene Setup Tool
    /// 在Unity编辑器中一键创建完整的UR5测试场景
    /// </summary>
    public class UR5SceneSetup : EditorWindow
    {
        private string sceneName = "UR5_TestScene";
        private bool createCamera = true;
        private bool createLighting = true;
        private bool createGround = true;
        private bool createRosConnection = true;
        private string rosIP = "127.0.0.1";
        private int rosPort = 10000;

        [MenuItem("Unity-MoveIt2/Setup UR5 Scene")]
        public static void ShowWindow()
        {
            GetWindow<UR5SceneSetup>("UR5 场景设置");
        }

        private void OnGUI()
        {
            GUILayout.Label("UR5 场景自动配置工具", EditorStyles.boldLabel);
            EditorGUILayout.Space();

            sceneName = EditorGUILayout.TextField("场景名称", sceneName);
            EditorGUILayout.Space();

            GUILayout.Label("场景元素", EditorStyles.boldLabel);
            createCamera = EditorGUILayout.Toggle("创建相机", createCamera);
            createLighting = EditorGUILayout.Toggle("创建光照", createLighting);
            createGround = EditorGUILayout.Toggle("创建地面", createGround);
            createRosConnection = EditorGUILayout.Toggle("创建ROS连接", createRosConnection);

            if (createRosConnection)
            {
                EditorGUI.indentLevel++;
                rosIP = EditorGUILayout.TextField("ROS IP", rosIP);
                rosPort = EditorGUILayout.IntField("ROS Port", rosPort);
                EditorGUI.indentLevel--;
            }

            EditorGUILayout.Space();

            if (GUILayout.Button("创建 UR5 测试场景", GUILayout.Height(30)))
            {
                CreateUR5Scene();
            }

            EditorGUILayout.Space();
            EditorGUILayout.HelpBox(
                "此工具将创建包含以下内容的完整场景:\n" +
                "• UR5 机器人对象（需要手动添加mesh）\n" +
                "• 相机和光照\n" +
                "• ROS TCP连接\n" +
                "• 地面和坐标网格",
                MessageType.Info
            );
        }

        private void CreateUR5Scene()
        {
            Debug.Log("[UR5SceneSetup] 开始创建UR5测试场景...");

            // 创建新场景
            var newScene = EditorSceneManager.NewScene(NewSceneSetup.DefaultGameObjects, NewSceneMode.Single);

            // 创建UR5机器人
            GameObject ur5Robot = CreateUR5Robot();

            // 创建相机
            if (createCamera)
            {
                CreateSceneCamera(ur5Robot.transform.position);
            }

            // 创建光照
            if (createLighting)
            {
                CreateSceneLighting();
            }

            // 创建地面
            if (createGround)
            {
                CreateGround();
            }

            // 创建ROS连接
            if (createRosConnection)
            {
                CreateROSConnection();
            }

            // 保存场景
            string scenePath = $"Assets/Scenes/{sceneName}.unity";
            System.IO.Directory.CreateDirectory("Assets/Scenes");
            EditorSceneManager.SaveScene(newScene, scenePath);

            Debug.Log($"[UR5SceneSetup] ✓ 场景创建完成: {scenePath}");
            EditorUtility.DisplayDialog(
                "场景创建完成",
                $"UR5测试场景已创建!\n\n位置: {scenePath}\n\n下一步:\n1. 选中UR5对象\n2. 在Inspector中配置UniversalRobotAdapter\n3. 为各个link添加mesh模型",
                "确定"
            );
        }

        private GameObject CreateUR5Robot()
        {
            GameObject ur5 = new GameObject("UR5_Robot");
            ur5.transform.position = Vector3.zero;
            ur5.transform.rotation = Quaternion.identity;

            // 添加UR5手动设置组件
            UR5ManualSetup setup = ur5.AddComponent<UR5ManualSetup>();

            // 自动创建UR5结构
            setup.CreateUR5Structure();

            Debug.Log("[UR5SceneSetup] ✓ 创建UR5机器人对象");
            return ur5;
        }

        private void CreateSceneCamera()
        {
            GameObject mainCamera = GameObject.Find("Main Camera");
            if (mainCamera == null)
            {
                mainCamera = new GameObject("Main Camera");
                mainCamera.AddComponent<Camera>();
                mainCamera.tag = "MainCamera";
            }

            // 设置相机位置以观察UR5
            mainCamera.transform.position = new Vector3(1.5f, 1.0f, -1.5f);
            mainCamera.transform.LookAt(new Vector3(0, 0.5f, 0));

            Debug.Log("[UR5SceneSetup] ✓ 配置场景相机");
        }

        private void CreateSceneCamera(Vector3 lookAtPosition)
        {
            GameObject mainCamera = GameObject.Find("Main Camera");
            if (mainCamera == null)
            {
                mainCamera = new GameObject("Main Camera");
                mainCamera.AddComponent<Camera>();
                mainCamera.tag = "MainCamera";
            }

            mainCamera.transform.position = lookAtPosition + new Vector3(1.5f, 1.0f, -1.5f);
            mainCamera.transform.LookAt(lookAtPosition + Vector3.up * 0.5f);

            Debug.Log("[UR5SceneSetup] ✓ 配置场景相机");
        }

        private void CreateSceneLighting()
        {
            // 创建定向光
            GameObject directionalLight = GameObject.Find("Directional Light");
            if (directionalLight == null)
            {
                directionalLight = new GameObject("Directional Light");
                Light light = directionalLight.AddComponent<Light>();
                light.type = LightType.Directional;
                light.intensity = 1.0f;
                directionalLight.transform.rotation = Quaternion.Euler(50, -30, 0);
            }

            Debug.Log("[UR5SceneSetup] ✓ 配置场景光照");
        }

        private void CreateGround()
        {
            GameObject ground = GameObject.CreatePrimitive(PrimitiveType.Plane);
            ground.name = "Ground";
            ground.transform.position = new Vector3(0, -0.01f, 0);
            ground.transform.localScale = new Vector3(2, 1, 2);

            // 设置材质
            Renderer renderer = ground.GetComponent<Renderer>();
            Material mat = new Material(Shader.Find("Standard"));
            mat.color = new Color(0.8f, 0.8f, 0.8f);
            renderer.material = mat;

            Debug.Log("[UR5SceneSetup] ✓ 创建地面");
        }

        private void CreateROSConnection()
        {
            GameObject rosConnection = new GameObject("ROS_Connection");
            rosConnection.transform.position = Vector3.zero;

            // 添加ROSTCPBridge组件
            var bridge = rosConnection.AddComponent<ROSTCPBridge>();

            Debug.Log($"[UR5SceneSetup] ✓ 创建ROS连接对象 ({rosIP}:{rosPort})");
            Debug.LogWarning("[UR5SceneSetup] 请在Inspector中手动配置ROSTCPBridge的IP和端口");
        }
    }
}