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

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

namespace LLM.Editor.Commands
{
    [Serializable]
    public class AddComponentParams
    {
        public string scriptName;
        public string componentName;
        public string targetIdentifier; // This is now a LOGICAL name
        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)))
            {
                context.ErrorMessage = "Invalid parameters. Target identifier and script/component name are required.";
                return CommandOutcome.Error;
            }
            
            _params.scriptName ??= _params.componentName;

            var targetObject = CommandUtility.ResolveIdentifier(context, _params.targetIdentifier);
            if (!targetObject)
            {
                var internalError = context.ErrorMessage;
                context.ErrorMessage = $"Could not find target with logical name '{_params.targetIdentifier}'. The object may have been deleted or the identifier map is out of sync. \nInternal error: {internalError}";
                return CommandOutcome.Error;
            }

            var scriptType = GetTypeByName(_params.scriptName);
            if (scriptType == null)
            {
                context.ErrorMessage = $"Could not find script type '{_params.scriptName}'. It may not have compiled correctly or does not exist.";
                return CommandOutcome.Error;
            }
            
            if (targetObject is not GameObject rootGo)
            {
                context.ErrorMessage = $"Target '{targetObject.name}' is not a GameObject and components cannot be added to it.";
                return CommandOutcome.Error;
            }
            
            var objectToAddComponentTo = rootGo;
            if (!string.IsNullOrEmpty(_params.childPath))
            {
                var childTransform = rootGo.transform.Find(_params.childPath);
                if (childTransform)
                {
                    objectToAddComponentTo = childTransform.gameObject;
                }
                else
                {
                    context.ErrorMessage = $"Could not find child at path '{_params.childPath}' on target '{rootGo.name}'.";
                    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);
                
                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);
                    context.ErrorMessage = $"Could not find equivalent of '{objectToAddComponentTo.name}' in prefab contents.";
                    return CommandOutcome.Error;
                }
            }
            else // It's a scene object
            {
                var result = Undo.AddComponent(objectToAddComponentTo, scriptType);
                if (!result)
                {
                    context.ErrorMessage = $"Failed to add component '{_params.scriptName}' to scene object '{objectToAddComponentTo.name}', mostly since it already has one.";
                    return CommandOutcome.Error;
                }
                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;
            var originalPath = AnimationUtility.CalculateTransformPath(original.transform, original.transform.root);
            var foundChild = prefabRoot.transform.Find(originalPath);
            return foundChild ? foundChild.gameObject : null;
        }

        private static Type GetTypeByName(string typeName)
        {
            return AppDomain.CurrentDomain.GetAssemblies()
                .Select(assembly => assembly.GetType(typeName, false)) // Be non-throwing
                .FirstOrDefault(type => type != null);
        }
    }
}