terra-llm/Assets/LLM/Editor/Commands/AddComponentToAssetCommand.cs

using System;
using System.Linq;
using UnityEngine;
using UnityEditor;
using LLM.Editor.Helper;
using JetBrains.Annotations;

namespace LLM.Editor.Commands
{
    [Serializable]
    public class AddComponentParams
    {
        public string scriptName;
        public string componentName;
        public string targetIdentifier;
        public string childPath; // Optional path to a child object, e.g., "Launcher/SpawnPoint"
    }

    [UsedImplicitly]
    public class AddComponentToAssetCommand : ICommand
    {
        private readonly AddComponentParams _params;

        public AddComponentToAssetCommand(string jsonParams)
        {
            _params = jsonParams?.FromJson<AddComponentParams>();
        }

        public CommandOutcome Execute(Data.CommandContext context)
        {
            if (_params == null || string.IsNullOrEmpty(_params.targetIdentifier) || (string.IsNullOrEmpty(_params.scriptName) && string.IsNullOrEmpty(_params.componentName)))
            {
                Debug.LogError("[AddComponentCommand] Invalid parameters. Target identifier and script name are required.");
                return CommandOutcome.Error;
            }
            
            _params.scriptName ??= _params.componentName;

            var targetObject = ResolveIdentifier(_params.targetIdentifier);
            if (!targetObject)
            {
                Debug.LogError($"[AddComponentCommand] Could not find target with identifier '{_params.targetIdentifier}'.");
                return CommandOutcome.Error;
            }

            var scriptType = GetTypeByName(_params.scriptName);
            if (scriptType == null)
            {
                Debug.LogError($"[AddComponentCommand] Could not find script type '{_params.scriptName}'. Did it compile?");
                return CommandOutcome.Error;
            }
            
            // Determine if we are adding to the root object or a child
            GameObject objectToAddComponentTo = null;
            if (targetObject is GameObject rootGo)
            {
                if (string.IsNullOrEmpty(_params.childPath))
                {
                    objectToAddComponentTo = rootGo;
                }
                else
                {
                    var childTransform = rootGo.transform.Find(_params.childPath);
                    if (childTransform)
                    {
                        objectToAddComponentTo = childTransform.gameObject;
                    }
                    else
                    {
                        Debug.LogError($"[AddComponentCommand] Could not find child at path '{_params.childPath}' on target '{rootGo.name}'.");
                        return CommandOutcome.Error;
                    }
                }
            }

            if (!objectToAddComponentTo)
            {
                Debug.LogError($"[AddComponentCommand] Target '{targetObject.name}' is not a GameObject or the child path is invalid.");
                return CommandOutcome.Error;
            }

            // Handle whether we're editing a prefab asset or a scene instance
            if (PrefabUtility.IsPartOfPrefabAsset(objectToAddComponentTo))
            {
                var prefabPath = PrefabUtility.GetPrefabAssetPathOfNearestInstanceRoot(targetObject);
                var prefabContents = PrefabUtility.LoadPrefabContents(prefabPath);
                
                // Find the equivalent object within the loaded prefab contents
                var targetInPrefab = FindEquivalentObjectInPrefab(objectToAddComponentTo, prefabContents);
                if (targetInPrefab)
                {
                    targetInPrefab.AddComponent(scriptType);
                    PrefabUtility.SaveAsPrefabAsset(prefabContents, prefabPath);
                    PrefabUtility.UnloadPrefabContents(prefabContents);
                    Debug.Log($"[AddComponentCommand] Added component '{_params.scriptName}' to prefab at '{prefabPath}'.");
                }
                else
                {
                    PrefabUtility.UnloadPrefabContents(prefabContents);
                    Debug.LogError($"[AddComponentCommand] Could not find equivalent of '{objectToAddComponentTo.name}' in prefab contents.");
                    return CommandOutcome.Error;
                }
            }
            else // It's a scene object
            {
                Undo.AddComponent(objectToAddComponentTo, scriptType);
                Debug.Log($"[AddComponentCommand] Added component '{_params.scriptName}' to scene object '{objectToAddComponentTo.name}'.");
            }

            return CommandOutcome.Success;
        }
        
        private static GameObject FindEquivalentObjectInPrefab(GameObject original, GameObject prefabRoot)
        {
            if (original.name == prefabRoot.name) return prefabRoot; // Simple case
            var originalPath = AnimationUtility.CalculateTransformPath(original.transform, original.transform.root);
            var foundChild = prefabRoot.transform.Find(originalPath);
            return foundChild ? foundChild.gameObject : null;
        }

        private static UnityEngine.Object ResolveIdentifier(string identifier)
        {
            if (string.IsNullOrEmpty(identifier)) return null;
            if (int.TryParse(identifier, out var instanceId))
            {
                return EditorUtility.InstanceIDToObject(instanceId);
            }
            var path = AssetDatabase.GUIDToAssetPath(identifier);
            return !string.IsNullOrEmpty(path) ? AssetDatabase.LoadAssetAtPath<GameObject>(path) : null;
        }

        private static Type GetTypeByName(string typeName)
        {
            return AppDomain.CurrentDomain.GetAssemblies().Select(assembly => assembly.GetType(typeName)).FirstOrDefault(type => type != null);
        }
    }
}